Pop the chunk only when received the complete packet

crates/agent/src/handshake_chunk.rs

@@ -0,0 +1,87 @@
+use std::error::Error;
+
+use protocol::cluster::AgentTunnelRequest;
+
+#[derive(Debug, Default)]
+pub struct HandshakedChunk {
+    data: Vec<u8>,
+}
+
+impl HandshakedChunk {
+    pub fn put_buf(&mut self, buf: &[u8]) {
+        self.data.extend_from_slice(buf);
+    }
+
+    pub fn pop_chunk(&self) -> Result<Option<(AgentTunnelRequest, Vec<u8>)>, Box<dyn Error>> {
+        let handshake_len = u16::from_be_bytes([self.data[0], self.data[1]]) as usize;
+        if handshake_len + 2 > self.data.len() {
+            return Ok(None);
+        }
+
+        match AgentTunnelRequest::try_from(&self.data[2..handshake_len + 2]) {
+            Ok(handshake) => {
+                let data = self.data[handshake_len + 2..].to_vec();
+                Ok(Some((handshake, data)))
+            }
+            Err(e) => Err(Box::new(e)),
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    fn build_sample() -> Vec<u8> {
+        let request = AgentTunnelRequest {
+            service: None,
+            tls: true,
+            domain: "ac5eb61578bebe98636148414efc482f.local.ha.8xff.io".to_string(),
+        };
+
+        let handshake_buf = Vec::from(&request);
+        let handshake_len = handshake_buf.len() as u16;
+        let mut buf = handshake_len.to_be_bytes().to_vec();
+        buf.append(&mut handshake_buf.clone());
+        buf.append([1, 2, 3].to_vec().as_mut());
+
+        // println!("buf: {} {handshake_len}", buf.len());
+
+        buf
+    }
+
+    #[test]
+    pub fn push_full_chunk() {
+        let mut chunk = HandshakedChunk::default();
+        let buf = build_sample();
+        chunk.put_buf(&buf);
+
+        let res = chunk.pop_chunk();
+        assert!(res.is_ok());
+        let (handshake, data) = res.unwrap().unwrap();
+        assert_eq!(
+            handshake.domain,
+            "ac5eb61578bebe98636148414efc482f.local.ha.8xff.io"
+        );
+        assert_eq!(data, [1, 2, 3]);
+    }
+
+    #[test]
+    pub fn push_part_chunk() {
+        let mut chunk = HandshakedChunk::default();
+        let buf = build_sample();
+
+        chunk.put_buf(&buf[0..2]);
+        let res = chunk.pop_chunk().unwrap();
+        assert!(res.is_none());
+        chunk.put_buf(&buf[2..]);
+        let res = chunk.pop_chunk();
+        assert!(res.is_ok());
+        let (handshake, data) = res.unwrap().unwrap();
+        assert_eq!(
+            handshake.domain,
+            "ac5eb61578bebe98636148414efc482f.local.ha.8xff.io"
+        );
+        assert_eq!(data, [1, 2, 3]);
+    }
+}

crates/agent/src/lib.rs

@@ -9,6 +9,7 @@ use protocol::cluster::AgentTunnelRequest;
 use crate::local_tunnel::LocalTunnel;
 
 mod connection;
+mod handshake_chunk;
 mod local_tunnel;
 
 pub use connection::{
@@ -27,51 +28,73 @@ where
     log::info!("sub_connection pipe to local_tunnel start");
     let (mut reader1, mut writer1) = sub_connection.split();
     let mut first_pkt = [0u8; 4096];
-    let (local_tunnel, first_pkt_start, first_pkt_end) = match reader1.read(&mut first_pkt).await {
-        Ok(first_pkt_len) => {
-            log::info!("first pkt size: {}", first_pkt_len);
-            if first_pkt_len < 2 {
-                log::error!("first pkt size is < 4 => close");
-                return;
-            }
-            let handshake_len = u16::from_be_bytes([first_pkt[0], first_pkt[1]]) as usize;
-            if handshake_len + 2 > first_pkt_len {
-                log::error!("first pkt size is < handshake {handshake_len} + 2 => close");
+
+    let mut chunk = handshake_chunk::HandshakedChunk::default();
+    let mut handshake: Option<AgentTunnelRequest> = None;
+    let mut data = Vec::new();
+
+    let timeout_duration = std::time::Duration::from_secs(10);
+
+    loop {
+        select! {
+            _ = async_std::task::sleep(timeout_duration).fuse() => {
+                log::error!("handshake timeout");
                 return;
             }
-            match AgentTunnelRequest::try_from(&first_pkt[2..(handshake_len + 2)]) {
-                Ok(handshake) => {
-                    if let Some(dest) =
-                        registry.dest_for(handshake.tls, handshake.service, &handshake.domain)
-                    {
-                        log::info!("create tunnel to dest {}", dest);
-                        (
-                            LocalTcpTunnel::new(dest).await,
-                            handshake_len + 2,
-                            first_pkt_len,
-                        )
-                    } else {
-                        log::warn!(
-                            "dest for service {:?} tls {} domain {} not found",
-                            handshake.service,
-                            handshake.tls,
-                            handshake.domain
-                        );
+            res = reader1.read(&mut first_pkt).fuse() => match res {
+                Ok(n) => {
+                    if n < 2 {
+                        log::error!("first pkt size is < 4 => close");
                         return;
                     }
+                    chunk.put_buf(&first_pkt[..n]);
+                    let res = chunk.pop_chunk();
+                    match res {
+                        Ok(Some((req, mut payload))) => {
+                            handshake = Some(req);
+                            data.append(&mut payload);
+                            break;
+                        }
+                        Ok(None) => {
+                            continue;
+                        }
+                        Err(e) => {
+                            log::error!("handshake parse error: {}", e);
+                            return;
+                        }
+                    }
                 }
                 Err(e) => {
-                    log::error!("handshake parse error: {}", e);
+                    log::error!("read first pkt error: {}", e);
                     return;
                 }
             }
         }
-        Err(e) => {
-            log::error!("read first pkt error: {}", e);
+    }
+
+    let handshake = match handshake {
+        Some(handshake) => handshake,
+        None => {
+            log::error!("handshake is None");
             return;
         }
     };
 
+    let local_tunnel = if let Some(dest) =
+        registry.dest_for(handshake.tls, handshake.service, &handshake.domain)
+    {
+        log::info!("create tunnel to dest {}", dest);
+        LocalTcpTunnel::new(dest).await
+    } else {
+        log::warn!(
+            "dest for service {:?} tls {} domain {} not found",
+            handshake.service,
+            handshake.tls,
+            handshake.domain
+        );
+        return;
+    };
+
     let local_tunnel = match local_tunnel {
         Ok(local_tunnel) => local_tunnel,
         Err(e) => {
@@ -82,10 +105,7 @@ where
 
     let (mut reader2, mut writer2) = local_tunnel.split();
 
-    if let Err(e) = writer2
-        .write_all(&first_pkt[first_pkt_start..first_pkt_end])
-        .await
-    {
+    if let Err(e) = writer2.write_all(&data).await {
         log::error!("write first pkt to local_tunnel error: {}", e);
         return;
     }