From b66b5b6b9042c3817ebb6d426c5ecd523b271c32 Mon Sep 17 00:00:00 2001
From: mat <git@matdoes.dev>
Date: Tue, 9 Apr 2024 07:15:30 +0000
Subject: [PATCH] add functions to ClientBuilder and SwarmBuilder for custom
 addresses

---
 azalea/src/lib.rs       | 21 +++++++++++++++++++++
 azalea/src/swarm/mod.rs | 42 ++++++++++++++++++++++++++++++++++++++---
 2 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/azalea/src/lib.rs b/azalea/src/lib.rs
index 84c215d54..6e18ff7d4 100644
--- a/azalea/src/lib.rs
+++ b/azalea/src/lib.rs
@@ -15,6 +15,8 @@ pub mod pathfinder;
 pub mod prelude;
 pub mod swarm;
 
+use std::net::SocketAddr;
+
 use app::Plugins;
 pub use azalea_auth as auth;
 pub use azalea_block as blocks;
@@ -189,6 +191,25 @@ where
         }
         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(
+        mut self,
+        account: Account,
+        address: impl TryInto<ServerAddress>,
+        resolved_address: SocketAddr,
+    ) -> Result<!, StartError> {
+        self.swarm.accounts = vec![account];
+        if self.swarm.states.is_empty() {
+            self.swarm.states = vec![S::default()];
+        }
+        self.swarm
+            .start_with_custom_resolved_address(address, resolved_address)
+            .await
+    }
 }
 impl Default for ClientBuilder<NoState> {
     fn default() -> Self {
diff --git a/azalea/src/swarm/mod.rs b/azalea/src/swarm/mod.rs
index 6d3885ef1..2be56567e 100644
--- a/azalea/src/swarm/mod.rs
+++ b/azalea/src/swarm/mod.rs
@@ -296,6 +296,28 @@ where
     ///
     /// [`ServerAddress`]: azalea_protocol::ServerAddress
     pub async fn start(self, address: impl TryInto<ServerAddress>) -> Result<!, StartError> {
+        // convert the TryInto<ServerAddress> into a ServerAddress
+        let address: ServerAddress = match address.try_into() {
+            Ok(address) => address,
+            Err(_) => return Err(StartError::InvalidAddress),
+        };
+
+        // resolve the address
+        let resolved_address = resolver::resolve_address(&address).await?;
+
+        self.start_with_custom_resolved_address(address, resolved_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(
+        self,
+        address: impl TryInto<ServerAddress>,
+        resolved_address: SocketAddr,
+    ) -> Result<!, StartError> {
         assert_eq!(
             self.accounts.len(),
             self.states.len(),
@@ -308,9 +330,6 @@ where
             Err(_) => return Err(StartError::InvalidAddress),
         };
 
-        // resolve the address
-        let resolved_address = resolver::resolve_address(&address).await?;
-
         let instance_container = Arc::new(RwLock::new(InstanceContainer::default()));
 
         // we can't modify the swarm plugins after this
@@ -528,6 +547,23 @@ impl Swarm {
         let address = self.address.read().clone();
         let resolved_address = *self.resolved_address.read();
 
+        self.add_with_custom_address(account, state, address, resolved_address)
+            .await
+    }
+    /// Add a new account to the swarm, using the given host and socket
+    /// address. This is useful if you want bots in the same swarm to connect to
+    /// different addresses. Usually you'll just want [`Self::add`] though.
+    ///
+    /// # Errors
+    ///
+    /// Returns an `Err` if the bot could not do a handshake successfully.
+    pub async fn add_with_custom_address<S: Component + Clone>(
+        &mut self,
+        account: &Account,
+        state: S,
+        address: ServerAddress,
+        resolved_address: SocketAddr,
+    ) -> Result<Client, JoinError> {
         let (bot, mut rx) = Client::start_client(
             self.ecs_lock.clone(),
             account,