1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

use std::time::Duration;
use std::io;
use std::str;
use std::net::{Ipv4Addr, Ipv6Addr};

use futures::future;
use futures::Future;
use tokio_io::io::{read_exact, write_all};
use tokio_core::net::TcpStream;
use tokio_timer::Timer;

mod v5 {
    pub const VERSION: u8 = 5;
    pub const METH_NO_AUTH: u8 = 0;
    pub const CMD_CONNECT: u8 = 1;
    pub const TYPE_IPV4: u8 = 1;
    pub const TYPE_IPV6: u8 = 4;
    pub const TYPE_DOMAIN: u8 = 3;
}

/// Creates a future which will handle socks5 connection.
///
/// if success The returned future will the handled `TcpStream` and address. if handle shake fail will
/// return `io::Error`.
pub fn serve(conn: TcpStream) -> Box<Future<Item = (TcpStream, String, u16), Error = io::Error>> {
    // socks version, only support version 5.
    let version = read_exact(conn, [0u8; 2]).and_then(|(conn, buf)| if buf[0] == v5::VERSION {
        Ok((conn, buf))
    } else {
        Err(other("unknown version"))
    });

    // ignore socks method
    let method = version.and_then(|(conn, buf)| read_exact(conn, [0u8, buf[1] as u8]));

    // send confirmation: version 5, no authentication required
    let part1 = method.and_then(|(conn, _)| write_all(conn, [v5::VERSION, v5::METH_NO_AUTH]));

    // check version
    let ack = part1.and_then(|(conn, _)| {
        read_exact(conn, [0u8]).and_then(|(conn, buf)| if buf[0] == v5::VERSION {
            Ok(conn)
        } else {
            Err(other("didn't confirm with v5 version"))
        })
    });
    // checkout cmd
    let command = ack.and_then(|conn| {
        read_exact(conn, [0u8]).and_then(|(conn, buf)| if buf[0] == v5::CMD_CONNECT {
            Ok(conn)
        } else {
            Err(other("unsupported command"))
        })
    });

    // there's one byte which is reserved for future use, so we read it and discard it.
    let resv = command.and_then(|c| read_exact(c, [0u8]));
    let adress_type = resv.and_then(|(conn, _)| read_exact(conn, [0u8]));
    let address = mybox(adress_type.and_then(move |(c, buf)| {
        match buf[0] {
            // For IPv4 addresses, we read the 4 bytes for the address as
            // well as 2 bytes for the port.
            v5::TYPE_IPV4 => mybox(read_exact(c, [0u8; 6]).and_then(|(c, buf)| {
                let addr = Ipv4Addr::new(buf[0], buf[1], buf[2], buf[3]);
                let port = ((buf[4] as u16) << 8) | (buf[5] as u16);
                mybox(future::ok((c, format!("{}", addr), port)))
            })),

            // For IPv6 addresses there's 16 bytes of an address plus two
            // bytes for a port, so we read that off and then keep going.
            v5::TYPE_IPV6 => mybox(read_exact(c, [0u8; 18]).and_then(|(conn, buf)| {
                let a = ((buf[0] as u16) << 8) | (buf[1] as u16);
                let b = ((buf[2] as u16) << 8) | (buf[3] as u16);
                let c = ((buf[4] as u16) << 8) | (buf[5] as u16);
                let d = ((buf[6] as u16) << 8) | (buf[7] as u16);
                let e = ((buf[8] as u16) << 8) | (buf[9] as u16);
                let f = ((buf[10] as u16) << 8) | (buf[11] as u16);
                let g = ((buf[12] as u16) << 8) | (buf[13] as u16);
                let h = ((buf[14] as u16) << 8) | (buf[15] as u16);
                let addr = Ipv6Addr::new(a, b, c, d, e, f, g, h);
                let port = ((buf[16] as u16) << 8) | (buf[17] as u16);
                mybox(future::ok((conn, format!("{}", addr), port)))
            })),

            // The SOCKSv5 protocol not only supports proxying to specific
            // IP addresses, but also arbitrary hostnames.
            v5::TYPE_DOMAIN => mybox(
                read_exact(c, [0u8])
                    .and_then(|(conn, buf)| {
                        read_exact(conn, vec![0u8; buf[0] as usize + 2])
                    })
                    .and_then(|(conn, buf)| {
                        let hostname = &buf[..buf.len() - 2];
                        let hostname = if let Ok(hostname) = str::from_utf8(hostname) {
                            hostname
                        } else {
                            return mybox(future::err(other("hostname include invalid utf8")));
                        };

                        let pos = buf.len() - 2;
                        let port = ((buf[pos] as u16) << 8) | (buf[pos + 1] as u16);
                        mybox(future::ok(
                            (conn, hostname.to_string(), port),
                        ))
                    }),
            ),
            n => {
                let msg = format!("unknown address type, received: {}", n);
                mybox(future::err(other(&msg)))
            }
        }
    }));

    // Sending connection established message immediately to client.
    // This some round trip time for creating socks connection with the client.
    // But if connection failed, the client will get connection reset error.
    let handshake_finish = mybox(address.and_then(move |(conn, addr, port)| {
        let resp = [0x05, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x08, 0x43];
        write_all(conn, resp).map(move |(conn, _)| (conn, addr, port))
    }));

    let timer = Timer::default();
    let timeout = timer
        .timeout(handshake_finish, Duration::new(10, 0))
        .map_err(|_| other("handshake timeout"));

    mybox(timeout.map(|(conn, addr, port)| (conn, addr, port)))
}

fn mybox<F: Future + 'static>(f: F) -> Box<Future<Item = F::Item, Error = F::Error>> {
    Box::new(f)
}


fn other(desc: &str) -> io::Error {
    io::Error::new(io::ErrorKind::Other, desc)
}