Skip to content

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

claude version 7

parent f4b3c34c
Pipeline #37398 passed with stages
      in 47 seconds
      Hide whitespace changes
      Inline Side-by-side
      Showing with 65 additions and 76 deletions
      src/main.rs
      View file @ d1fd29ce
      ......@@ -8,6 +8,7 @@ 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 {
      ......@@ -38,6 +39,7 @@ pub struct Config {
      }
      use crossbeam_utils::thread;
      use grouping_by::GroupingBy;
      use pnet::packet::ipv4::Ipv4Packet;
      use socket2::Socket;
      fn main() -> Result<(), Box<dyn Error>> {
      ......@@ -74,52 +76,40 @@ fn main() -> Result<(), Box<dyn Error>> {
      thread::scope(|s| {
      for router in router_readers.values_mut() {
      #[cfg(target_os = "linux")]
      let mark_set = std::sync::atomic::AtomicBool::new(false);
      // 每个线程一个 mark_set 标志
      let mark_set = Arc::new(AtomicBool::new(false));
      let mark_set_clone = mark_set.clone();
      s.spawn(move |_| {
      // 使用更大的缓冲区以支持巨帧
      let mut buffer = vec![0u8; 9000];
      let mut buffer = [0u8; 1500 - 20]; // 保持原始大小
      let meta_size = size_of::<Meta>();
      // 预初始化 Meta 头部（local -> remote）
      // 预初始化 Meta 头部
      let meta = Meta {
      src_id: config.local_id,
      dst_id: router.config.remote_id,
      reversed: 0,
      };
      // 直接写入缓冲区，避免额外的切片操作
      unsafe {
      let meta_ptr = buffer.as_mut_ptr() as *mut Meta;
      *meta_ptr = meta;
      }
      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);
      loop {
      match router.tun_reader.read(&mut buffer[meta_size..]) {
      Ok(n) => {
      // 使用 try_read 减少锁争用
      if let Ok(endpoint_guard) = router.endpoint.try_read() {
      if let Some(ref addr) = *endpoint_guard {
      router.encrypt(&mut buffer[meta_size..meta_size + n]);
      #[cfg(target_os = "linux")]
      {
      // 只在第一次设置 mark
      use std::sync::atomic::Ordering;
      if !mark_set.load(Ordering::Relaxed) {
      let _ = router.socket.set_mark(router.config.mark);
      mark_set.store(true, Ordering::Relaxed);
      }
      }
      let _ = router.socket.send_to(&buffer[..meta_size + n], addr);
      }
      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) {
      let _ = router.socket.set_mark(router.config.mark);
      mark_set_clone.store(true, Ordering::Relaxed);
      }
      }
      Err(_) => {
      // TUN 读取失败时短暂休眠，避免 CPU 空转
      std::thread::sleep(std::time::Duration::from_millis(1));
      }
      let _ = router.socket.send_to(&buffer[..meta_size + n], addr);
      }
      }
      });
      ......@@ -127,54 +117,53 @@ fn main() -> Result<(), Box<dyn Error>> {
      for (socket, mut router_writers) in router_writers3 {
      s.spawn(move |_| {
      // 使用更大的缓冲区和重用内存
      let mut recv_buf = vec![MaybeUninit::uninit(); 9000];
      let mut recv_buf = [MaybeUninit::uninit(); 1500];
      // 预先计算常用值
      let meta_size = size_of::<Meta>();
      let local_id = config.local_id;
      loop {
      match socket.recv_from(&mut recv_buf) {
      Ok((len, addr)) => {
      // 快速边界检查
      if len < 20 + meta_size {
      continue;
      }
      let data: &mut [u8] = unsafe { transmute(&mut recv_buf[..len]) };
      // 优化：直接计算 IP 头部长度，避免创建 Ipv4Packet
      let header_len = ((data[0] & 0x0f) as usize) * 4;
      if len < header_len + meta_size {
      continue;
      }
      // 直接从内存读取 Meta，避免额外的切片操作
      let meta: &Meta = unsafe {
      &*(data.as_ptr().add(header_len) as *const Meta)
      };
      if meta.dst_id == config.local_id && meta.reversed == 0 {
      if let Some(router) = router_writers.get_mut(&meta.src_id) {
      // 使用 try_write 减少锁争用
      if let Ok(mut endpoint) = router.endpoint.try_write() {
      *endpoint = Some(addr);
      }
      let payload_start = header_len + meta_size;
      let payload = &mut data[payload_start..];
      router.decrypt(payload, &local_secret);
      // 忽略写入错误，继续处理下一个数据包
      let _ = router.tun_writer.write_all(payload);
      }
      }
      let _ = (|| {
      let (len, addr) = 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")?;
      let header_len = packet.get_header_length() as usize * 4;
      // 使用不安全的指针操作减少边界检查
      if len < header_len + meta_size {
      return Err("packet too small".into());
      }
      Err(_) => {
      // Socket 接收失败时短暂休眠
      std::thread::sleep(std::time::Duration::from_millis(1));
      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>>(())
      })();
      }
      });
      }
      ......
      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