From 70dee329d1022ef4512b466415203b325ce5842f Mon Sep 17 00:00:00 2001
From: "Sijie.Sun" <sunsijie@buaa.edu.cn>
Date: Mon, 29 Apr 2024 21:02:05 +0800
Subject: [PATCH] fix crash bugs (#64)

---
 easytier/src/peers/peer_conn.rs      | 14 ++++++++++++++
 easytier/src/vpn_portal/wireguard.rs | 24 ++++++++++++++----------
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/easytier/src/peers/peer_conn.rs b/easytier/src/peers/peer_conn.rs
index 9043910..960397a 100644
--- a/easytier/src/peers/peer_conn.rs
+++ b/easytier/src/peers/peer_conn.rs
@@ -135,6 +135,20 @@ impl PeerConn {
         *need_retry = true;
 
         let rsp = rsp?;
+        let Some(peer_mgr_hdr) = rsp.peer_manager_header() else {
+            return Err(Error::WaitRespError(format!(
+                "unexpected packet: {:?}, cannot decode peer manager hdr",
+                rsp
+            )));
+        };
+
+        if peer_mgr_hdr.packet_type != PacketType::HandShake as u8 {
+            return Err(Error::WaitRespError(format!(
+                "unexpected packet type: {:?}",
+                peer_mgr_hdr.packet_type
+            )));
+        }
+
         let rsp = HandshakeRequest::decode(rsp.payload()).map_err(|e| {
             Error::WaitRespError(format!("decode handshake response error: {:?}", e))
         })?;
diff --git a/easytier/src/vpn_portal/wireguard.rs b/easytier/src/vpn_portal/wireguard.rs
index 7cbcd9b..2423ccd 100644
--- a/easytier/src/vpn_portal/wireguard.rs
+++ b/easytier/src/vpn_portal/wireguard.rs
@@ -302,16 +302,20 @@ Endpoint = {listenr_addr} # should be the public ip of the vpn server
     async fn list_clients(&self) -> Vec<String> {
         self.inner
             .as_ref()
-            .unwrap()
-            .wg_peer_ip_table
-            .iter()
-            .map(|x| {
-                x.value()
-                    .endpoint_addr
-                    .as_ref()
-                    .map(|x| x.to_string())
-                    .unwrap_or_default()
+            .and_then(|w| {
+                Some(
+                    w.wg_peer_ip_table
+                        .iter()
+                        .map(|x| {
+                            x.value()
+                                .endpoint_addr
+                                .as_ref()
+                                .map(|x| x.to_string())
+                                .unwrap_or_default()
+                        })
+                        .collect(),
+                )
             })
-            .collect()
+            .unwrap_or_default()
     }
 }