minor performance improvements

Author: ylxdzsw

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "v2socks"
-version = "0.2.0"
+version = "0.2.1"
 authors = ["ylxdzsw <[email protected]>"]
 edition = "2018"
 

PKGBUILD

@@ -3,7 +3,7 @@
 pkgname=v2socks
 pkgdesc="An opinioned lightweight socks5 server and vmess (v2ray) client."
 pkgrel=1
-pkgver=0.2.0
+pkgver=0.2.1
 arch=(any)
 url='https://github.com/ylxdzsw/v2socks'
 license=(GPL3)

readme.md

@@ -1,4 +1,4 @@
 V2socks
 =======
 
-An opinioned lightweight socks5 server and vmess (v2ray) client implemented in Rust.
+An opinionated lightweight socks5 server and vmess (v2ray) client implemented in Rust.

src/lib.rs

@@ -1,10 +1,13 @@
+#![allow(irrefutable_let_patterns)]
 #![allow(dead_code, unused_imports)]
-#![deny(bare_trait_objects)]
 #![allow(non_camel_case_types)]
+#![deny(bare_trait_objects)]
+#![warn(clippy::all)]
 
 mod socks;
 mod vmess;
 
+use oh_my_rust::*;
 pub use socks::*;
 pub use vmess::*;
 
@@ -18,9 +21,9 @@ pub enum Addr {
 impl std::fmt::Display for Addr {
     fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            Addr::V4(x) => std::fmt::Display::fmt(&std::net::Ipv4Addr::from(x.clone()), fmt),
-            Addr::V6(x) => std::fmt::Display::fmt(&std::net::Ipv6Addr::from(x.clone()), fmt),
-            Addr::Domain(x) => std::fmt::Display::fmt(&String::from_utf8(x.to_vec()).unwrap(), fmt) // somewhat ugly, need more time to fight the type checker
+            Addr::V4(x) => std::fmt::Display::fmt(&std::net::Ipv4Addr::from(*x), fmt),
+            Addr::V6(x) => std::fmt::Display::fmt(&std::net::Ipv6Addr::from(*x), fmt),
+            Addr::Domain(x) => std::fmt::Display::fmt(std::str::from_utf8(x).msg(std::fmt::Error)?, fmt)
         }
     }
 }

src/main.rs

@@ -29,18 +29,14 @@ fn main() {
             let server = Socks5Server::new(port);
             vmess(&server, proxy.into(), parse_uid(user_id).unwrap())
         },
