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

use crate::addresses::ProxyAddress;
use crate::socks6::{self, Socks6Reply};
use crate::{Socks6Client, SocksHandler};
use anyhow::Result;
use async_trait::async_trait;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;

#[derive(Clone)]
pub struct Socks6Handler {
    static_links: Vec<ProxyAddress>,
}

impl Default for Socks6Handler {
    fn default() -> Self {
        Self::new(vec![])
    }
}

impl Socks6Handler {
    ///
    ///
    ///
    pub fn new(static_links: Vec<ProxyAddress>) -> Self {
        Socks6Handler { static_links }
    }
}

#[async_trait]
impl SocksHandler for Socks6Handler {
    ///
    ///
    ///
    async fn accept_request(
        &self,
        source: &mut TcpStream,
    ) -> Result<()> {
        let mut destination = self.setup(source).await?;

        // Start bidirectional copy, after this the connection closes.
        tokio::io::copy_bidirectional(source, &mut destination).await?;

        Ok(())
    }

    ///
    ///
    ///
    async fn refuse_request(
        &self,
        source: &mut TcpStream,
    ) -> Result<()> {
        // Notify source that the connection is refused.
        socks6::write_reply(source, Socks6Reply::ConnectionRefused).await?;

        Ok(())
    }

    ///
    ///
    ///
    async fn setup(
        &self,
        source: &mut TcpStream,
    ) -> Result<TcpStream> {
        // Receive SOCKS request, and allow unauthenticated access.
        let request = socks6::read_request(source).await?;
        socks6::write_no_authentication(source).await?;

        let destination = request.destination.to_string();
        let chain = request.chain(&self.static_links)?;

        let mut destination = if let Some(mut chain) = chain {
            if let Some(next) = chain.next_link() {
                let next = next.clone();

                let proxy_addr = format!("{}:{}", next.host, next.port);
                let client = Socks6Client::new(proxy_addr, next.credentials).await?;

                let (outgoing, _) = client.connect(destination, None, Some(chain.as_options())).await?;
                outgoing
            } else {
                TcpStream::connect(destination).await?
            }
        } else {
            TcpStream::connect(destination).await?
        };

        // Send initial data
        if request.initial_data_length > 0 {
            let mut initial_data = vec![0; request.initial_data_length as usize];
            source.read_exact(&mut initial_data).await?;
            destination.write(&initial_data).await?;
        }

        // Notify source that the connection has been set up.
        socks6::write_reply(source, Socks6Reply::Success).await?;
        source.flush().await?;

        Ok(destination)
    }
}