socks5 support (#113)

azalea-client/Cargo.toml

@@ -43,6 +43,7 @@ azalea-entity = { version = "0.9.0", path = "../azalea-entity" }
 serde_json = "1.0.113"
 serde = "1.0.196"
 minecraft_folder_path = "0.1.2"
+socks5-impl = "0.5.6"
 
 [features]
 default = ["log"]

azalea-client/src/client.rs

@@ -34,7 +34,7 @@ use azalea_entity::{
 };
 use azalea_physics::PhysicsPlugin;
 use azalea_protocol::{
-    connect::{Connection, ConnectionError},
+    connect::{Connection, ConnectionError, Proxy},
     packets::{
         configuration::{
             serverbound_client_information_packet::ClientInformation,
@@ -183,6 +183,7 @@ impl Client {
     pub async fn join(
         account: &Account,
         address: impl TryInto<ServerAddress>,
+        proxy: Option<Proxy>,
     ) -> Result<(Self, mpsc::UnboundedReceiver<Event>), JoinError> {
         let address: ServerAddress = address.try_into().map_err(|_| JoinError::InvalidAddress)?;
         let resolved_address = resolver::resolve_address(&address).await?;
@@ -200,6 +201,7 @@ impl Client {
             account,
             &address,
             &resolved_address,
+            proxy,
             run_schedule_sender,
         )
         .await
@@ -212,6 +214,7 @@ impl Client {
         account: &Account,
         address: &ServerAddress,
         resolved_address: &SocketAddr,
+        proxy: Option<Proxy>,
         run_schedule_sender: mpsc::UnboundedSender<()>,
     ) -> Result<(Self, mpsc::UnboundedReceiver<Event>), JoinError> {
         // check if an entity with our uuid already exists in the ecs and if so then
@@ -239,7 +242,11 @@ impl Client {
             entity
         };
 
-        let conn = Connection::new(resolved_address).await?;
+        let conn = if let Some(proxy) = proxy {
+            Connection::new_with_proxy(resolved_address, proxy).await?
+        } else {
+            Connection::new(resolved_address).await?
+        };
         let (conn, game_profile) =
             Self::handshake(ecs_lock.clone(), entity, conn, account, address).await?;
 

azalea-client/src/ping.rs

@@ -1,9 +1,12 @@
 //! Ping Minecraft servers.
 
 use azalea_protocol::{
-    connect::{Connection, ConnectionError},
+    connect::{Connection, ConnectionError, Proxy},
     packets::{
-        handshaking::client_intention_packet::ClientIntentionPacket,
+        handshaking::{
+            client_intention_packet::ClientIntentionPacket, ClientboundHandshakePacket,
+            ServerboundHandshakePacket,
+        },
         status::{
             clientbound_status_response_packet::ClientboundStatusResponsePacket,
             serverbound_status_request_packet::ServerboundStatusRequestPacket,
@@ -47,11 +50,29 @@ pub async fn ping_server(
     address: impl TryInto<ServerAddress>,
 ) -> Result<ClientboundStatusResponsePacket, PingError> {
     let address: ServerAddress = address.try_into().map_err(|_| PingError::InvalidAddress)?;
-
     let resolved_address = resolver::resolve_address(&address).await?;
+    let conn = Connection::new(&resolved_address).await?;
+    ping_server_with_connection(address, conn).await
+}
 
-    let mut conn = Connection::new(&resolved_address).await?;
+/// Ping a Minecraft server through a Socks5 proxy.
+pub async fn ping_server_with_proxy(
+    address: impl TryInto<ServerAddress>,
+    proxy: Proxy,
+) -> Result<ClientboundStatusResponsePacket, PingError> {
+    let address: ServerAddress = address.try_into().map_err(|_| PingError::InvalidAddress)?;
+    let resolved_address = resolver::resolve_address(&address).await?;
+    let conn = Connection::new_with_proxy(&resolved_address, proxy).await?;
+    ping_server_with_connection(address, conn).await
+}
 
+/// Ping a Minecraft server after we've already created a [`Connection`]. The
+/// `Connection` must still be in the handshake state (which is the state it's
+/// in immediately after it's created).
+pub async fn ping_server_with_connection(
+    address: ServerAddress,
+    mut conn: Connection<ClientboundHandshakePacket, ServerboundHandshakePacket>,
+) -> Result<ClientboundStatusResponsePacket, PingError> {
     // send the client intention packet and switch to the status state
     conn.write(
         ClientIntentionPacket {

azalea-protocol/Cargo.toml

@@ -48,6 +48,8 @@ trust-dns-resolver = { version = "^0.23.2", default-features = false, features =
 uuid = "1.7.0"
 log = "0.4.20"
 
+socks5-impl = "0.5.6"
+
 [features]
 connecting = []
 default = ["packets"]

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 {

azalea/examples/testbot/main.rs

@@ -181,10 +181,12 @@ async fn swarm_handle(
     _state: SwarmState,
 ) -> anyhow::Result<()> {
     match &event {
-        SwarmEvent::Disconnect(account) => {
+        SwarmEvent::Disconnect(account, join_opts) => {
             println!("bot got kicked! {}", account.username);
             tokio::time::sleep(Duration::from_secs(5)).await;
-            swarm.add_and_retry_forever(account, State::default()).await;
+            swarm
+                .add_and_retry_forever_with_opts(account, State::default(), join_opts)
+                .await;
         }
         SwarmEvent::Chat(chat) => {
             if chat.message().to_string() == "The particle was not visible for anybody" {

azalea/src/lib.rs

@@ -38,6 +38,7 @@ pub use azalea_world as world;
 pub use bot::*;
 use ecs::component::Component;
 use futures::{future::BoxFuture, Future};
+use protocol::connect::Proxy;
 use protocol::{resolver::ResolverError, ServerAddress};
 use swarm::SwarmBuilder;
 use thiserror::Error;
@@ -185,30 +186,26 @@ where
         account: Account,
         address: impl TryInto<ServerAddress>,
     ) -> Result<!, StartError> {
-        self.swarm.accounts = vec![account];
+        self.swarm.accounts = vec![(account, JoinOpts::default())];
         if self.swarm.states.is_empty() {
             self.swarm.states = vec![S::default()];
         }
         self.swarm.start(address).await
     }
 
-    /// Do the same as [`Self::start`], but allow passing in a custom resolved
-    /// address. This is useful if the address you're connecting to doesn't
-    /// resolve to anything, like if the server uses the address field to pass
-    /// custom data (like Bungeecord or Forge).
-    pub async fn start_with_custom_resolved_address(
+    /// Do the same as [`Self::start`], but allow passing in custom join
+    /// options.
+    pub async fn start_with_opts(
         mut self,
         account: Account,
         address: impl TryInto<ServerAddress>,
-        resolved_address: SocketAddr,
+        opts: JoinOpts,
     ) -> Result<!, StartError> {
-        self.swarm.accounts = vec![account];
+        self.swarm.accounts = vec![(account, opts.clone())];
         if self.swarm.states.is_empty() {
             self.swarm.states = vec![S::default()];
         }
-        self.swarm
-            .start_with_custom_resolved_address(address, resolved_address)
-            .await
+        self.swarm.start_with_default_opts(address, opts).await
     }
 }
 impl Default for ClientBuilder<NoState> {
@@ -224,3 +221,55 @@ impl Default for ClientBuilder<NoState> {
 /// [`SwarmBuilder`]: swarm::SwarmBuilder
 #[derive(Component, Clone, Default)]
 pub struct NoState;
+
+/// Optional settings when adding an account to a swarm or client.
+#[derive(Clone, Debug, Default)]
+#[non_exhaustive]
+pub struct JoinOpts {
+    /// The Socks5 proxy that this bot will use.
+    pub proxy: Option<Proxy>,
+    /// Override the server address that this specific bot will send in the
+    /// handshake packet.
+    pub custom_address: Option<ServerAddress>,
+    /// Override the socket address that this specific bot will use to connect
+    /// to the server.
+    pub custom_resolved_address: Option<SocketAddr>,
+}
+
+impl JoinOpts {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn update(&mut self, other: &Self) {
+        if let Some(proxy) = other.proxy.clone() {
+            self.proxy = Some(proxy);
+        }
+        if let Some(custom_address) = other.custom_address.clone() {
+            self.custom_address = Some(custom_address);
+        }
+        if let Some(custom_resolved_address) = other.custom_resolved_address {
+            self.custom_resolved_address = Some(custom_resolved_address);
+        }
+    }
+
+    /// Set the proxy that this bot will use.
+    #[must_use]
+    pub fn proxy(mut self, proxy: Proxy) -> Self {
+        self.proxy = Some(proxy);
+        self
+    }
+    /// Set the custom address that this bot will send in the handshake packet.
+    #[must_use]
+    pub fn custom_address(mut self, custom_address: ServerAddress) -> Self {
+        self.custom_address = Some(custom_address);
+        self
+    }
+    /// Set the custom resolved address that this bot will use to connect to the
+    /// server.
+    #[must_use]
+    pub fn custom_resolved_address(mut self, custom_resolved_address: SocketAddr) -> Self {
+        self.custom_resolved_address = Some(custom_resolved_address);
+        self
+    }
+}