CLI

src/main.rs

@@ -10,15 +10,43 @@ use std::io::prelude::*;
 
 // todo:
 // 1. use thread pool or async io
+// 2. less unwraps and better error handling
 
 fn main() {
-    let server = Socks5Server::new();
-    vmess(&server, Addr::V4([127,0,0,1]), 1080, [219, 131, 173, 224, 50, 114, 78, 197, 160, 203, 164, 175, 6, 31, 23, 48])
+    let args: Vec<_> = std::env::args().skip(1).collect();
+    println!("{:?}", args);
+    let args: Vec<_> = args.iter().map(|x| &x[..]).collect();
+    match args[..] {
+        ["plain"] | ["plain", _] => {
+            let port: u16 = args.get(1).map(|x| x.parse().unwrap()).unwrap_or(1080);
+            let server = Socks5Server::new(port);
+            plain(&server)
+        },
+        ["vmess", proxy, user_id] | ["vmess", proxy, user_id, _] => {
+            let port: u16 = args.get(4).map(|x| x.parse().unwrap()).unwrap_or(1080);
+            let server = Socks5Server::new(port);
+            vmess(&server, proxy.into(), parse_uid(user_id).unwrap())
+        },
+        _ => {
+            eprintln!("
+                Usage: v2socks plain [local_port=1080]
+                       v2socks vmess <server_addr>:<server_port> <userid> [local_port=1080]
+            ")
+        },
+    }
 }
 
-fn vmess(server: &Socks5Server, proxy_addr: Addr, proxy_port: u16, user_id: [u8; 16]) {
+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)
+}
+
+fn vmess(server: &Socks5Server, proxy: String, user_id: [u8; 16]) {
     let connect = Box::leak(Box::new(move |dest, port| {
-        let client = std::net::TcpStream::connect(format!("{}:{}", proxy_addr, proxy_port)).unwrap();
+        let client = std::net::TcpStream::connect(&proxy).unwrap();
         let local = client.local_addr().unwrap();
         let local_addr = match local.ip() {
             std::net::IpAddr::V4(x) => Addr::V4(x.octets()),
@@ -30,6 +58,7 @@ fn vmess(server: &Socks5Server, proxy_addr: Addr, proxy_port: u16, user_id: [u8;
         (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), stream: std::net::TcpStream| {
         let mut key = [0; 16];
         thread_rng().fill_bytes(&mut key);

src/socks.rs

@@ -21,11 +21,13 @@ macro_rules! close_on_error {
     }}
 }
 
-pub struct Socks5Server;
+pub struct Socks5Server {
+    port: u16
+}
 
 impl Socks5Server {
-    pub fn new() -> Socks5Server {
-        return Socks5Server
+    pub fn new(port: u16) -> Socks5Server {
+        Socks5Server { port }
     }
 
     pub fn listen<T>(&self, connect: &'static (impl Fn(Addr, u16) -> (Addr, u16, T) + Sync), pass: &'static (impl Fn(T, TcpStream) + Sync)) {

src/vmess.rs

@@ -126,6 +126,7 @@ impl VmessWriter<std::net::TcpStream> {
 }
 
 impl<W: Write> VmessWriter<W> {
+    #[allow(non_snake_case, non_upper_case_globals)]
     fn handshake(&mut self, user_id: [u8; 16], addr: Addr, port: u16, key: [u8; 16], IV: [u8; 16]) -> std::io::Result<()> {
         let time = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs().to_be_bytes();
         let mut hmac = crypto::hmac::Hmac::new(crypto::md5::Md5::new(), &user_id);