rust-next

Cargo.toml

@@ -4,12 +4,12 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-tun = "0.7"
-socket2 = { version = "0.5.8", features = ["all"] }
+tun = "0.8"
+socket2 = { version = "0.6.1", features = ["all"] }
 pnet = "0.35.0"
-serde = { version = "1.0.217", features = ["derive"] }
+serde = { version = "1.0.228", features = ["derive"] }
 serde_json = "1.0"
 base64 = "0.22.1"
 crossbeam = "0.8.4"
-crossbeam-utils = "0.8.20"
-grouping_by = "0.2.2"
+crossbeam-utils = "0.8.21"
+libc = "0.2.178"

src/main.rs

 mod router;
-
-use crate::router::{Router, RouterReader, RouterWriter, SECRET_LENGTH};
-use std::collections::HashMap;
+use crate::router::{Router, SECRET_LENGTH};
 use std::env;
 use std::error::Error;
-use std::intrinsics::transmute;
 use std::io::{Read, Write};
-use std::mem::MaybeUninit;
-use std::sync::Arc;
+use std::mem::size_of;
+use std::mem::{transmute, MaybeUninit};
+use std::sync::atomic::Ordering;
 
 #[repr(C)]
 pub struct Meta {
@@ -18,10 +16,16 @@ pub struct Meta {
 
 use serde::Deserialize;
 
+#[derive(Deserialize)]
+pub struct Config {
+    pub local_id: u8,
+    pub local_secret: String,
+    pub routers: Vec<ConfigRouter>,
+}
 #[derive(Deserialize)]
 pub struct ConfigRouter {
     pub remote_id: u8,
-    pub proto: i32,
+    pub proto: u8,
     pub family: u8,
     pub mark: u32,
     pub endpoint: String,
@@ -30,85 +34,55 @@ pub struct ConfigRouter {
     pub up: String,
 }
 
-#[derive(Deserialize)]
-pub struct Config {
-    pub local_id: u8,
-    pub local_secret: String,
-    pub routers: Vec<ConfigRouter>,
-}
 use crossbeam_utils::thread;
-use grouping_by::GroupingBy;
 use pnet::packet::ipv4::Ipv4Packet;
-use socket2::Socket;
 
 fn main() -> Result<(), Box<dyn Error>> {
     let config: Config = serde_json::from_str(env::args().nth(1).ok_or("need param")?.as_str())?;
+
     let local_secret: [u8; SECRET_LENGTH] = Router::create_secret(config.local_secret.as_str())?;
-    let mut sockets: HashMap<u16, Arc<Socket>> = HashMap::new();
-    let routers: HashMap<u8, Router> = config
+    let routers: Vec<Router> = config
         .routers
-        .iter()
-        .map(|c| Router::new(c, &mut sockets).map(|router| (c.remote_id, router)))
-        .collect::<Result<_, _>>()?;
-    let (mut router_readers, router_writers): (
-        HashMap<u8, RouterReader>,
-        HashMap<u8, RouterWriter>,
-    ) = routers
         .into_iter()
-        .map(|(id, router)| {
-            let (reader, writer) = router.split();
-            ((id, reader), (id, writer))
-        })
-        .unzip();
-    let router_writers3: Vec<(Arc<Socket>, HashMap<u8, RouterWriter>)> = router_writers
-        .into_iter()
-        .grouping_by(|(_, v)| v.key())
-        .into_iter()
-        .map(|(k, v)| {
-            (
-                Arc::clone(sockets.get_mut(&k).unwrap()),
-                v.into_iter().collect(),
-            )
-        })
-        .collect();
+        .map(|c| Router::new(c, config.local_id))
+        .collect::<Result<Vec<_>, _>>()?;
+
     println!("created tuns");
 
     thread::scope(|s| {
-        for router in router_readers.values_mut() {
-            s.spawn(|_| {
-                let mut buffer = [0u8; 1500 - 20]; // minus typical IP header space
-                let meta_size = size_of::<Meta>();
+        for router in routers {
+            let (mut reader, mut writer) = router.split();
 
+            s.spawn(move |_| {
+                let mut buffer = [0u8; 1500 - 20]; // minus typical IP header space
+                const META_SIZE: usize = size_of::<Meta>();
                 // Pre-initialize with our Meta header (local -> remote)
                 let meta = Meta {
                     src_id: config.local_id,
-                    dst_id: router.config.remote_id,
+                    dst_id: reader.config.remote_id,
                     reversed: 0,
                 };
                 // Turn the Meta struct into bytes
-                let meta_bytes = unsafe {
-                    std::slice::from_raw_parts(&meta as *const Meta as *const u8, meta_size)
-                };
-                buffer[..meta_size].copy_from_slice(meta_bytes);
+                let meta_bytes: &[u8; META_SIZE] =
+                    unsafe { &*(&meta as *const Meta as *const [u8; META_SIZE]) };
+                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")]
-                        let _ = router.socket.set_mark(router.config.mark);
-                        let _ = router.socket.send_to(&buffer[..meta_size + n], addr);
+                    let n = reader.tun_reader.read(&mut buffer[META_SIZE..]).unwrap();
+                    let guard = crossbeam::epoch::pin();
+                    let shared = reader.endpoint.load(Ordering::Acquire, &guard);
+                    if let Some(addr) = unsafe { shared.as_ref() } {
+                        reader.encrypt(&mut buffer[META_SIZE..META_SIZE + n]);
+                        let _ = reader.socket.send_to(&buffer[..META_SIZE + n], addr);
                     }
                 }
             });
-        }
 
-        for (socket, mut router_writers) in router_writers3 {
             s.spawn(move |_| {
                 let mut recv_buf = [MaybeUninit::uninit(); 1500];
                 loop {
                     let _ = (|| {
-                        let (len, addr) = socket.recv_from(&mut recv_buf).unwrap();
+                        let (len, addr) = writer.socket.recv_from(&mut recv_buf).unwrap();
                         let data: &mut [u8] = unsafe { transmute(&mut recv_buf[..len]) };
 
                         let packet = Ipv4Packet::new(data).ok_or("malformed packet")?;
@@ -119,16 +93,24 @@ fn main() -> Result<(), Box<dyn Error>> {
                         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 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 guard = crossbeam::epoch::pin();
+                        let current_shared = writer.endpoint.load(Ordering::SeqCst, &guard);
+                        let is_same = unsafe { current_shared.as_ref() }
+                            .map(|c| *c == addr)
+                            .unwrap_or(false);
+                        if !is_same {
+                            let new_shared = crossbeam::epoch::Owned::new(addr).into_shared(&guard);
+                            let old_shared =
+                                writer.endpoint.swap(new_shared, Ordering::SeqCst, &guard);
+                            unsafe {
+                                guard.defer_destroy(old_shared);
+                            }
                         }
 
+                        writer.decrypt(payload, &local_secret);
+                        writer.tun_writer.write_all(payload)?;
+
                         Ok::<(), Box<dyn Error>>(())
                     })();
                 }

src/router.rs

-use socket2::{Domain, Protocol, SockAddr, Socket, Type};
-use std::collections::hash_map::Entry;
-use std::collections::HashMap;
+use base64::prelude::BASE64_STANDARD;
+use base64::Engine;
+use socket2::{Domain, Protocol, SockAddr, SockFilter, Socket, Type};
 use std::net::ToSocketAddrs;
 use std::process::{Command, ExitStatus};
-use std::sync::{Arc, RwLock};
+use std::sync::Arc;
 use tun::{Reader, Writer};
 pub const SECRET_LENGTH: usize = 32;
-use crate::ConfigRouter;
-use base64::prelude::*;
+use crate::{ConfigRouter, Meta};
+use crossbeam::epoch::Atomic;
 
 // tun -> raw
-pub struct RouterReader<'a> {
-    pub config: &'a ConfigRouter,
+pub struct RouterReader {
+    pub config: ConfigRouter,
     pub secret: [u8; SECRET_LENGTH],
     pub tun_reader: Reader,
     pub socket: Arc<Socket>,
-    pub endpoint: Arc<RwLock<Option<SockAddr>>>,
+    pub endpoint: Arc<Atomic<SockAddr>>,
 }
 
-impl<'a> RouterReader<'a> {
+impl RouterReader {
     pub(crate) fn encrypt(&self, data: &mut [u8]) {
         for (i, b) in data.iter_mut().enumerate() {
             *b ^= self.secret[i % SECRET_LENGTH];
@@ -27,34 +27,30 @@ impl<'a> RouterReader<'a> {
 }
 
 // raw -> tun
-pub struct RouterWriter<'a> {
-    pub config: &'a ConfigRouter,
+pub struct RouterWriter {
     pub tun_writer: Writer,
-    pub endpoint: Arc<RwLock<Option<SockAddr>>>,
+    pub socket: Arc<Socket>,
+    pub endpoint: Arc<Atomic<SockAddr>>,
 }
 
-impl<'a> RouterWriter<'a> {
+impl RouterWriter {
     pub(crate) fn decrypt(&self, data: &mut [u8], secret: &[u8; SECRET_LENGTH]) {
         for (i, b) in data.iter_mut().enumerate() {
             *b ^= secret[i % SECRET_LENGTH];
         }
     }
-
-    pub(crate) fn key(&self) -> u16 {
-        Router::key(self.config)
-    }
 }
 
-pub struct Router<'a> {
-    pub config: &'a ConfigRouter,
+pub struct Router {
+    pub config: ConfigRouter,
     pub secret: [u8; SECRET_LENGTH],
     pub tun_reader: Reader,
     pub tun_writer: Writer,
     pub socket: Arc<Socket>,
-    pub endpoint: Arc<RwLock<Option<SockAddr>>>,
+    pub endpoint: Arc<Atomic<SockAddr>>,
 }
 
-impl<'a> Router<'a> {
+impl Router {
     pub(crate) fn create_secret(
         config: &str,
     ) -> Result<[u8; SECRET_LENGTH], Box<dyn std::error::Error>> {
@@ -65,30 +61,78 @@ impl<'a> Router<'a> {
         Ok(secret)
     }
 
-    fn key(config: &ConfigRouter) -> u16 {
-        (config.family as u16) << 8 | config.proto as u16
-    }
     fn create_raw_socket(
         config: &ConfigRouter,
-        sockets: &mut HashMap<u16, Arc<Socket>>,
-    ) -> Result<Arc<Socket>, Box<dyn std::error::Error>> {
-        let key = Router::key(config);
-        let result = match sockets.entry(key) {
-            Entry::Occupied(entry) => entry.get().clone(),
-            Entry::Vacant(entry) => entry
-                .insert(Arc::new(Socket::new(
+        local_id: u8,
+    ) -> Result<Arc<Socket>, std::io::Error> {
+        let result = Socket::new(
             if config.family == 6 {
                 Domain::IPV6
             } else {
                 Domain::IPV4
             },
             Type::RAW,
-                    Some(Protocol::from(config.proto)),
-                )?))
-                .clone(),
+            Some(Protocol::from(config.proto as i32)),
+        )?;
+        #[cfg(target_os = "linux")]
+        result.set_mark(config.mark)?;
+        Self::attach_readable_filter(config, local_id, &result)?;
+        Ok(Arc::new(result))
+    }
+
+    fn attach_readable_filter(
+        config: &ConfigRouter,
+        local_id: u8,
+        socket: &Socket,
+    ) -> std::io::Result<()> {
+        // 由于多个对端可能会使用相同的 ipprpto 号，这里确保每个 socket 上只会收到自己对应的对端发来的消息
+        const META_SIZE: usize = size_of::<Meta>();
+        let meta = Meta {
+            src_id: config.remote_id,
+            dst_id: local_id,
+            reversed: 0,
         };
-        Ok(result)
+        let meta_bytes: [u8; META_SIZE] =
+            unsafe { *(&meta as *const Meta as *const [u8; META_SIZE]) };
+        let target_val = u32::from_be_bytes(meta_bytes);
+
+        // 如果你的协议是 UDP，这里必须是 8 (跳过 UDP 头: SrcPort(2)+DstPort(2)+Len(2)+Sum(2))
+        // 如果是纯自定义 IP 协议，这里是 0
+        let payload_offset = 0;
+
+        socket.attach_filter(&[
+            // 1. 【加载 IP 头长度到寄存器 X】
+            // BPF_LDX (Load X) | BPF_B (Byte) | BPF_MSH (Magic Shift)
+            // 这是一个特殊的 BPF 指令，专门用于计算 IPv4 头长度：X = 4 * (IP[0] & 0xf)
+            bpf_stmt(libc::BPF_LDX | libc::BPF_B | libc::BPF_MSH, 0),
+            // 2. 【读取 Payload 数据到累加器 A】
+            // BPF_LD (Load) | BPF_W (Word, 4 bytes) | BPF_IND (Indirect, relative to X)
+            // 逻辑：A = Packet[X + k]
+            bpf_stmt(libc::BPF_LD | libc::BPF_W | libc::BPF_IND, payload_offset),
+            // 3. 【比较并跳转】
+            // BPF_JMP (Jump) | BPF_JEQ (Jump if Equal) | BPF_K (Const)
+            // 逻辑：if (A == target_val) goto True(0); else goto False(1);
+            // jt=0: 继续执行下一条
+            // jf=1: 跳过下一条 (直接跳到 Reject)
+            bpf_jump(
+                libc::BPF_JMP | libc::BPF_JEQ | libc::BPF_K,
+                target_val,
+                0,
+                1,
+            ),
+            // 4. 【接受 (True 路径)】
+            // BPF_RET (Return) | BPF_K (Constant)
+            // 返回 -1 (0xFFFFFFFF) 表示截取整个包的最大长度（即接收包）
+            bpf_stmt(libc::BPF_RET | libc::BPF_K, u32::MAX),
+            // 5. 【拒绝 (False 路径)】
+            // BPF_RET (Return) | BPF_K (Constant)
+            // 返回 0 表示截取 0 字节（即丢弃包）
+            bpf_stmt(libc::BPF_RET | libc::BPF_K, 0),
+        ])?;
+
+        Ok(())
     }
+
     fn create_tun_device(
         config: &ConfigRouter,
     ) -> Result<(Reader, Writer), Box<dyn std::error::Error>> {
@@ -104,21 +148,18 @@ impl<'a> Router<'a> {
 
     fn create_endpoint(
         config: &ConfigRouter,
-    ) -> Result<Arc<RwLock<Option<SockAddr>>>, Box<dyn std::error::Error>> {
+    ) -> Result<Arc<Atomic<SockAddr>>, Box<dyn std::error::Error>> {
         let parsed = (config.endpoint.clone(), 0u16)
             .to_socket_addrs()?
             .next()
             .ok_or(config.endpoint.clone())?;
-        Ok(Arc::new(RwLock::new(Some(parsed.into()))))
+        Ok(Arc::new(Atomic::new(parsed.into())))
     }
 
-    pub fn new(
-        config: &'a ConfigRouter,
-        sockets: &mut HashMap<u16, Arc<Socket>>,
-    ) -> Result<Router<'a>, Box<dyn std::error::Error>> {
+    pub fn new(config: ConfigRouter, local_id: u8) -> Result<Router, Box<dyn std::error::Error>> {
         let secret = Self::create_secret(config.remote_secret.as_str())?;
         let endpoint = Self::create_endpoint(&config)?;
-        let socket = Self::create_raw_socket(&config, sockets)?;
+        let socket = Self::create_raw_socket(&config, local_id)?;
         let (tun_reader, tun_writer) = Self::create_tun_device(&config)?;
         Self::run_up_script(&config)?;
 
@@ -134,11 +175,11 @@ impl<'a> Router<'a> {
         Ok(router)
     }
 
-    pub fn split(self) -> (RouterReader<'a>, RouterWriter<'a>) {
+    pub fn split(self) -> (RouterReader, RouterWriter) {
         let writer = RouterWriter {
-            config: self.config,
-            endpoint: Arc::clone(&self.endpoint),
+            endpoint: self.endpoint.clone(),
             tun_writer: self.tun_writer,
+            socket: self.socket.clone(),
         };
 
         let reader = RouterReader {
@@ -152,3 +193,11 @@ impl<'a> Router<'a> {
         (reader, writer)
     }
 }
+
+fn bpf_stmt(code: u32, k: u32) -> SockFilter {
+    SockFilter::new(code as u16, 0, 0, k)
+}
+
+fn bpf_jump(code: u32, k: u32, jt: u8, jf: u8) -> SockFilter {
+    SockFilter::new(code as u16, jt, jf, k)
+}