add Client::join_with_proxy and fix tests

azalea-rs · Apr 20, 2024 · 6d9d1a4 · 6d9d1a4
1 parent 353eda2
commit 6d9d1a4
Show file tree

Hide file tree

Showing 3 changed files with 77 additions and 41 deletions.
diff --git a/azalea-client/src/client.rs b/azalea-client/src/client.rs
@@ -138,6 +138,45 @@ pub enum JoinError {
     Disconnect { reason: FormattedText },
 }
 
+pub struct StartClientOpts<'a> {
+    pub ecs_lock: Arc<Mutex<World>>,
+    pub account: &'a Account,
+    pub address: &'a ServerAddress,
+    pub resolved_address: &'a SocketAddr,
+    pub proxy: Option<Proxy>,
+    pub run_schedule_sender: mpsc::UnboundedSender<()>,
+}
+
+impl<'a> StartClientOpts<'a> {
+    pub fn new(
+        account: &'a Account,
+        address: &'a ServerAddress,
+        resolved_address: &'a SocketAddr,
+    ) -> StartClientOpts<'a> {
+        // An event that causes the schedule to run. This is only used internally.
+        let (run_schedule_sender, run_schedule_receiver) = mpsc::unbounded_channel();
+
+        let mut app = App::new();
+        app.add_plugins(DefaultPlugins);
+
+        let ecs_lock = start_ecs_runner(app, run_schedule_receiver, run_schedule_sender.clone());
+
+        Self {
+            ecs_lock,
+            account,
+            address,
+            resolved_address,
+            proxy: None,
+            run_schedule_sender,
+        }
+    }
+
+    pub fn proxy(mut self, proxy: Proxy) -> Self {
+        self.proxy = Some(proxy);
+        self
+    }
+}
+
 impl Client {
     /// Create a new client from the given [`GameProfile`], ECS Entity, ECS
     /// World, and schedule runner function.
@@ -183,39 +222,36 @@ 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?;
 
-        // An event that causes the schedule to run. This is only used internally.
-        let (run_schedule_sender, run_schedule_receiver) = mpsc::unbounded_channel();
-
-        let mut app = App::new();
-        app.add_plugins(DefaultPlugins);
+        Self::start_client(StartClientOpts::new(account, &address, &resolved_address)).await
+    }
 
