Skip to content

T
tun
Commit 74855653 authored by nanahira's avatar nanahira
1 Browse files
Options

try

parent 3b8eeedb
Pipeline #37397 failed with stages
      in 22 seconds
      Show whitespace changes
      Inline Side-by-side
      Showing with 78 additions and 38 deletions
      src/main.rs
      View file @ 74855653
      ...@@ -8,6 +8,9 @@ use std::intrinsics::transmute; ...@@ -8,6 +8,9 @@ use std::intrinsics::transmute;
      use std::io::{Read, Write}; use std::io::{Read, Write};
      use std::mem::MaybeUninit; use std::mem::MaybeUninit;
      use std::sync::Arc; use std::sync::Arc;
      use std::sync::atomic::{AtomicBool, Ordering};
      use std::net::SocketAddr;
      use parking_lot::RwLock;
      #[repr(C)] #[repr(C)]
      pub struct Meta { pub struct Meta {
      ...@@ -38,8 +41,16 @@ pub struct Config { ...@@ -38,8 +41,16 @@ pub struct Config {
      } }
      use crossbeam_utils::thread; use crossbeam_utils::thread;
      use grouping_by::GroupingBy; use grouping_by::GroupingBy;
      use pnet::packet::ipv4::Ipv4Packet;
      use socket2::Socket; use socket2::Socket;
      // 优化的 RouterWriter 包装器，用于缓存 mark 状态
      struct OptimizedRouterWriter {
      writer: RouterWriter,
      #[cfg(target_os = "linux")]
      mark_set: AtomicBool,
      }
      fn main() -> Result<(), Box<dyn Error>> { fn main() -> Result<(), Box<dyn Error>> {
      let config: Config = serde_json::from_str(env::args().nth(1).ok_or("need param")?.as_str())?; 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 local_secret: [u8; SECRET_LENGTH] = Router::create_secret(config.local_secret.as_str())?;
      ...@@ -59,36 +70,51 @@ fn main() -> Result<(), Box<dyn Error>> { ...@@ -59,36 +70,51 @@ fn main() -> Result<(), Box<dyn Error>> {
      ((id, reader), (id, writer)) ((id, reader), (id, writer))
      }) })
      .unzip(); .unzip();
      let router_writers3: Vec<(Arc<Socket>, HashMap<u8, RouterWriter>)> = router_writers
      // 使用 parking_lot 的 RwLock 替换标准库的 RwLock
      let router_writers = router_writers
      .into_iter()
      .map(|(id, writer)| {
      (id, Arc::new(RwLock::new(OptimizedRouterWriter {
      writer,
      #[cfg(target_os = "linux")]
      mark_set: AtomicBool::new(false),
      })))
      })
      .collect::<HashMap<_, _>>();
      let router_writers3: Vec<(Arc<Socket>, HashMap<u8, Arc<RwLock<OptimizedRouterWriter>>>)> = router_writers
      .iter()
      .map(|(id, writer)| (*id, writer.read().writer.key(), Arc::clone(writer)))
      .into_iter() .into_iter()
      .grouping_by(|(_, v)| v.key()) .grouping_by(|(_, key, _)| *key)
      .into_iter() .into_iter()
      .map(|(k, v)| { .map(|(k, v)| {
      ( (
      Arc::clone(sockets.get_mut(&k).unwrap()), Arc::clone(sockets.get_mut(&k).unwrap()),
      v.into_iter().collect(), v.into_iter().map(|(id, _, writer)| (id, writer)).collect(),
      ) )
      }) })
      .collect(); .collect();
      println!("created tuns"); println!("created tuns");
      thread::scope(|s| { thread::scope(|s| {
      // 发送线程池
      for router in router_readers.values_mut() { for router in router_readers.values_mut() {
      #[cfg(target_os = "linux")] s.spawn(|_| {
      let mark_set = std::sync::atomic::AtomicBool::new(false);
      s.spawn(move |_| {
      // 使用更大的缓冲区以支持巨帧 // 使用更大的缓冲区以支持巨帧
      let mut buffer = vec![0u8; 9000]; let mut buffer = vec![0u8; 9000];
      let meta_size = size_of::<Meta>(); let meta_size = size_of::<Meta>();
      // 预初始化 Meta 头部（local -> remote） // 预初始化 Meta 头部
      let meta = Meta { let meta = Meta {
      src_id: config.local_id, src_id: config.local_id,
      dst_id: router.config.remote_id, dst_id: router.config.remote_id,
      reversed: 0, reversed: 0,
      }; };
      // 直接写入缓冲区，避免额外的切片操作
      // 直接写入缓冲区，避免额外的内存分配
      unsafe { unsafe {
      let meta_ptr = buffer.as_mut_ptr() as *mut Meta; let meta_ptr = buffer.as_mut_ptr() as *mut Meta;
      *meta_ptr = meta; *meta_ptr = meta;
      ...@@ -100,15 +126,15 @@ fn main() -> Result<(), Box<dyn Error>> { ...@@ -100,15 +126,15 @@ fn main() -> Result<(), Box<dyn Error>> {
      // 使用 try_read 减少锁争用 // 使用 try_read 减少锁争用
      if let Ok(endpoint_guard) = router.endpoint.try_read() { if let Ok(endpoint_guard) = router.endpoint.try_read() {
      if let Some(ref addr) = *endpoint_guard { if let Some(ref addr) = *endpoint_guard {
      // 原地加密，避免额外的内存拷贝
      router.encrypt(&mut buffer[meta_size..meta_size + n]); router.encrypt(&mut buffer[meta_size..meta_size + n]);
      #[cfg(target_os = "linux")] #[cfg(target_os = "linux")]
      { {
      // 只在第一次设置 mark // 只在第一次或 mark 改变时设置
      use std::sync::atomic::Ordering; if !router.mark_set.load(Ordering::Relaxed) {
      if !mark_set.load(Ordering::Relaxed) {
      let _ = router.socket.set_mark(router.config.mark); let _ = router.socket.set_mark(router.config.mark);
      mark_set.store(true, Ordering::Relaxed); router.mark_set.store(true, Ordering::Relaxed);
      } }
      } }
      ...@@ -117,62 +143,76 @@ fn main() -> Result<(), Box<dyn Error>> { ...@@ -117,62 +143,76 @@ fn main() -> Result<(), Box<dyn Error>> {
      } }
      } }
      Err(_) => { Err(_) => {
      // TUN 读取失败时短暂休眠，避免 CPU 空转 // 忽略读取错误，继续循环
      std::thread::sleep(std::time::Duration::from_millis(1)); continue;
      } }
      } }
      } }
      }); });
      } }
      for (socket, mut router_writers) in router_writers3 { // 接收线程池
      for (socket, router_writers_map) in router_writers3 {
      s.spawn(move |_| { s.spawn(move |_| {
      // 使用更大的缓冲区和重用内存 // 使用更大的缓冲区
      let mut recv_buf = vec![MaybeUninit::uninit(); 9000]; let mut recv_buf = vec![MaybeUninit::uninit(); 9000];
      let meta_size = size_of::<Meta>(); let meta_size = size_of::<Meta>();
      loop { loop {
      match socket.recv_from(&mut recv_buf) { match socket.recv_from(&mut recv_buf) {
      Ok((len, addr)) => { Ok((len, addr)) => {
      // 快速边界检查 // 快速路径：直接处理数据，不创建 Ipv4Packet
      if len < 20 + meta_size { if len < 20 + meta_size {
      continue; continue; // 数据包太小，跳过
      } }
      let data: &mut [u8] = unsafe { transmute(&mut recv_buf[..len]) }; let data: &mut [u8] = unsafe {
      std::slice::from_raw_parts_mut(
      recv_buf.as_mut_ptr() as *mut u8,
      len
      )
      };
      // 优化：直接计算 IP 头部长度，避免创建 Ipv4Packet // 快速获取 IP 头部长度
      let header_len = ((data[0] & 0x0f) as usize) * 4; let header_len = ((data[0] & 0x0f) as usize) * 4;
      if len < header_len + meta_size { if len < header_len + meta_size {
      continue; continue; // 数据包格式错误，跳过
      } }
      // 直接从内存读取 Meta，避免额外的切片操作 // 直接解析 Meta 结构
      let meta: &Meta = unsafe { let meta: &Meta = unsafe {
      &*(data.as_ptr().add(header_len) as *const Meta) &*(data[header_len..].as_ptr() as *const Meta)
      }; };
      if meta.dst_id == config.local_id && meta.reversed == 0 { if meta.dst_id == config.local_id && meta.reversed == 0 {
      if let Some(router) = router_writers.get_mut(&meta.src_id) { // 使用 try_read 减少锁争用
      // 使用 try_write 减少锁争用 if let Some(router_lock) = router_writers_map.get(&meta.src_id) {
      if let Ok(mut endpoint) = router.endpoint.try_write() { if let Ok(mut router) = router_lock.try_write() {
      // 更新端点地址
      if let Ok(mut endpoint) = router.writer.endpoint.try_write() {
      *endpoint = Some(addr); *endpoint = Some(addr);
      } }
      // 原地解密
      let payload_start = header_len + meta_size; let payload_start = header_len + meta_size;
      let payload = &mut data[payload_start..]; let payload_len = len - payload_start;
      router.writer.decrypt(
      router.decrypt(payload, &local_secret); &mut data[payload_start..len],
      &local_secret
      );
      // 忽略写入错误，继续处理下一个数据包 // 写入 TUN 设备
      let _ = router.tun_writer.write_all(payload); let _ = router.writer.tun_writer.write_all(
      &data[payload_start..len]
      );
      }
      } }
      } }
      } }
      Err(_) => { Err(_) => {
      // Socket 接收失败时短暂休眠 // 忽略接收错误，继续循环
      std::thread::sleep(std::time::Duration::from_millis(1)); continue;
      } }
      } }
      } }
      ......
      Markdown is supported
      0% or
      You are about to add 0 people to the discussion. Proceed with caution.
      Finish editing this message first!
      Please register or to comment