#![feature(proc_macro, conservative_impl_trait, generators)]

extern crate arrayvec;
extern crate byteorder;
extern crate futures_await as futures;
extern crate tokio_core;
extern crate tokio_io;

use futures::prelude::*;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6, ToSocketAddrs};
use arrayvec::{ArrayString, ArrayVec};
use tokio_core::reactor::Handle;
use tokio_core::net::TcpStream;
use tokio_io::io::{read, read_exact, write_all};
use byteorder::{NetworkEndian, WriteBytesExt};
use std::io::{self, Read, Write};
use tokio_io::{AsyncRead, AsyncWrite};

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum TargetAddr {
    Ip(SocketAddr),
    Domain(ArrayString<[u8; 256]>, u16),
}

pub trait ToTargetAddr {
    fn to_target_addr(&self) -> io::Result<TargetAddr>;
}

impl ToTargetAddr for TargetAddr {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        Ok(*self)
    }
}

impl ToTargetAddr for SocketAddr {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        Ok(TargetAddr::Ip(*self))
    }
}

impl ToTargetAddr for SocketAddrV4 {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        SocketAddr::V4(*self).to_target_addr()
    }
}

impl ToTargetAddr for SocketAddrV6 {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        SocketAddr::V6(*self).to_target_addr()
    }
}

impl ToTargetAddr for (Ipv4Addr, u16) {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        SocketAddrV4::new(self.0, self.1).to_target_addr()
    }
}

impl ToTargetAddr for (Ipv6Addr, u16) {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        SocketAddrV6::new(self.0, self.1, 0, 0).to_target_addr()
    }
}

impl<'a> ToTargetAddr for (&'a str, u16) {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        // try to parse as an IP first
        if let Ok(addr) = self.0.parse::<Ipv4Addr>() {
            return (addr, self.1).to_target_addr();
        }

        if let Ok(addr) = self.0.parse::<Ipv6Addr>() {
            return (addr, self.1).to_target_addr();
        }

        Ok(TargetAddr::Domain(
            ArrayString::from(self.0)
                .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "domain name too long"))?,
            self.1,
        ))
    }
}

impl<'a> ToTargetAddr for &'a str {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        // try to parse as an IP first
        if let Ok(addr) = self.parse::<SocketAddrV4>() {
            return addr.to_target_addr();
        }

        if let Ok(addr) = self.parse::<SocketAddrV6>() {
            return addr.to_target_addr();
        }

        // split the string by ':' and convert the second part to u16
        let mut parts_iter = self.rsplitn(2, ':');
        let port_str = match parts_iter.next() {
            Some(s) => s,
            None => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "invalid socket address",
                ))
            }
        };

        let host = match parts_iter.next() {
            Some(s) => s,
            None => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "invalid socket address",
                ))
            }
        };

        let port: u16 = match port_str.parse() {
            Ok(p) => p,
            Err(_) => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "invalid port value",
                ))
            }
        };

        (host, port).to_target_addr()
    }
}

#[derive(Debug)]
pub struct SocksStream {
    stream: TcpStream,
}

#[repr(u8)]
enum Version {
    Five = 5u8,
}

#[repr(u8)]
enum AuthType {
    None = 0x00u8,
    Invalid = 0xffu8,
}

#[repr(u8)]
enum Command {
    Connect = 1u8,
}

#[repr(u8)]
enum AddrType {
    Ipv4 = 1u8,
    Domain = 3u8,
    Ipv6 = 4u8,
}

#[repr(u8)]
enum ReplyStatus {
    Success = 0u8,
}

impl SocksStream {
    #[async]
    pub fn connect<T: 'static>(
        handle: Handle,
        proxy: T,
        target_addr: TargetAddr,
    ) -> io::Result<SocksStream>
    where
        T: ToSocketAddrs,
    {
        let mut stream = await!(TcpStream::connect(
            &proxy.to_socket_addrs()?.next().unwrap(),
            &handle
        ))?;

        let mut buf = ArrayVec::<[u8; 384]>::new();
        buf.write_u8(Version::Five as u8)?;
        buf.write_u8(1)?; // number of methods
        buf.write_u8(AuthType::None as u8)?;

        let (new_stream, mut buf) = await!(write_all(stream, buf))?;
        stream = new_stream;

        let mut reply = [0; 2];
        let (new_stream, reply) = await!(read_exact(stream, reply))?;
        stream = new_stream;

        if reply[1] == AuthType::Invalid as u8 {
            return Err(io::Error::new(
                io::ErrorKind::Other,
                "No acceptable authentication methods",
            ));
        }

        buf.clear();
        buf.write_u8(Version::Five as u8)?;
        buf.write_u8(Command::Connect as u8)?;
        buf.write_u8(0)?; // reserved
        match target_addr {
            TargetAddr::Ip(socket_addr) => {
                match socket_addr.ip() {
                    IpAddr::V4(addr) => {
                        buf.write_u8(AddrType::Ipv4 as u8)?;
                        buf.write_u32::<NetworkEndian>(addr.into())?;
                    }
                    IpAddr::V6(addr) => {
                        buf.write_u8(AddrType::Ipv6 as u8)?;
                        for &segment in &addr.segments() {
                            buf.write_u16::<NetworkEndian>(segment)?;
                        }
                    }
                }
                buf.write_u16::<NetworkEndian>(socket_addr.port())?;
            }
            TargetAddr::Domain(domain, port) => {
                buf.write_u8(AddrType::Domain as u8)?;
                buf.write_u8(domain.len() as u8)?;
                buf.write_all(domain.as_bytes())?;
                buf.write_u16::<NetworkEndian>(port)?;
            }
        };

        let (new_stream, _) = await!(write_all(stream, buf))?;
        stream = new_stream;

        let mut reply = ArrayVec::<[u8; 128]>::new();
        for _ in 0..128 {
            reply.push(0);
        }
        let (new_stream, reply, _) = await!(read(stream, reply))?;
        stream = new_stream;

        if reply[1] != ReplyStatus::Success as u8 {
            return Err(io::Error::new(
                io::ErrorKind::Other,
                "General failure".to_owned(),
            ));
        }

        Ok(SocksStream { 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()
    }
}

impl AsyncRead for SocksStream {
    unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
        self.stream.prepare_uninitialized_buffer(buf)
    }
}

impl AsyncWrite for SocksStream {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        let stream = &mut self.stream;
        stream.shutdown()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use tokio_core::reactor::Core;
    use tokio_io::io::read_to_end;

    fn test<T: 'static>(target: T, version: &'static str) -> io::Result<()>
    where
        T: ToTargetAddr,
    {
        let mut core = Core::new()?;
        let future =
            SocksStream::connect(core.handle(), "127.0.0.1:9050", target.to_target_addr()?);

        let mut proxy = core.run(future)?;

        let (new_proxy, _) = core.run(write_all(
            proxy,
            format!("GET / HTTP/{}\n\n", version).as_bytes(),
        ))?;
        proxy = new_proxy;

        let response = Vec::new();
        let (_, response) = core.run(read_to_end(proxy, response))?;
        assert!(String::from_utf8_lossy(&response).starts_with(&format!("HTTP/{}", version)));

        Ok(())
    }

    #[test]
    fn ipv4() {
        test("216.58.216.238:80", "1.0").unwrap();
    }

    #[test]
    fn domain_port() {
        test(("52g5y5karruvc7bz.onion", 80), "1.1").unwrap();
    }

    #[test]
    fn domain() {
        test("52g5y5karruvc7bz.onion:80", "1.1").unwrap();
    }
}