socks5x 0.1.3

use socks5x::Socks5Address;
use socks5x::server::{ClientHandler, ConnectionCreator, SocksSplitable};
use std::io;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::net::{TcpListener, TcpStream};

#[derive(Clone)]
enum DataDirection {
    Read,
    Write,
}

impl DataDirection {
    fn arrow(&self) -> &'static str {
        match self {
            DataDirection::Read => "->",
            DataDirection::Write => "<-",
        }
    }

    fn name(&self) -> &'static str {
        match self {
            DataDirection::Read => "READ",
            DataDirection::Write => "WRITE",
        }
    }
}

struct LoggingReader {
    inner: Box<dyn AsyncRead + Unpin + Send>,
    target: String,
    direction: DataDirection,
}

impl LoggingReader {
    fn new(
        inner: Box<dyn AsyncRead + Unpin + Send>,
        target: String,
        direction: DataDirection,
    ) -> Self {
        Self {
            inner,
            target,
            direction,
        }
    }
}

impl AsyncRead for LoggingReader {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let before_len = buf.filled().len();

        match Pin::new(&mut self.inner).poll_read(cx, buf) {
            Poll::Ready(Ok(())) => {
                let after_len = buf.filled().len();
                let bytes_read = after_len - before_len;

                if bytes_read > 0 {
                    let data = &buf.filled()[before_len..after_len];
                    log_data(&self.target, &self.direction, data);
                }

                Poll::Ready(Ok(()))
            }
            poll => poll,
        }
    }
}

struct LoggingWriter {
    inner: Box<dyn AsyncWrite + Unpin + Send>,
    target: String,
    direction: DataDirection,
}

impl LoggingWriter {
    fn new(
        inner: Box<dyn AsyncWrite + Unpin + Send>,
        target: String,
        direction: DataDirection,
    ) -> Self {
        Self {
            inner,
            target,
            direction,
        }
    }
}

impl AsyncWrite for LoggingWriter {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, io::Error>> {
        match Pin::new(&mut self.inner).poll_write(cx, buf) {
            Poll::Ready(Ok(bytes_written)) => {
                if bytes_written > 0 {
                    let data = &buf[..bytes_written];
                    log_data(&self.target, &self.direction, data);
                }

                Poll::Ready(Ok(bytes_written))
            }
            poll => poll,
        }
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        Pin::new(&mut self.inner).poll_flush(cx)
    }

    fn poll_shutdown(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), io::Error>> {
        Pin::new(&mut self.inner).poll_shutdown(cx)
    }
}

fn hex_dump(data: &[u8]) -> String {
    if data.is_empty() {
        return String::new();
    }

    let preview: String = data
        .iter()
        .take(16)
        .map(|b| format!("{:02x} ", b))
        .collect();

    let ascii: String = data
        .iter()
        .take(16)
        .map(|&b| if b >= 32 && b <= 126 { b as char } else { '.' })
        .collect();

    format!("{} {} (total {} bytes)", preview, ascii, data.len())
}

fn log_data(target: &str, direction: &DataDirection, data: &[u8]) {
    let arrow = direction.arrow();
    let direction_name = direction.name();
    let byte_count = data.len();

    println!(
        "[{}] {} {} {} bytes: {}",
        direction_name,
        target,
        arrow,
        byte_count,
        hex_dump(data)
    );
}

struct LoggingStream {
    inner: TcpStream,
    target: String,
}

impl SocksSplitable for LoggingStream {
    fn split_for_socks(
        self,
    ) -> io::Result<(
        Box<dyn AsyncRead + Unpin + Send>,
        Box<dyn AsyncWrite + Unpin + Send>,
    )> {
        println!("Splitting connection to {} for data logging", self.target);
        let (read, write) = self.inner.into_split();

        let logging_reader =
            LoggingReader::new(Box::new(read), self.target.clone(), DataDirection::Read);

        let logging_writer =
            LoggingWriter::new(Box::new(write), self.target.clone(), DataDirection::Write);

        Ok((Box::new(logging_reader), Box::new(logging_writer)))
    }
}
struct LoggedCreator;

impl ConnectionCreator for LoggedCreator {
    type Stream = LoggingStream;
    async fn create_stream(&self, address: Socks5Address, port: u16) -> io::Result<Self::Stream> {
        let target = format!("{}:{}", address, port);
        println!("🔗 Establishing connection to: {}", target);

        let stream = match TcpStream::connect(&target).await {
            Ok(stream) => {
                println!("Successfully connected to {}", target);
                stream
            }
            Err(e) => {
                println!("Failed to connect to {}: {}", target, e);
                return Err(e);
            }
        };

        Ok(LoggingStream {
            inner: stream,
            target,
        })
    }
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let listener = TcpListener::bind("127.0.0.1:1080").await?;
    println!("SOCKS5 proxy with data logging listening on 127.0.0.1:1080");
    let client_handler = Arc::new(ClientHandler::no_auth(LoggedCreator));
    loop {
        let (socket, addr) = listener.accept().await?;
        println!("New client connection from: {}", addr);
        let client_handler = client_handler.clone();

        tokio::spawn(async move {
            if let Err(e) = client_handler.handle(socket).await {
                eprintln!("Error handling client {}: {}", addr, e);
            } else {
                println!("Client {} disconnected", addr);
            }
        });
    }
}