-        let ecs_lock = start_ecs_runner(app, run_schedule_receiver, run_schedule_sender.clone());
+    pub async fn join_with_proxy(
+        account: &Account,
+        address: impl TryInto<ServerAddress>,
+        proxy: 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?;
 
-        Self::start_client(
-            ecs_lock,
-            account,
-            &address,
-            &resolved_address,
-            proxy,
-            run_schedule_sender,
-        )
-        .await
+        Self::start_client(StartClientOpts::new(account, &address, &resolved_address).proxy(proxy))
+            .await
     }
 
     /// Create a [`Client`] when you already have the ECS made with
     /// [`start_ecs_runner`]. You'd usually want to use [`Self::join`] instead.
     pub async fn start_client(
-        ecs_lock: Arc<Mutex<World>>,
-        account: &Account,
-        address: &ServerAddress,
-        resolved_address: &SocketAddr,
-        proxy: Option<Proxy>,
-        run_schedule_sender: mpsc::UnboundedSender<()>,
+        StartClientOpts {
+            ecs_lock,
+            account,
+            address,
+            resolved_address,
+            proxy,
+            run_schedule_sender,
+        }: StartClientOpts<'_>,
     ) -> Result<(Self, mpsc::UnboundedReceiver<Event>), JoinError> {
         // check if an entity with our uuid already exists in the ecs and if so then
         // just use that

diff --git a/azalea-client/src/lib.rs b/azalea-client/src/lib.rs
@@ -32,7 +32,8 @@ pub mod task_pool;
 pub use account::{Account, AccountOpts};
 pub use azalea_protocol::packets::configuration::serverbound_client_information_packet::ClientInformation;
 pub use client::{
-    start_ecs_runner, Client, DefaultPlugins, JoinError, JoinedClientBundle, TickBroadcast,
+    start_ecs_runner, Client, DefaultPlugins, JoinError, JoinedClientBundle, StartClientOpts,
+    TickBroadcast,
 };
 pub use events::Event;
 pub use local_player::{GameProfileComponent, InstanceHolder, TabList};

diff --git a/azalea/src/swarm/mod.rs b/azalea/src/swarm/mod.rs
@@ -6,6 +6,7 @@ pub mod prelude;
 
 use azalea_client::{
     chat::ChatPacket, start_ecs_runner, Account, Client, DefaultPlugins, Event, JoinError,
+    StartClientOpts,
 };
 use azalea_protocol::{resolver, ServerAddress};
 use azalea_world::InstanceContainer;
@@ -535,10 +536,10 @@ pub type BoxSwarmHandleFn<SS> =
 ///     _state: SwarmState,
 /// ) -> anyhow::Result<()> {
 ///     match &event {
-///         SwarmEvent::Disconnect(account) => {
+///         SwarmEvent::Disconnect(account, join_opts) => {
 ///             // automatically reconnect after 5 seconds
 ///             tokio::time::sleep(Duration::from_secs(5)).await;
-///             swarm.add(account, State::default()).await?;
+///             swarm.add_with_opts(account, State::default(), join_opts).await?;
 ///         }
 ///         SwarmEvent::Chat(m) => {
 ///             println!("{}", m.message().to_ansi());
@@ -560,7 +561,7 @@ impl Swarm {
         account: &Account,
         state: S,
     ) -> Result<Client, JoinError> {
-        self.add_with_opts(account, state, JoinOpts::default())
+        self.add_with_opts(account, state, &JoinOpts::default())
             .await
     }
     /// Add a new account to the swarm, using custom options. This is useful if
@@ -574,24 +575,24 @@ impl Swarm {
         &mut self,
         account: &Account,
         state: S,
-        opts: JoinOpts,
+        join_opts: &JoinOpts,
     ) -> Result<Client, JoinError> {
-        let address = opts
+        let address = join_opts
             .custom_address
             .clone()
             .unwrap_or_else(|| self.address.read().clone());
-        let resolved_address = opts
+        let resolved_address = join_opts
             .custom_resolved_address
             .unwrap_or_else(|| *self.resolved_address.read());
 
-        let (bot, mut rx) = Client::start_client(
-            self.ecs_lock.clone(),
+        let (bot, mut rx) = Client::start_client(StartClientOpts {
+            ecs_lock: self.ecs_lock.clone(),
             account,
-            &address,
-            &resolved_address,
-            opts.proxy.clone(),
-            self.run_schedule_sender.clone(),
-        )
+            address: &address,
+            resolved_address: &resolved_address,
+            proxy: join_opts.proxy.clone(),
+            run_schedule_sender: self.run_schedule_sender.clone(),
+        })
         .await?;
         // add the state to the client
         {
@@ -605,6 +606,7 @@ impl Swarm {
         let cloned_bots_tx = self.bots_tx.clone();
         let cloned_bot = bot.clone();
         let swarm_tx = self.swarm_tx.clone();
+        let join_opts = join_opts.clone();
         tokio::spawn(async move {
             while let Some(event) = rx.recv().await {
                 // we can't handle events here (since we can't copy the handler),
@@ -618,7 +620,7 @@ impl Swarm {
                 .get_component::<Account>()
                 .expect("bot is missing required Account component");
             swarm_tx
-                .send(SwarmEvent::Disconnect(Box::new(account), opts))
+                .send(SwarmEvent::Disconnect(Box::new(account), join_opts))
                 .unwrap();
         });
 
@@ -649,10 +651,7 @@ impl Swarm {
     ) -> Client {
         let mut disconnects = 0;
         loop {
-            match self
-                .add_with_opts(account, state.clone(), opts.clone())
-                .await
-            {
+            match self.add_with_opts(account, state.clone(), opts).await {
                 Ok(bot) => return bot,
                 Err(e) => {
                     disconnects += 1;