From 96bed38ad0f3aa3323cf349adfabad87de7f8354 Mon Sep 17 00:00:00 2001
From: Prabhpreet Dua <615318+prabhpreet@users.noreply.github.com>
Date: Thu, 13 Jun 2024 17:34:48 +0530
Subject: [PATCH] Enable privileged only on linux

---
 wireguard-broker/src/bin/priviledged.rs    | 101 +++++++++++----------
 wireguard-broker/src/bin/socket_handler.rs |   1 -
 2 files changed, 54 insertions(+), 48 deletions(-)

diff --git a/wireguard-broker/src/bin/priviledged.rs b/wireguard-broker/src/bin/priviledged.rs
index 57f78e6..34fe251 100644
--- a/wireguard-broker/src/bin/priviledged.rs
+++ b/wireguard-broker/src/bin/priviledged.rs
@@ -1,57 +1,64 @@
-#![cfg(target_os = "linux")]
-use std::io::{stdin, stdout, Read, Write};
-use std::result::Result;
-
-use rosenpass_wireguard_broker::api::msgs;
-use rosenpass_wireguard_broker::api::server::BrokerServer;
-use rosenpass_wireguard_broker::brokers::netlink as wg;
-
-#[derive(thiserror::Error, Debug)]
-pub enum BrokerAppError {
-    #[error(transparent)]
-    IoError(#[from] std::io::Error),
-    #[error(transparent)]
-    WgConnectError(#[from] wg::ConnectError),
-    #[error(transparent)]
-    WgSetPskError(#[from] wg::SetPskError),
-    #[error("Oversized message {}; something about the request is fatally wrong", .0)]
-    OversizedMessage(u64),
+fn main() {
+    #[cfg(target_os = "linux")]
+    linux::main().unwrap();
 }
 
-fn main() -> Result<(), BrokerAppError> {
-    let mut broker = BrokerServer::new(wg::NetlinkWireGuardBroker::new()?);
+#[cfg(target_os = "linux")]
+pub mod linux {
+    use std::io::{stdin, stdout, Read, Write};
+    use std::result::Result;
 
-    let mut stdin = stdin().lock();
-    let mut stdout = stdout().lock();
-    loop {
-        // Read the message length
-        let mut len = [0u8; 8];
-        stdin.read_exact(&mut len)?;
+    use rosenpass_wireguard_broker::api::msgs;
+    use rosenpass_wireguard_broker::api::server::BrokerServer;
+    use rosenpass_wireguard_broker::brokers::netlink as wg;
 
-        // Parse the message length
-        let len = u64::from_le_bytes(len);
-        if (len as usize) > msgs::REQUEST_MSG_BUFFER_SIZE {
-            return Err(BrokerAppError::OversizedMessage(len));
-        }
+    #[derive(thiserror::Error, Debug)]
+    pub enum BrokerAppError {
+        #[error(transparent)]
+        IoError(#[from] std::io::Error),
+        #[error(transparent)]
+        WgConnectError(#[from] wg::ConnectError),
+        #[error(transparent)]
+        WgSetPskError(#[from] wg::SetPskError),
+        #[error("Oversized message {}; something about the request is fatally wrong", .0)]
+        OversizedMessage(u64),
+    }
 
-        // Read the message itself
-        let mut req_buf = [0u8; msgs::REQUEST_MSG_BUFFER_SIZE];
-        let req_buf = &mut req_buf[..(len as usize)];
-        stdin.read_exact(req_buf)?;
+    pub fn main() -> Result<(), BrokerAppError> {
+        let mut broker = BrokerServer::new(wg::NetlinkWireGuardBroker::new()?);
 
-        // Process the message
-        let mut res_buf = [0u8; msgs::RESPONSE_MSG_BUFFER_SIZE];
-        let res = match broker.handle_message(req_buf, &mut res_buf) {
-            Ok(len) => &res_buf[..len],
-            Err(e) => {
-                eprintln!("Error processing message for wireguard PSK broker: {e:?}");
-                continue;
+        let mut stdin = stdin().lock();
+        let mut stdout = stdout().lock();
+        loop {
+            // Read the message length
+            let mut len = [0u8; 8];
+            stdin.read_exact(&mut len)?;
+
+            // Parse the message length
+            let len = u64::from_le_bytes(len);
+            if (len as usize) > msgs::REQUEST_MSG_BUFFER_SIZE {
+                return Err(BrokerAppError::OversizedMessage(len));
             }
-        };
 
-        // Write the response
-        stdout.write_all(&(res.len() as u64).to_le_bytes())?;
-        stdout.write_all(&res)?;
-        stdout.flush()?;
+            // Read the message itself
+            let mut req_buf = [0u8; msgs::REQUEST_MSG_BUFFER_SIZE];
+            let req_buf = &mut req_buf[..(len as usize)];
+            stdin.read_exact(req_buf)?;
+
+            // Process the message
+            let mut res_buf = [0u8; msgs::RESPONSE_MSG_BUFFER_SIZE];
+            let res = match broker.handle_message(req_buf, &mut res_buf) {
+                Ok(len) => &res_buf[..len],
+                Err(e) => {
+                    eprintln!("Error processing message for wireguard PSK broker: {e:?}");
+                    continue;
+                }
+            };
+
+            // Write the response
+            stdout.write_all(&(res.len() as u64).to_le_bytes())?;
+            stdout.write_all(&res)?;
+            stdout.flush()?;
+        }
     }
 }
diff --git a/wireguard-broker/src/bin/socket_handler.rs b/wireguard-broker/src/bin/socket_handler.rs
index 5e46f17..e1693c5 100644
--- a/wireguard-broker/src/bin/socket_handler.rs
+++ b/wireguard-broker/src/bin/socket_handler.rs
@@ -1,4 +1,3 @@
-#![cfg(target_os = "linux")]
 use std::process::Stdio;
 
 use tokio::io::{AsyncReadExt, AsyncWriteExt};