-        _ => {
-            eprint!("{}", USAGE)
-        },
+        _ => eprint!("{}", USAGE)
     }
 }
 
 fn parse_uid(x: &str) -> Option<[u8; 16]> {
-    let mut r = [0; 16];
     let x = x.replace('-', "");
     let list: Vec<_> = (0..32).step_by(2).map(|i| u8::from_str_radix(&x[i..i+2], 16).unwrap()).collect();
-    r.clone_from_slice(list.get(0..16)?);
-    Some(r)
+    list.get(0..16).map(|x| [0; 16].apply(|buf| buf.copy_from_slice(x)))
 }
 
 fn is_normal_close(e: &std::io::Error) -> bool {
@@ -51,7 +47,7 @@ fn is_normal_close(e: &std::io::Error) -> bool {
 }
 
 fn vmess(server: &Socks5Server, proxy: String, user_id: [u8; 16]) {
-    let connect = Box::leak(Box::new(move |dest, port| {
+    let connect = leak(move |dest, port| {
         let client = std::net::TcpStream::connect(&proxy)?;
         debug!("connect {}:{} through proxy", &dest, port);
 
@@ -63,22 +59,19 @@ fn vmess(server: &Socks5Server, proxy: String, user_id: [u8; 16]) {
         let local_port = local.port();
 
         Ok((local_addr, local_port, (dest, port, client)))
-    }));
+    });
 
     #[allow(non_snake_case)]
-    let pass = Box::leak(Box::new(move |(dest, port, conn): (Addr, u16, std::net::TcpStream), mut stream: std::net::TcpStream| {
-        let mut key = [0; 16];
-        thread_rng().fill_bytes(&mut key);
-
-        let mut IV = [0; 16];
-        thread_rng().fill_bytes(&mut IV);
+    let pass = leak(move |(dest, port, conn): (Addr, u16, std::net::TcpStream), mut stream: std::net::TcpStream| {
+        let key = [0; 16].apply(|x| thread_rng().fill_bytes(x));
+        let IV = [0; 16].apply(|x| thread_rng().fill_bytes(x));
 
         {
             let conn = conn.try_clone().expect("failed to clone TCP handle");
             let mut stream = stream.try_clone().expect("failed to clone TCP handle");
 
             std::thread::spawn(move || {
-                let mut buffer = Box::new( unsafe { std::mem::uninitialized::<[u8; 16384]>() } );
+                let mut buffer = Box::new([0; 16384]);
                 let mut reader = match VmessReader::new(conn, key, IV) {
                     Some(x) => x,
                     None => return warn!("reader handshake failed")
@@ -88,21 +81,21 @@ fn vmess(server: &Socks5Server, proxy: String, user_id: [u8; 16]) {
                         Ok(0) => break,
                         Ok(x) => x,
                         Err(ref e) if is_normal_close(e) => break,
-                        Err(e) => { error!("{}", e); break }
+                        Err(e) => { warn!("{}", e); break }
                     };
 
                     match stream.write_all(&buffer[..len]) {
                         Ok(_) => debug!("read {} bytes", len),
                         Err(ref e) if is_normal_close(e) => break,
-                        Err(e) => { error!("{}", e); break }
+                        Err(e) => { warn!("{}", e); break }
                     }
                 }
                 reader.close();
                 debug!("closed reading")
             });
         }
 
-        let mut buffer = Box::new( unsafe { std::mem::uninitialized::<[u8; 16384]>() } );
+        let mut buffer = Box::new([0; 16384]);
         let mut writer = match VmessWriter::new(conn, user_id, dest, port, key, IV) {
             Some(x) => x,
             None => return warn!("writer handshake failed")
@@ -112,19 +105,19 @@ fn vmess(server: &Socks5Server, proxy: String, user_id: [u8; 16]) {
                 Ok(0) => break,
                 Ok(x) => x,
                 Err(ref e) if is_normal_close(e) => break,
-                Err(e) => { error!("{}", e); break }
+                Err(e) => { warn!("{}", e); break }
             };
 
             match writer.write_all(&buffer[..len]) {
                 Ok(_) => debug!("sent {} bytes", len),
                 Err(ref e) if is_normal_close(e) => break,
-                Err(e) => { error!("{}", e); break }
+                Err(e) => { warn!("{}", e); break }
             }
         }
 
         writer.close();
         debug!("closed writing");
-    }));
+    });
 
     server.listen(connect, pass)
 }

src/socks.rs

@@ -15,7 +15,7 @@ macro_rules! close_on_error {
         match $ex {
             Ok(x) => x,
             Err(e) => {
-                return error!("{}", e)
+                return warn!("{}", e)
             }
         }
     }}
@@ -30,7 +30,7 @@ impl Socks5Server {
         Socks5Server { port }
     }
 
-    pub fn listen<T>(&self, connect: &'static (impl Fn(Addr, u16) -> std::io::Result<(Addr, u16, T)> + Sync), pass: &'static (impl Fn(T, TcpStream) + Sync)) {
+    pub fn listen<T>(&self, connect: &'static (impl Fn(Addr, u16) -> std::io::Result<(Addr, u16, T)> + Sync), pass: &'static (impl Fn(T, TcpStream) + Sync)) -> ! {
         let socket = TcpListener::bind(format!("0.0.0.0:{}", self.port)).expect("Address already in use");
         info!("v2socks starts listening at 0.0.0.0:{}", self.port);
 
@@ -44,32 +44,31 @@ impl Socks5Server {
                 pass(proxy, stream);
             });
         }
+
+        unreachable!()
     }
 }
 
 fn initialize(stream: &mut (impl ReadExt + Write)) -> Result<(), String> {
-    let mut header = [0, 0];
-    stream.read_exact(&mut header).map_err(|_| "read initial bits failed")?;
-    
+    let header = read_exact!(stream, [0, 0]).msg("read initial bits failed")?;
+
     if header[0] != 5 {
         return Err(format!("unsupported socks version {}", header[0]))
     }
 
-    let list: Vec<u8> = stream.read_exact_alloc(header[1] as usize).map_err(|_| "read methods failed")?;
+    let list: Vec<u8> = stream.read_exact_alloc(header[1] as usize).msg("read methods failed")?;
 
     if !list.contains(&0) {
-        stream.write(&[5, 0xff]).map_err(|_| "write response failed")?;
-        return Err("client do not support NO AUTH method".to_owned())
+        stream.write(&[5, 0xff]).msg("write response failed")?;
+        return Err("client do not support NO AUTH method".to_string())
     }
 
-    stream.write(&[5, 0]).map_err(|_| "write response failed")?;
+    stream.write(&[5, 0]).msg("write response failed")?;
     Ok(())
 }
 
 fn read_request(stream: &mut (impl ReadExt + Write)) -> Result<(Addr, u16), String> {
-    let mut header = [0; 4];
-    stream.read_exact(&mut header).map_err(|_| "read request header failed")?;
-    let [ver, cmd, _rev, atyp] = header;
+    let [ver, cmd, _rev, atyp] = read_exact!(stream, [0; 4]).msg("read request header failed")?;
 
     if ver != 5 {
         return Err(format!("unsupported socks version {}", ver))
@@ -80,27 +79,24 @@ fn read_request(stream: &mut (impl ReadExt + Write)) -> Result<(Addr, u16), Stri
     }
 
     let addr = match atyp {
-        0x01 => Addr::V4(read_exact!(stream, [0; 4]).map_err(|_| "read v4 address failed")?),
-        0x04 => Addr::V6(read_exact!(stream, [0; 16]).map_err(|_| "read v6 address failed")?),
+        0x01 => Addr::V4(read_exact!(stream, [0; 4]).msg("read v4 address failed")?),
+        0x04 => Addr::V6(read_exact!(stream, [0; 16]).msg("read v6 address failed")?),
         0x03 => {
-            let mut len = [0];
-            stream.read_exact(&mut len).map_err(|_| "read domain length failed")?;
-            let len = len[0];
-
-            Addr::Domain(stream.read_exact_alloc(len as usize).map_err(|_| "read domain failed")?.into_boxed_slice())
+            let len = read_exact!(stream, [0]).msg("read domain length failed")?[0];
+            Addr::Domain(stream.read_exact_alloc(len as usize).msg("read domain failed")?.into_boxed_slice())
         },
-        _ => return Err("unknown ATYP".to_owned())
+        _ => return Err("unknown ATYP".to_string())
     };
 
-    let mut port: [u8; 2] = [0; 2];
-    stream.read_exact(&mut port).map_err(|_| "read port failed")?;
-    let port: u16 = (port[0] as u16) << 8 | port[1] as u16;
+    let port = read_exact!(stream, [0; 2]).msg("read port failed")?;
+    let port = (port[0] as u16) << 8 | port[1] as u16;
 
     Ok((addr, port))
 }
 
 fn reply_request(stream: &mut (impl ReadExt + Write), addr: Addr, port: u16) -> Result<(), String> {
-    let mut reply = vec![5, 0, 0];
+    let mut reply = Vec::with_capacity(22); // cover V4 and V6
+    reply.extend_from_slice(&[5, 0, 0]);
 
     match addr {
         Addr::V4(x) => {
@@ -121,7 +117,7 @@ fn reply_request(stream: &mut (impl ReadExt + Write), addr: Addr, port: u16) ->
     reply.push((port >> 8) as u8);
     reply.push(port as u8);
 
-    stream.write(&reply).map_err(|_| "write reply failed")?;
+    stream.write(&reply).msg("write reply failed")?;
 
     Ok(())
 }

src/vmess.rs

@@ -83,7 +83,7 @@ impl<R: ReadExt> Read for VmessReader<R> {
         // 4. verify checksum
         let checksum = fnv1a(&buf[..len-4]);
         if checksum.to_be_bytes() != temp {
-            panic!("fuck")
+            return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "invalid checksum!"))
         }
 
         Ok(len-4)
@@ -182,8 +182,7 @@ impl<W: Write> VmessWriter<W> {
 
         AES128CFB::new(header_key, header_IV).encode(&mut buffer);
 
-        self.writer.write_all(&buffer)?;
-        Ok(())
+        self.writer.write_all(&buffer)
     }
 }
 
