From 13a5ab83521edb389747019b3ea857b50aea4e9a Mon Sep 17 00:00:00 2001
From: Jett Chen <jettchen12345@gmail.com>
Date: Wed, 22 Nov 2023 21:33:22 +0800
Subject: [PATCH] add support for preshared keys

---
 burrow/src/daemon/mod.rs      | 14 ++++++++++----
 burrow/src/wireguard/iface.rs |  4 ----
 burrow/src/wireguard/pcb.rs   | 35 ++++++++++++++++++++++-------------
 burrow/src/wireguard/peer.rs  |  1 +
 4 files changed, 33 insertions(+), 21 deletions(-)

diff --git a/burrow/src/daemon/mod.rs b/burrow/src/daemon/mod.rs
index 112f48b..9c3bd14 100644
--- a/burrow/src/daemon/mod.rs
+++ b/burrow/src/daemon/mod.rs
@@ -21,18 +21,22 @@ pub use net::start_srv;
 pub use response::{DaemonResponseData, DaemonResponse, ServerInfo};
 
 #[throws]
-fn parse_secret_key(string: &str) -> StaticSecret {
+fn parse_key(string: &str) -> [u8; 32] {
     let value = general_purpose::STANDARD.decode(string)?;
     let mut key = [0u8; 32];
     key.copy_from_slice(&value[..]);
+    key
+}
+
+#[throws]
+fn parse_secret_key(string: &str) -> StaticSecret {
+    let key = parse_key(string)?;
     StaticSecret::from(key)
 }
 
 #[throws]
 fn parse_public_key(string: &str) -> PublicKey {
-    let value = general_purpose::STANDARD.decode(string)?;
-    let mut key = [0u8; 32];
-    key.copy_from_slice(&value[..]);
+    let key = parse_key(string)?;
     PublicKey::from(key)
 }
 
@@ -47,11 +51,13 @@ pub async fn daemon_main() -> Result<()> {
 
     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 {
         endpoint,
         private_key,
         public_key,
+        preshared_key,
         allowed_ips: vec![IpNetwork::V4(Ipv4Network::DEFAULT_ROUTE)],
     }])?;
 
diff --git a/burrow/src/wireguard/iface.rs b/burrow/src/wireguard/iface.rs
index 0f82440..a427bf8 100755
--- a/burrow/src/wireguard/iface.rs
+++ b/burrow/src/wireguard/iface.rs
@@ -142,10 +142,6 @@ impl Interface {
                     Ok(siz) => {
                         log::info!("received {} bytes from peer",siz);
                         log::debug!("bytes: {:?}", &recv_buf[..siz]);
-                        match tun.send(&recv_buf[..siz]).await{
-                            Ok(..) => log::debug!("sent packet to interface"),
-                            Err(e) => log::error!("failed to send packet {}", e),
-                        }
                     },
                     Err(e) => {
                         log::error!("failed to receive packet {}", e);
diff --git a/burrow/src/wireguard/pcb.rs b/burrow/src/wireguard/pcb.rs
index 78bfad0..8d7cf4a 100755
--- a/burrow/src/wireguard/pcb.rs
+++ b/burrow/src/wireguard/pcb.rs
@@ -23,7 +23,7 @@ pub struct PeerPcb {
 impl PeerPcb {
     #[throws]
     pub fn new(peer: Peer) -> Self {
-        let tunnel = Tunnel::new(peer.private_key, peer.public_key, None, None, 1, None)
+        let tunnel = Tunnel::new(peer.private_key, peer.public_key, peer.preshared_key, None, 1, None)
             .map_err(|s| anyhow::anyhow!("{}", s))?;
 
         Self {
@@ -72,22 +72,31 @@ impl PeerPcb {
             };
             let mut res_buf = [0;1500];
             let (len, addr) = socket.recv_from(&mut res_buf).await?;
-            let res_dat = &res_buf[..len];
+            let mut res_dat = &res_buf[..len];
             tracing::debug!("Decapsulating {} bytes from {}", len, addr);
             tracing::debug!("{:?}", &res_dat);
-            match self.tunnel.decapsulate(None, res_dat, &mut buf[..]) {
-                TunnResult::Done => {tracing::debug!("Decapsulate done")}
-                TunnResult::Err(e) => {
-                    tracing::error!(message = "Decapsulate error", error = ?e)
+            loop {
+                match self.tunnel.decapsulate(None, res_dat, &mut buf[..]) {
+                    TunnResult::Done => {
+                        tracing::debug!("Decapsulate done");
+                        break;
+                    }
+                    TunnResult::Err(e) => {
+                        tracing::error!(message = "Decapsulate error", error = ?e);
+                        break;
+                    }
+                    TunnResult::WriteToNetwork(packet) => {
+                        tracing::debug!("WriteToNetwork: {:?}", packet);
+                        res_dat = &[];
+                        continue;
+                    }
+                    TunnResult::WriteToTunnelV4(packet, addr) => {
+                        tracing::debug!("WriteToTunnelV4: {:?}, {:?}", packet, addr);
+                        continue;
+                    }
+                    e => panic!("Unexpected result from decapsulate: {:?}", e),
                 }
-                TunnResult::WriteToNetwork(packet) => {
-                    tracing::debug!("sending {} bytes to {}", packet.len(), addr);
-                    let socket = self.socket().await?;
-                    socket.send(packet).await?;
-                }
-                _ => panic!("Unexpected result from decapsulate"),
             }
-            tracing::debug!("received {} bytes from {}", len, addr);
             return Ok(len)
         }
     }
diff --git a/burrow/src/wireguard/peer.rs b/burrow/src/wireguard/peer.rs
index 8a74ce1..cc8a296 100755
--- a/burrow/src/wireguard/peer.rs
+++ b/burrow/src/wireguard/peer.rs
@@ -10,6 +10,7 @@ pub struct Peer {
     pub private_key: StaticSecret,
     pub public_key: PublicKey,
     pub allowed_ips: Vec<IpNetwork>,
+    pub preshared_key: Option<[u8; 32]>
 }
 
 impl fmt::Debug for Peer {