socks5 support (#113)

1 files changed, 37 insertions, 1 deletions

diff --git a/azalea-protocol/src/connect.rs b/azalea-protocol/src/connect.rs
index 86b92693..cc1b71c1 100755
--- a/azalea-protocol/src/connect.rs
+++ b/azalea-protocol/src/connect.rs
@@ -20,7 +20,7 @@ use std::io::Cursor;
 use std::marker::PhantomData;
 use std::net::SocketAddr;
 use thiserror::Error;
-use tokio::io::AsyncWriteExt;
+use tokio::io::{AsyncWriteExt, BufStream};
 use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf, ReuniteError};
 use tokio::net::TcpStream;
 use tracing::{error, info};
@@ -257,6 +257,20 @@ pub enum ConnectionError {
     Io(#[from] std::io::Error),
 }
 
+use socks5_impl::protocol::UserKey;
+
+#[derive(Debug, Clone)]
+pub struct Proxy {
+    pub addr: SocketAddr,
+    pub auth: Option<UserKey>,
+}
+
+impl Proxy {
+    pub fn new(addr: SocketAddr, auth: Option<UserKey>) -> Self {
+        Self { addr, auth }
+    }
+}
+
 impl Connection<ClientboundHandshakePacket, ServerboundHandshakePacket> {
     /// Create a new connection to the given address.
     pub async fn new(address: &SocketAddr) -> Result<Self, ConnectionError> {
@@ -265,6 +279,28 @@ impl Connection<ClientboundHandshakePacket, ServerboundHandshakePacket> {
         // enable tcp_nodelay
         stream.set_nodelay(true)?;
 
+        Self::new_from_stream(stream).await
+    }
+
+    /// Create a new connection to the given address and Socks5 proxy. If you're
+    /// not using a proxy, use [`Self::new`] instead.
+    pub async fn new_with_proxy(
+        address: &SocketAddr,
+        proxy: Proxy,
+    ) -> Result<Self, ConnectionError> {
+        let proxy_stream = TcpStream::connect(proxy.addr).await?;
+        let mut stream = BufStream::new(proxy_stream);
+
+        let _ = socks5_impl::client::connect(&mut stream, address, proxy.auth)
+            .await
+            .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
+
+        Self::new_from_stream(stream.into_inner()).await
+    }
+
+    /// Create a new connection from an existing stream. Useful if you want to
+    /// set custom options on the stream. Otherwise, just use [`Self::new`].
+    pub async fn new_from_stream(stream: TcpStream) -> Result<Self, ConnectionError> {
         let (read_stream, write_stream) = stream.into_split();
 
         Ok(Connection {