From 4a9fab62db133dadc93aba980d3f8dad4dafc944 Mon Sep 17 00:00:00 2001
From: Jett Chen <jettchen12345@gmail.com>
Date: Thu, 7 Dec 2023 11:45:42 +0800
Subject: [PATCH] Update daemon

---
 burrow/src/daemon/instance.rs | 13 +++++++++----
 burrow/src/daemon/mod.rs      | 12 ++++++++++--
 burrow/src/wireguard/iface.rs |  6 +++---
 tun/src/tokio/mod.rs          |  2 +-
 4 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/burrow/src/daemon/instance.rs b/burrow/src/daemon/instance.rs
index f807ba2..bb94897 100644
--- a/burrow/src/daemon/instance.rs
+++ b/burrow/src/daemon/instance.rs
@@ -1,13 +1,14 @@
+use std::ops::Deref;
 use tracing::{debug, info, warn};
 use DaemonResponse;
-use tun::TunInterface;
+use tun::tokio::TunInterface;
 use crate::daemon::response::{DaemonResponseData, ServerConfig, ServerInfo};
 use super::*;
 
 pub struct DaemonInstance {
     rx: async_channel::Receiver<DaemonCommand>,
     sx: async_channel::Sender<DaemonResponse>,
-    tun_interface: Option<TunInterface>,
+    tun_interface: Option<Arc<RwLock<TunInterface>>>,
 }
 
 impl DaemonInstance {
@@ -19,13 +20,17 @@ impl DaemonInstance {
         }
     }
 
+    pub fn set_tun_interface(&mut self, tun_interface: Arc<RwLock<TunInterface>>) {
+        self.tun_interface = Some(tun_interface);
+    }
+
     async fn proc_command(&mut self, command: DaemonCommand) -> Result<DaemonResponseData> {
         info!("Daemon got command: {:?}", command);
         match command {
             DaemonCommand::Start(st) => {
                 if self.tun_interface.is_none() {
                     debug!("Daemon attempting start tun interface.");
-                    self.tun_interface = Some(st.tun.open()?);
+                    self.tun_interface = Some(Arc::new(RwLock::new(TunInterface::new(st.tun.open()?)?)));
                     info!("Daemon started tun interface");
                 } else {
                     warn!("Got start, but tun interface already up.");
@@ -39,7 +44,7 @@ impl DaemonInstance {
                         info!("{:?}", ti);
                         Ok(
                             DaemonResponseData::ServerInfo(
-                                ServerInfo::try_from(ti)?
+                                ServerInfo::try_from(ti.read().await.inner.get_ref())?
                             )
                         )
                     }
diff --git a/burrow/src/daemon/mod.rs b/burrow/src/daemon/mod.rs
index b42e350..b44efc1 100644
--- a/burrow/src/daemon/mod.rs
+++ b/burrow/src/daemon/mod.rs
@@ -1,4 +1,5 @@
 use std::net::{Ipv4Addr, SocketAddr, ToSocketAddrs};
+use std::sync::Arc;
 
 
 mod command;
@@ -11,6 +12,7 @@ use base64::{engine::general_purpose, Engine as _};
 pub use command::{DaemonCommand, DaemonStartOptions};
 use fehler::throws;
 use ip_network::{IpNetwork, Ipv4Network};
+use tokio::sync::RwLock;
 use instance::DaemonInstance;
 use crate::wireguard::{StaticSecret, Peer, Interface, PublicKey};
 pub use net::DaemonClient;
@@ -19,6 +21,7 @@ pub use net::DaemonClient;
 pub use net::start_srv;
 
 pub use response::{DaemonResponseData, DaemonResponse, ServerInfo};
+use crate::daemon::net::listen;
 
 #[throws]
 fn parse_key(string: &str) -> [u8; 32] {
@@ -49,12 +52,16 @@ pub async fn daemon_main() -> Result<()> {
     _tun.set_ipv4_addr(Ipv4Addr::from([10,13,13,2]))?;
     _tun.set_nonblocking(true)?;
     let tun = tun::tokio::TunInterface::new(_tun)?;
+    let tun_ref = Arc::new(RwLock::new(tun));
 
     let private_key = parse_secret_key("GNqIAOCRxjl/cicZyvkvpTklgQuUmGUIEkH7IXF/sEE=")?;
     let public_key = parse_public_key("uy75leriJay0+oHLhRMpV+A5xAQ0hCJ+q7Ww81AOvT4=")?;
     let preshared_key = Some(parse_key("s7lx/mg+reVEMnGnqeyYOQkzD86n2+gYnx1M9ygi08k=")?);
     let endpoint = "wg.burrow.rs:51820".to_socket_addrs()?.next().unwrap();
-    let iface = Interface::new(tun, vec![Peer {
+
+    inst.set_tun_interface(tun_ref.clone());
+
+    let iface = Interface::new(tun_ref, vec![Peer {
         endpoint,
         private_key,
         public_key,
@@ -62,6 +69,7 @@ pub async fn daemon_main() -> Result<()> {
         allowed_ips: vec![IpNetwork::V4(Ipv4Network::DEFAULT_ROUTE)],
     }])?;
 
-    iface.run().await;
+    tokio::try_join!(iface.run(), inst.run(), listen(commands_tx, response_rx))
+        .map(|_| {()});
     Ok(())
 }
diff --git a/burrow/src/wireguard/iface.rs b/burrow/src/wireguard/iface.rs
index 7d1b1ec..9f5dae4 100755
--- a/burrow/src/wireguard/iface.rs
+++ b/burrow/src/wireguard/iface.rs
@@ -83,8 +83,7 @@ pub struct Interface {
 
 impl Interface {
     #[throws]
-    pub fn new<I: IntoIterator<Item = Peer>>(tun: TunInterface, peers: I) -> Self {
-        let tun = Arc::new(RwLock::new(tun));
+    pub fn new<I: IntoIterator<Item = Peer>>(tun: Arc<RwLock<TunInterface>>, peers: I) -> Self {
         let pcbs: IndexedPcbs = peers
             .into_iter()
             .map(|peer| PeerPcb::new(peer, tun.clone()))
@@ -94,7 +93,7 @@ impl Interface {
         Self { tun, pcbs }
     }
 
-    pub async fn run(self) {
+    pub async fn run(self) -> anyhow::Result<()> {
         let pcbs = self.pcbs.clone();
         let tun = self.tun.clone();
         log::info!("starting interface");
@@ -187,5 +186,6 @@ impl Interface {
         }
         log::debug!("preparing to join..");
         join_all(tsks).await;
+        Ok(())
     }
 }
diff --git a/tun/src/tokio/mod.rs b/tun/src/tokio/mod.rs
index 2ade0a1..599e92c 100644
--- a/tun/src/tokio/mod.rs
+++ b/tun/src/tokio/mod.rs
@@ -5,7 +5,7 @@ use tracing::instrument;
 
 #[derive(Debug)]
 pub struct TunInterface {
-    inner: AsyncFd<crate::TunInterface>,
+    pub inner: AsyncFd<crate::TunInterface>,
 }
 
 impl TunInterface {