@@ -195,7 +194,7 @@ impl<W: Write> Write for VmessWriter<W> {
         buf.extend_from_slice(&fnv1a(data).to_be_bytes());
         buf.extend_from_slice(data);
         self.encoder.encode(&mut buf); // this is the right code. the fucking protocol document is misleading!
-        self.writer.write_all(&mut buf)?;
+        self.writer.write_all(&buf)?;
         Ok(data.len())
     }
 
@@ -204,6 +203,7 @@ impl<W: Write> Write for VmessWriter<W> {
     }
 }
 
+#[allow(clippy::unreadable_literal)]
 fn fnv1a(x: &[u8]) -> u32 {
     let prime = 16777619;
     let mut hash = 0x811c9dc5;
@@ -230,8 +230,7 @@ impl AES128CFB {
     fn encode(&mut self, data: &mut [u8]) {
         for byte in data.iter_mut() {
             if self.p == 16 {
-                let temp = self.state.clone();
-                crypto::aessafe::AesSafe128Encryptor::new(&self.key).encrypt_block(&temp, &mut self.state);
+                crypto::aessafe::AesSafe128Encryptor::new(&self.key).encrypt_block(&self.state.clone(), &mut self.state);
                 self.p = 0;
             }
             *byte ^= self.state[self.p];
@@ -243,8 +242,7 @@ impl AES128CFB {
     fn decode(&mut self, data: &mut [u8]) {
         for byte in data.iter_mut() {
             if self.p == 16 {
-                let temp = self.state.clone();
-                crypto::aessafe::AesSafe128Encryptor::new(&self.key).encrypt_block(&temp, &mut self.state); // yes it's encrypt
+                crypto::aessafe::AesSafe128Encryptor::new(&self.key).encrypt_block(&self.state.clone(), &mut self.state); // yes it's encrypt
                 self.p = 0;
             }
             let temp = *byte;