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

extern crate byteorder;

use byteorder::{WriteBytesExt, BigEndian};
use std::net::{TcpStream, IpAddr, ToSocketAddrs};
use std::io::{self, Read, Write};

pub struct SocksStream {
    stream: TcpStream
}

enum Version {
    Five = 5,
}

enum AuthType {
    None = 0x00,
    Invalid = 0xff
}

enum Command {
    Connect = 1,
}

enum AddrType {
    Ipv4 = 1,
    Ipv6 = 4,
}

enum ReplyStatus {
    Success = 0,
}

impl SocksStream {
    pub fn connect<T: ToSocketAddrs>(proxy: T, target: T) -> io::Result<SocksStream> {
        let target_addr = target.to_socket_addrs()?.next().unwrap();
        let mut stream = TcpStream::connect(proxy)?;

        let mut buf = Vec::new();
        let _ = buf.write_u8(Version::Five as u8);
        let _ = buf.write_u8(1); // number of methods
        let _ = buf.write_u8(AuthType::None as u8);
        let _ = stream.write_all(&buf);
        let mut reply = [0; 2];
        let _ = stream.read(&mut reply);
        if reply[1] == AuthType::Invalid as u8 {
            return Err(io::Error::new(io::ErrorKind::Other, "No acceptable methods"));
        }

        buf.clear();
        let _ = buf.write_u8(Version::Five as u8);
        let _ = buf.write_u8(Command::Connect as u8);
        let _ = buf.write_u8(0); // reserved
        match target_addr.ip() {
            IpAddr::V4(addr) => {
                let _ = buf.write_u8(AddrType::Ipv4 as u8);
                let _ = buf.write_u32::<BigEndian>(addr.into());
            },
            IpAddr::V6(addr) => {
                let _ = buf.write_u8(AddrType::Ipv6 as u8);
                for &segment in addr.segments().iter() {
                    let _ = buf.write_u16::<BigEndian>(segment);
                }
            },
        };
        let _ = buf.write_u16::<BigEndian>(target_addr.port());
        let _ = stream.write_all(&buf);
        let mut reply = [0; 128];
        let _ = stream.read(&mut reply);
        if reply[1] != ReplyStatus::Success as u8 {
            return Err(io::Error::new(io::ErrorKind::Other, format!("Failure: {}", reply[1])));
        }

        Ok(SocksStream {
            stream: stream
        })
    }
}

impl Read for SocksStream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.stream.read(buf)
    }
}

impl Write for SocksStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.stream.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.stream.flush()
    }
}

#[cfg(test)]
mod tests {
    use super::SocksStream;
    use std::io::{Read, Write};

    #[test]
    fn ipv4() {
        let mut proxy = SocksStream::connect("127.0.0.1:9050", "216.58.216.238:80").unwrap();
        let _ = proxy.write_all(b"GET / HTTP/1.0\n\n");
        let mut response = Vec::new();
        let _ = proxy.read_to_end(&mut response);
        assert!(String::from_utf8_lossy(&response).starts_with("HTTP/1.0"));
    }
}