async_proxies/proxies/socks4/

general.rs

use crate::proxy::AsyncProxy;
use async_trait::async_trait;

use tokio::io::{self, AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::time::timeout;

use std::net::SocketAddr;
use std::net::SocketAddrV4;
use std::fmt::{Formatter, Display, Error};
use std::time::Duration;

pub struct Socks4General;

pub enum ErrorKind {
    ConnectionFailed,
    RawStreamIOFailed(io::Error),
    GotBadBuffer,
    RequestDenied,
    IdentIsUnavailable,
    BadIdent,
    OperationTimeoutReached
}

#[repr(u8)]
pub enum Command {
    TcpConnectionEstablishment = 1,
    TcpPortBinding = 2
}

pub struct Timeouts {
    connection_timeout: Duration,
    write_timeout: Duration,
    read_timeout: Duration
}

pub struct ConnParams {
    // the address of a proxy
    socks_addr: SocketAddr,
    // the address of a destination service
    dest_addr: SocketAddrV4,
    ident: String,
    // the way to establish connection (bind a port/connect)
    command: Command,
    // just timeouts
    timeouts: Timeouts,
}

impl Timeouts {
    pub fn new(connection_timeout: Duration,
               write_timeout: Duration,
               read_timeout: Duration)
        -> Timeouts
    {
        Timeouts { connection_timeout, write_timeout, read_timeout }
    }
}

impl Display for ErrorKind {
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        match self {
            ErrorKind::ConnectionFailed => f.write_str("connection failed"),
            ErrorKind::RawStreamIOFailed(e) 
                => f.write_str(&format!("I/O error: {}", e)),
            ErrorKind::GotBadBuffer => f.write_str("bad buffer has been received"),
            ErrorKind::RequestDenied => f.write_str("request denied"),
            ErrorKind::IdentIsUnavailable => f.write_str("ident is unavailable"),
            ErrorKind::BadIdent => f.write_str("bad ident"),
            ErrorKind::OperationTimeoutReached => f.write_str("operation timeout reached")
        }
    }
}

impl ConnParams {
    pub fn new(socks_addr: SocketAddr, dest_addr: SocketAddrV4,
        ident: String, command: Command, timeouts: Timeouts)
        -> ConnParams
    {
        ConnParams { socks_addr, dest_addr, ident, command, timeouts }
    }
}

#[async_trait]
impl AsyncProxy for Socks4General {
    type OutputStream = TcpStream;
    type ErrorKind = ErrorKind;
    type ConnParams = ConnParams;

    async fn connect(&mut self, params: Self::ConnParams)
        -> Result<Self::OutputStream, Self::ErrorKind>
    {
        // Connecting to the proxy itself
        let future = TcpStream::connect(params.socks_addr);
        let future = timeout(params.timeouts.connection_timeout, future);
        // You also can see that this code template
        // (with creating a future and wrapping it with a timeout)
        // repeats many times in this code. The repetitions
        // will be ofc reduced, don't worry.
        let stream = future.await
                           .map_err(|_| ErrorKind::OperationTimeoutReached)?
                           .map_err(|_| ErrorKind::ConnectionFailed)?;
        
        // Buffer length in bytes is determined this way:
        //  (+1) for the number of the version of the socks protocol (4 in this case)
        //  (+1) for the command number (1 or 2)
        //  (+2) for port (in the network byte order)
        //  (+4) for the IPv4 address
        //  (+n) where `n` is the length of the given ident
        //  (+1) for the NULL-termination byte (0x00)
        let buffer_len = 1 + 1 + 2 + 4 + params.ident.len() + 1;
        let mut buffer = Vec::with_capacity(buffer_len);

        // the version of the socks protocol being used
        buffer.push(4);
        // the command number
        buffer.push(params.command as u8);
            
        let port_in_bytes = params.dest_addr.port().to_be_bytes();
        // the destination port
        buffer.extend_from_slice(&port_in_bytes[..]);

        let ipaddr_in_bytes = params.dest_addr.ip().octets();
        // the ipv4 address of the destination host
        buffer.extend_from_slice(&ipaddr_in_bytes[..]);

        // the given ident
        buffer.extend_from_slice(&params.ident.as_bytes());
        // and, finally, the NULL-termination byte
        buffer.push(0);
            
        let future = self.write_buffer(&buffer[..], stream);
        let future = timeout(params.timeouts.write_timeout, future);
        let mut stream = future.await
                               .map_err(|_| ErrorKind::OperationTimeoutReached)??;

        let future = stream.read(&mut buffer[..]);
        let future = timeout(params.timeouts.read_timeout, future);
        let read_bytes = future.await
                               .map_err(|_| ErrorKind::OperationTimeoutReached)?
                               .map_err(|e| ErrorKind::RawStreamIOFailed(e))?;

        if read_bytes < 2 {
            return Err(ErrorKind::GotBadBuffer)
        }

        match buffer[1] {
            0x5a => Ok(stream),
            0x5b => Err(ErrorKind::RequestDenied),
            0x5c => Err(ErrorKind::IdentIsUnavailable),
            0x5d => Err(ErrorKind::BadIdent),
            _ => Err(ErrorKind::GotBadBuffer)
        }
    }

    async fn write_buffer(&mut self, buffer: &[u8], mut stream: Self::OutputStream)
        -> Result<Self::OutputStream, Self::ErrorKind>
    {
        match stream.write_all(buffer).await {
            Ok(_) => Ok(stream),
            Err(e) => Err(ErrorKind::RawStreamIOFailed(e))
        }

        // also could be written as:
        //  stream.write_all(buffer)
        //         .await
        //         .map(|_| stream)
        //         .map_err(|e| ErrorKind::RawStreamIOFailed(e))
    }

    async fn read_buffer(&mut self, buffer: &mut [u8], mut stream: Self::OutputStream)
        -> Result<Self::OutputStream, Self::ErrorKind>
    {
        match stream.read(buffer).await {
            Ok(_) => Ok(stream),
            Err(e) => Err(ErrorKind::RawStreamIOFailed(e))
        }
    }

    async fn drop_stream(&mut self, _: Self::OutputStream)
        -> Result<(), Self::ErrorKind>
    {
        Ok(())
    }
}