Add options to request Mojang sessionserver with a proxy (#293)

* add options to request mojang sessionserver with a socks5 proxy

* update changelog

* rename auth_proxy to sessionserver_proxy

12 files changed, 182 insertions, 54 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6953a9c8..bdc68c2e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -30,6 +30,7 @@ is breaking anyways, semantic versioning is not followed.
 - Rename `azalea_protocol::resolver` to `resolve` and `ResolverError` to `ResolveError`.
 - Refactor `RegistryHolder` to pre-deserialize some registries.
 - The handler function is now automatically single-threaded, making `#[tokio::main(flavor = "current_thread")]` unnecessary.
+- Mojang's sessionserver is now requested using the SOCKS5 proxy given in `JoinOpts::proxy`.
 
 ### Fixed
 
diff --git a/Cargo.lock b/Cargo.lock
index 380b9bcf..5f1ca525 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -498,6 +498,7 @@ dependencies = [
  "futures-lite",
  "hickory-resolver",
  "indexmap",
+ "reqwest",
  "serde",
  "serde_json",
  "simdnbt",
diff --git a/azalea-auth/src/sessionserver.rs b/azalea-auth/src/sessionserver.rs
index e8ab5de7..6c304942 100644
--- a/azalea-auth/src/sessionserver.rs
+++ b/azalea-auth/src/sessionserver.rs
@@ -51,28 +51,40 @@ pub struct ForbiddenError {
     pub path: String,
 }
 
-static REQWEST_CLIENT: LazyLock<reqwest::Client> = LazyLock::new(reqwest::Client::new);
+pub struct SessionServerJoinOpts<'a> {
+    pub access_token: &'a str,
+    /// Given to us by the
+    pub public_key: &'a [u8],
+    pub private_key: &'a [u8],
+    pub uuid: &'a Uuid,
+    /// This is given to us by the server, but it's typically an empty string.
+    pub server_id: &'a str,
+
+    pub proxy: Option<reqwest::Proxy>,
+}
 
 /// Tell Mojang's servers that you are going to join a multiplayer server,
 /// which is required to join online-mode servers.
-///
-/// The server ID should typically be an empty string.
-pub async fn join(
-    access_token: &str,
-    public_key: &[u8],
-    private_key: &[u8],
-    uuid: &Uuid,
-    server_id: &str,
-) -> Result<(), ClientSessionServerError> {
-    let client = REQWEST_CLIENT.clone();
+pub async fn join(opts: SessionServerJoinOpts<'_>) -> Result<(), ClientSessionServerError> {
+    let client = if let Some(proxy) = opts.proxy {
+        // reusing the client is too complicated if we're using proxies, so don't bother
+        reqwest::ClientBuilder::new().proxy(proxy).build()?
+    } else {
+        // no_proxy so we don't check reqwest's proxy env variables (because azalea
+        // doesn't handle them when connecting to servers anyways)
+        static REQWEST_CLIENT: LazyLock<reqwest::Client> =
+            LazyLock::new(|| reqwest::ClientBuilder::new().no_proxy().build().unwrap());
+
+        REQWEST_CLIENT.clone()
+    };
 
     let server_hash = azalea_crypto::hex_digest(&azalea_crypto::digest_data(
-        server_id.as_bytes(),
-        public_key,
-        private_key,
+        opts.server_id.as_bytes(),
+        opts.public_key,
+        opts.private_key,
     ));
 
-    join_with_server_id_hash(&client, access_token, uuid, &server_hash).await
+    join_with_server_id_hash(&client, opts.access_token, opts.uuid, &server_hash).await
 }
 
 pub async fn join_with_server_id_hash(
diff --git a/azalea-client/Cargo.toml b/azalea-client/Cargo.toml
index a237804a..f3c1bce5 100644
--- a/azalea-client/Cargo.toml
+++ b/azalea-client/Cargo.toml
@@ -32,7 +32,7 @@ minecraft_folder_path.workspace = true
 parking_lot.workspace = true
 pastey.workspace = true
 regex.workspace = true
-reqwest = { workspace = true, optional = true }
+reqwest = { workspace = true, optional = true, features = ["socks"] }
 simdnbt.workspace = true
 thiserror.workspace = true
 tokio = { workspace = true, features = ["sync"] }
diff --git a/azalea-client/src/client.rs b/azalea-client/src/client.rs
index 18a54125..970d8da0 100644
--- a/azalea-client/src/client.rs
+++ b/azalea-client/src/client.rs
@@ -121,14 +121,39 @@ impl StartClientOpts {
             connect_opts: ConnectOpts {
                 address,
                 resolved_address,
-                proxy: None,
+                server_proxy: None,
+                sessionserver_proxy: None,
             },
             event_sender,
         }
     }
 
-    pub fn proxy(mut self, proxy: Proxy) -> Self {
-        self.connect_opts.proxy = Some(proxy);
+    /// Configure the SOCKS5 proxy used for connecting to the server and for
+    /// authenticating with Mojang.
+    ///
+    /// To configure these separately, for example to only use the proxy for the
+    /// Minecraft server and not for authentication, you may use
+    /// [`Self::server_proxy`] and [`Self::sessionserver_proxy`] individually.
+    pub fn proxy(self, proxy: Proxy) -> Self {
+        self.server_proxy(proxy.clone()).sessionserver_proxy(proxy)
+    }
+    /// Configure the SOCKS5 proxy that will be used for connecting to the
+    /// Minecraft server.
+    ///
+    /// To avoid errors on servers with the "prevent-proxy-connections" option
+    /// set, you should usually use [`Self::proxy`] instead.
+    ///
+    /// Also see [`Self::sessionserver_proxy`].
+    pub fn server_proxy(mut self, proxy: Proxy) -> Self {
+        self.connect_opts.server_proxy = Some(proxy);
+        self
+    }
+    /// Configure the SOCKS5 proxy that this bot will use for authenticating the
+    /// server join with Mojang's API.
+    ///
+    /// Also see [`Self::proxy`] and [`Self::server_proxy`].
+    pub fn sessionserver_proxy(mut self, proxy: Proxy) -> Self {
+        self.connect_opts.sessionserver_proxy = Some(proxy);
         self
     }
 }
diff --git a/azalea-client/src/ping.rs b/azalea-client/src/ping.rs
index 93018a82..2c3ddfcf 100644
--- a/azalea-client/src/ping.rs
+++ b/azalea-client/src/ping.rs
@@ -56,7 +56,7 @@ pub async fn ping_server(
     ping_server_with_connection(address, conn).await
 }
 
-/// Ping a Minecraft server through a Socks5 proxy.
+/// Ping a Minecraft server through a SOCKS5 proxy.
 pub async fn ping_server_with_proxy(
     address: impl TryInto<ServerAddress>,
     proxy: Proxy,
diff --git a/azalea-client/src/plugins/join.rs b/azalea-client/src/plugins/join.rs
index a8f6e3c5..538369b0 100644
--- a/azalea-client/src/plugins/join.rs
+++ b/azalea-client/src/plugins/join.rs
@@ -64,9 +64,22 @@ pub struct StartJoinServerEvent {
 /// This is inserted as a component on clients to make auto-reconnecting work.
 #[derive(Debug, Clone, Component)]
 pub struct ConnectOpts {
+    /// The unresolved address that we're going to tell the server that we used
+    /// to connect.
     pub address: ServerAddress,
+    /// The actual IP and port that we're going to make a connection to.
     pub resolved_address: SocketAddr,
-    pub proxy: Option<Proxy>,
+    /// The SOCKS5 proxy used for connecting to the Minecraft server.
+    pub server_proxy: Option<Proxy>,
+    /// The SOCKS5 proxy that will be used when authenticating our server join
+    /// with Mojang.
+    ///
+    /// This should typically be either the same as [`Self::server_proxy`], or
+    /// `None`.
+    ///
+    /// This is useful to set if a server has `prevent-proxy-connections`
+    /// enabled.
+    pub sessionserver_proxy: Option<Proxy>,
 }
 
 /// An event that's sent when creating the TCP connection and sending the first
@@ -158,7 +171,7 @@ pub fn handle_start_join_server_event(
 async fn create_conn_and_send_intention_packet(
     opts: ConnectOpts,
 ) -> Result<LoginConn, ConnectionError> {
-    let mut conn = if let Some(proxy) = opts.proxy {
+    let mut conn = if let Some(proxy) = opts.server_proxy {
         Connection::new_with_proxy(&opts.resolved_address, proxy).await?
     } else {
         Connection::new(&opts.resolved_address).await?
diff --git a/azalea-client/src/plugins/login.rs b/azalea-client/src/plugins/login.rs
index 46ff8ea9..46ab20c4 100644
--- a/azalea-client/src/plugins/login.rs
+++ b/azalea-client/src/plugins/login.rs
@@ -1,19 +1,20 @@
 #[cfg(feature = "online-mode")]
 use azalea_auth::sessionserver::ClientSessionServerError;
-use azalea_protocol::packets::login::{
-    ClientboundHello, ServerboundCustomQueryAnswer, ServerboundKey,
+use azalea_protocol::{
+    connect::Proxy,
+    packets::login::{ClientboundHello, ServerboundCustomQueryAnswer, ServerboundKey},
 };
 use bevy_app::prelude::*;
 use bevy_ecs::prelude::*;
 use bevy_tasks::{IoTaskPool, Task, futures_lite::future};
 use thiserror::Error;
-use tracing::{debug, error, trace};
+use tracing::{debug, error, trace, warn};
 
 use super::{
     connection::RawConnection,
     packet::login::{ReceiveCustomQueryEvent, ReceiveHelloEvent, SendLoginPacketEvent},
 };
-use crate::Account;
+use crate::{Account, join::ConnectOpts};
 
 /// Some systems that run during the `login` state.
 pub struct LoginPlugin;
@@ -24,15 +25,29 @@ impl Plugin for LoginPlugin {
     }
 }
 
-fn handle_receive_hello_event(receive_hello: On<ReceiveHelloEvent>, mut commands: Commands) {
+fn handle_receive_hello_event(
+    receive_hello: On<ReceiveHelloEvent>,
+    mut commands: Commands,
+    query: Query<&ConnectOpts>,
+) {
     let task_pool = IoTaskPool::get();
 
     let account = receive_hello.account.clone();
     let packet = receive_hello.packet.clone();
-    let player = receive_hello.entity;
+    let client = receive_hello.entity;
+
+    // we store the auth proxy in the ConnectOpts component to make it easily
+    // configurable. that component should've definitely been inserted by now, but
+    // if it somehow wasn't then we should let the user know.
+    let connect_opts = if let Ok(opts) = query.get(client) {
+        opts.sessionserver_proxy.clone()
+    } else {
+        warn!("ConnectOpts component missing on a client ({client}) that got ReceiveHelloEvent");
+        None
+    };
 
-    let task = task_pool.spawn(auth_with_account(account, packet));
-    commands.entity(player).insert(AuthTask(task));
+    let task = task_pool.spawn(auth_with_account(account, packet, connect_opts));
+    commands.entity(client).insert(AuthTask(task));
 }
 
 /// A marker component on our clients that indicates that the server is
@@ -92,6 +107,7 @@ pub enum AuthWithAccountError {
 pub async fn auth_with_account(
     account: Account,
     packet: ClientboundHello,
+    proxy: Option<Proxy>,
 ) -> Result<(ServerboundKey, PrivateKey), AuthWithAccountError> {
     let Ok(encrypt_res) = azalea_crypto::encrypt(&packet.public_key, &packet.challenge) else {
         return Err(AuthWithAccountError::Encryption(packet));
@@ -116,22 +132,29 @@ pub async fn auth_with_account(
         // after too many
         let mut attempts: usize = 1;
 
+        let proxy = proxy.map(Proxy::into);
+
         while let Err(err) = {
+            use azalea_auth::sessionserver::{self, SessionServerJoinOpts};
+
             let access_token = access_token.lock().clone();
 
             let uuid = &account
                 .uuid
                 .expect("Uuid must be present if access token is present.");
 
+            let proxy = proxy.clone();
+
             // this is necessary since reqwest usually depends on tokio and we're using
             // `futures` here
-            async_compat::Compat::new(azalea_auth::sessionserver::join(
-                &access_token,
-                &packet.public_key,
-                &private_key,
+            async_compat::Compat::new(sessionserver::join(SessionServerJoinOpts {
+                access_token: &access_token,
+                public_key: &packet.public_key,
+                private_key: &private_key,
                 uuid,
-                &packet.server_id,
-            ))
+                server_id: &packet.server_id,
+                proxy,
+            }))
             .await
         } {
             if attempts >= 2 {
diff --git a/azalea-protocol/Cargo.toml b/azalea-protocol/Cargo.toml
index 42507901..1c299b49 100644
--- a/azalea-protocol/Cargo.toml
+++ b/azalea-protocol/Cargo.toml
@@ -41,12 +41,13 @@ tracing.workspace = true
 hickory-resolver = { workspace = true, features = ["tokio", "system-config"] }
 uuid.workspace = true
 indexmap.workspace = true
+reqwest = { workspace = true, optional = true, features = ["socks"] }
 
 [features]
 default = ["packets", "online-mode"]
 connecting = []
 packets = ["connecting", "dep:azalea-core"]
-online-mode = ["azalea-auth/online-mode"]
+online-mode = ["azalea-auth/online-mode", "dep:reqwest"]
 
 [lints]
 workspace = true
diff --git a/azalea-protocol/src/connect.rs b/azalea-protocol/src/connect.rs
index 556f0f1c..fd599ff8 100644
--- a/azalea-protocol/src/connect.rs
+++ b/azalea-protocol/src/connect.rs
@@ -277,7 +277,7 @@ pub enum ConnectionError {
 
 use socks5_impl::protocol::UserKey;
 
-/// An address and authentication method for connecting to a Socks5 proxy.
+/// An address and authentication method for connecting to a SOCKS5 proxy.
 #[derive(Debug, Clone)]
 pub struct Proxy {
     pub addr: SocketAddr,
@@ -299,6 +299,13 @@ impl Display for Proxy {
     }
 }
 
+impl From<Proxy> for reqwest::Proxy {
+    fn from(proxy: Proxy) -> Self {
+        reqwest::Proxy::all(proxy.to_string())
+            .expect("azalea proxies should not fail to parse as reqwest proxies")
+    }
+}
+
 impl Connection<ClientboundHandshakePacket, ServerboundHandshakePacket> {
     /// Create a new connection to the given address.
     pub async fn new(address: &SocketAddr) -> Result<Self, ConnectionError> {
@@ -310,7 +317,7 @@ impl Connection<ClientboundHandshakePacket, ServerboundHandshakePacket> {
         Self::new_from_stream(stream).await
     }
 
-    /// Create a new connection to the given address and Socks5 proxy.
+    /// 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(
@@ -443,7 +450,7 @@ impl Connection<ClientboundLoginPacket, ServerboundLoginPacket> {
     ///     ClientboundLoginPacket::Hello(p) => {
     ///         // tell Mojang we're joining the server & enable encryption
     ///         let e = azalea_crypto::encrypt(&p.public_key, &p.challenge).unwrap();
-    ///         conn.authenticate(&access_token, &profile.id, e.secret_key, &p)
+    ///         conn.authenticate(&access_token, &profile.id, e.secret_key, &p, None)
     ///             .await?;
     ///         conn.write(ServerboundKey {
     ///             key_bytes: e.encrypted_public_key,
@@ -464,14 +471,18 @@ impl Connection<ClientboundLoginPacket, ServerboundLoginPacket> {
         uuid: &Uuid,
         private_key: [u8; 16],
         packet: &ClientboundHello,
+        sessionserver_proxy: Option<Proxy>,
     ) -> Result<(), ClientSessionServerError> {
-        azalea_auth::sessionserver::join(
+        use azalea_auth::sessionserver::{self, SessionServerJoinOpts};
+
+        sessionserver::join(SessionServerJoinOpts {
             access_token,
-            &packet.public_key,
-            &private_key,
+            public_key: &packet.public_key,
+            private_key: &private_key,
             uuid,
-            &packet.server_id,
-        )
+            server_id: &packet.server_id,
+            proxy: sessionserver_proxy.map(Proxy::into),
+        })
         .await
     }
 }
diff --git a/azalea/src/lib.rs b/azalea/src/lib.rs
index daed8451..bf188791 100644
--- a/azalea/src/lib.rs
+++ b/azalea/src/lib.rs
@@ -264,8 +264,18 @@ pub struct NoState;
 #[derive(Clone, Debug, Default)]
 #[non_exhaustive]
 pub struct JoinOpts {
-    /// The Socks5 proxy that this bot will use.
-    pub proxy: Option<Proxy>,
+    /// The SOCKS5 proxy that this bot will use for connecting to the Minecraft
+    /// server.
+    pub server_proxy: Option<Proxy>,
+    /// The SOCKS5 proxy that will be used when authenticating the bot's join
+    /// with Mojang.
+    ///
+    /// This should typically be either the same as [`Self::server_proxy`] or
+    /// `None`.
+    ///
+    /// This is useful to set if a server has `prevent-proxy-connections`
+    /// enabled.
+    pub sessionserver_proxy: Option<Proxy>,
     /// Override the server address that this specific bot will send in the
     /// handshake packet.
     pub custom_address: Option<ServerAddress>,
@@ -280,8 +290,11 @@ impl JoinOpts {
     }
 
     pub fn update(&mut self, other: &Self) {
-        if let Some(proxy) = other.proxy.clone() {
-            self.proxy = Some(proxy);
+        if let Some(proxy) = other.server_proxy.clone() {
+            self.server_proxy = Some(proxy);
+        }
+        if let Some(proxy) = other.sessionserver_proxy.clone() {
+            self.sessionserver_proxy = Some(proxy);
         }
         if let Some(custom_address) = other.custom_address.clone() {
             self.custom_address = Some(custom_address);
@@ -291,12 +304,38 @@ impl JoinOpts {
         }
     }
 
-    /// Set the proxy that this bot will use.
+    /// Configure the SOCKS5 proxy used for connecting to the server and for
+    /// authenticating with Mojang.
+    ///
+    /// To configure these separately, for example to only use the proxy for the
+    /// Minecraft server and not for authentication, you may use
+    /// [`Self::server_proxy`] and [`Self::sessionserver_proxy`] individually.
+    #[must_use]
+    pub fn proxy(self, proxy: Proxy) -> Self {
+        self.server_proxy(proxy.clone()).sessionserver_proxy(proxy)
+    }
+    /// Configure the SOCKS5 proxy that will be used for connecting to the
+    /// Minecraft server.
+    ///
+    /// To avoid errors on servers with the "prevent-proxy-connections" option
+    /// set, you should usually use [`Self::proxy`] instead.
+    ///
+    /// Also see [`Self::sessionserver_proxy`].
     #[must_use]
-    pub fn proxy(mut self, proxy: Proxy) -> Self {
-        self.proxy = Some(proxy);
+    pub fn server_proxy(mut self, proxy: Proxy) -> Self {
+        self.server_proxy = Some(proxy);
         self
     }
+    /// Configure the SOCKS5 proxy that this bot will use for authenticating the
+    /// server join with Mojang's API.
+    ///
+    /// Also see [`Self::proxy`] and [`Self::server_proxy`].
+    #[must_use]
+    pub fn sessionserver_proxy(mut self, proxy: Proxy) -> Self {
+        self.sessionserver_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 {
diff --git a/azalea/src/swarm/mod.rs b/azalea/src/swarm/mod.rs
index a744e0bc..e77f3713 100644
--- a/azalea/src/swarm/mod.rs
+++ b/azalea/src/swarm/mod.rs
@@ -748,7 +748,8 @@ impl Swarm {
         let resolved_address = join_opts
             .custom_resolved_address
             .unwrap_or_else(|| *self.resolved_address.read());
-        let proxy = join_opts.proxy.clone();
+        let server_proxy = join_opts.server_proxy.clone();
+        let sessionserver_proxy = join_opts.sessionserver_proxy.clone();
 
         let (tx, rx) = mpsc::unbounded_channel();
 
@@ -758,7 +759,8 @@ impl Swarm {
             connect_opts: ConnectOpts {
                 address,
                 resolved_address,
-                proxy,
+                server_proxy,
+                sessionserver_proxy,
             },
             event_sender: Some(tx),
         })