claude version 6

src/main.rs

@@ -8,7 +8,6 @@ use std::intrinsics::transmute;
 use std::io::{Read, Write};
 use std::mem::MaybeUninit;
 use std::sync::Arc;
-use std::sync::atomic::{AtomicBool, Ordering};
 
 #[repr(C)]
 pub struct Meta {
@@ -76,12 +75,9 @@ fn main() -> Result<(), Box<dyn Error>> {
 
     thread::scope(|s| {
         for router in router_readers.values_mut() {
-            // 每个线程一个 mark_set 标志
-            let mark_set = Arc::new(AtomicBool::new(false));
-            let mark_set_clone = mark_set.clone();
-            
-            s.spawn(move |_| {
-                let mut buffer = [0u8; 1500 - 20]; // 保持原始大小
+            s.spawn(|_| {
+                // 使用 2048 字节缓冲区，足够处理大多数包且避免过大的内存占用
+                let mut buffer = vec![0u8; 2048];
                 let meta_size = size_of::<Meta>();
 
                 // 预初始化 Meta 头部
@@ -96,20 +92,16 @@ fn main() -> Result<(), Box<dyn Error>> {
                 buffer[..meta_size].copy_from_slice(meta_bytes);
 
                 loop {
-                    let n = router.tun_reader.read(&mut buffer[meta_size..]).unwrap();
-                    if let Some(ref addr) = *router.endpoint.read().unwrap() {
-                        router.encrypt(&mut buffer[meta_size..meta_size + n]);
-                        
-                        #[cfg(target_os = "linux")]
-                        {
-                            // 只在第一次设置 mark，避免重复的系统调用
-                            if !mark_set_clone.load(Ordering::Relaxed) {
+                    match router.tun_reader.read(&mut buffer[meta_size..]) {
+                        Ok(n) => {
+                            if let Some(ref addr) = *router.endpoint.read().unwrap() {
+                                router.encrypt(&mut buffer[meta_size..meta_size + n]);
+                                #[cfg(target_os = "linux")]
                                 let _ = router.socket.set_mark(router.config.mark);
-                                mark_set_clone.store(true, Ordering::Relaxed);
+                                let _ = router.socket.send_to(&buffer[..meta_size + n], addr);
                             }
                         }
-                        
-                        let _ = router.socket.send_to(&buffer[..meta_size + n], addr);
+                        Err(_) => continue,
                     }
                 }
             });
@@ -117,52 +109,34 @@ fn main() -> Result<(), Box<dyn Error>> {
 
         for (socket, mut router_writers) in router_writers3 {
             s.spawn(move |_| {
-                let mut recv_buf = [MaybeUninit::uninit(); 1500];
-                
-                // 预先计算常用值
-                let meta_size = size_of::<Meta>();
-                let local_id = config.local_id;
+                // 使用 2048 字节接收缓冲区
+                let mut recv_buf = vec![MaybeUninit::uninit(); 2048];
                 
                 loop {
-                    let _ = (|| {
-                        let (len, addr) = socket.recv_from(&mut recv_buf).unwrap();
+                    let _ = (|| -> Result<(), Box<dyn Error>> {
+                        let (len, addr) = socket.recv_from(&mut recv_buf)?;
                         let data: &mut [u8] = unsafe { transmute(&mut recv_buf[..len]) };
 
                         let packet = Ipv4Packet::new(data).ok_or("malformed packet")?;
                         let header_len = packet.get_header_length() as usize * 4;
+                        let (_ip_header, rest) = data
+                            .split_at_mut_checked(header_len)
+                            .ok_or("malformed packet")?;
+                        let (meta_bytes, payload) = rest
+                            .split_at_mut_checked(size_of::<Meta>())
+                            .ok_or("malformed packet")?;
+                        let meta: &Meta = unsafe { transmute(meta_bytes.as_ptr()) };
                         
-                        // 使用不安全的指针操作减少边界检查
-                        if len < header_len + meta_size {
-                            return Err("packet too small".into());
+                        if meta.dst_id == config.local_id && meta.reversed == 0 {
+                            let router = router_writers
+                                .get_mut(&meta.src_id)
+                                .ok_or("missing router")?;
+                            *router.endpoint.write().unwrap() = Some(addr);
+                            router.decrypt(payload, &local_secret);
+                            router.tun_writer.write_all(payload)?;
                         }
-                        
-                        let meta: &Meta = unsafe {
-                            &*(data.as_ptr().add(header_len) as *const Meta)
-                        };
-                        
-                        // 快速检查
-                        if meta.dst_id != local_id || meta.reversed != 0 {
-                            return Ok(());
-                        }
-                        
-                        let router = router_writers
-                            .get_mut(&meta.src_id)
-                            .ok_or("missing router")?;
-                            
-                        *router.endpoint.write().unwrap() = Some(addr);
-                        
-                        let payload_offset = header_len + meta_size;
-                        let payload = unsafe {
-                            std::slice::from_raw_parts_mut(
-                                data.as_mut_ptr().add(payload_offset),
-                                len - payload_offset
-                            )
-                        };
-                        
-                        router.decrypt(payload, &local_secret);
-                        router.tun_writer.write_all(payload)?;
 
-                        Ok::<(), Box<dyn Error>>(())
+                        Ok(())
                     })();
                 }
             });