From 8f0d0d92808286e92b7c8aea704be3c00f5afce4 Mon Sep 17 00:00:00 2001 From: mat Date: Wed, 25 Dec 2024 07:27:09 +0000 Subject: [PATCH] close tcp connection on bot disconnect and add swarms to testbot cli --- azalea-client/src/disconnect.rs | 4 ++ azalea-protocol/src/read.rs | 8 ++-- azalea/examples/testbot/main.rs | 68 +++++++++++++++++++-------------- azalea/src/swarm/mod.rs | 8 ++-- 4 files changed, 51 insertions(+), 37 deletions(-) diff --git a/azalea-client/src/disconnect.rs b/azalea-client/src/disconnect.rs index ea8479d3..37bb37dd 100644 --- a/azalea-client/src/disconnect.rs +++ b/azalea-client/src/disconnect.rs @@ -48,8 +48,12 @@ pub fn remove_components_from_disconnected_players( commands .entity(*entity) .remove::() + // this makes it close the tcp connection + .remove::() // swarm detects when this tx gets dropped to fire SwarmEvent::Disconnect .remove::(); + // note that we don't remove the client from the ECS, so if they decide + // to reconnect they'll keep their state } } diff --git a/azalea-protocol/src/read.rs b/azalea-protocol/src/read.rs index 6f9b754a..c324740f 100755 --- a/azalea-protocol/src/read.rs +++ b/azalea-protocol/src/read.rs @@ -230,7 +230,7 @@ pub async fn read_packet( cipher: &mut Option, ) -> Result> where - R: AsyncRead + std::marker::Unpin + std::marker::Send + std::marker::Sync, + R: AsyncRead + Unpin + Send + Sync, { let raw_packet = read_raw_packet(stream, buffer, compression_threshold, cipher).await?; let packet = deserialize_packet(&mut Cursor::new(&raw_packet))?; @@ -265,7 +265,7 @@ pub async fn read_raw_packet( cipher: &mut Option, ) -> Result, Box> where - R: AsyncRead + std::marker::Unpin + std::marker::Send + std::marker::Sync, + R: AsyncRead + Unpin + Send + Sync, { loop { if let Some(buf) = read_raw_packet_from_buffer::(buffer, compression_threshold)? { @@ -284,7 +284,7 @@ pub fn try_read_raw_packet( cipher: &mut Option, ) -> Result>, Box> where - R: AsyncRead + std::marker::Unpin + std::marker::Send + std::marker::Sync, + R: AsyncRead + Unpin + Send + Sync, { loop { if let Some(buf) = read_raw_packet_from_buffer::(buffer, compression_threshold)? { @@ -355,7 +355,7 @@ pub fn read_raw_packet_from_buffer( compression_threshold: Option, ) -> Result>, Box> where - R: AsyncRead + std::marker::Unpin + std::marker::Send + std::marker::Sync, + R: AsyncRead + Unpin + Send + Sync, { let Some(mut buf) = frame_splitter(buffer).map_err(ReadPacketError::from)? else { // no full packet yet :( diff --git a/azalea/examples/testbot/main.rs b/azalea/examples/testbot/main.rs index 3d9f999e..a9ba978a 100644 --- a/azalea/examples/testbot/main.rs +++ b/azalea/examples/testbot/main.rs @@ -45,22 +45,26 @@ async fn main() { thread::spawn(deadlock_detection_thread); - let account = if args.name.contains('@') { - Account::microsoft(&args.name).await.unwrap() - } else { - Account::offline(&args.name) - }; + let join_address = args.server.clone(); - let mut commands = CommandDispatcher::new(); - register_commands(&mut commands); - - let join_address = args.address.clone(); - - let builder = SwarmBuilder::new(); - builder + let mut builder = SwarmBuilder::new() .set_handler(handle) - .set_swarm_handler(swarm_handle) - .add_account_with_state(account, State::new(args, commands)) + .set_swarm_handler(swarm_handle); + + for username_or_email in &args.accounts { + let account = if username_or_email.contains('@') { + Account::microsoft(&username_or_email).await.unwrap() + } else { + Account::offline(&username_or_email) + }; + + let mut commands = CommandDispatcher::new(); + register_commands(&mut commands); + + builder = builder.add_account_with_state(account, State::new(args.clone(), commands)); + } + + builder .join_delay(Duration::from_millis(100)) .start(join_address) .await @@ -138,7 +142,7 @@ async fn handle(bot: Client, event: azalea::Event, state: State) -> anyhow::Resu let (Some(username), content) = chat.split_sender_and_content() else { return Ok(()); }; - if username != state.args.owner { + if username != state.args.owner_username { return Ok(()); } @@ -211,29 +215,31 @@ async fn swarm_handle( #[derive(Debug, Clone, Default)] pub struct Args { - pub owner: String, - pub name: String, - pub address: String, + pub owner_username: String, + pub accounts: Vec, + pub server: String, pub pathfinder_debug_particles: bool, } fn parse_args() -> Args { - let mut owner_username = None; - let mut bot_username = None; - let mut address = None; + let mut owner_username = "admin".to_string(); + let mut accounts = Vec::new(); + let mut server = "localhost".to_string(); let mut pathfinder_debug_particles = false; let mut args = env::args().skip(1); while let Some(arg) = args.next() { match arg.as_str() { "--owner" | "-O" => { - owner_username = args.next(); + owner_username = args.next().expect("Missing owner username"); } - "--name" | "-N" => { - bot_username = args.next(); + "--account" | "-A" => { + for account in args.next().expect("Missing account").split(',') { + accounts.push(account.to_string()); + } } - "--address" | "-A" => { - address = args.next(); + "--server" | "-S" => { + server = args.next().expect("Missing server address"); } "--pathfinder-debug-particles" | "-P" => { pathfinder_debug_particles = true; @@ -245,10 +251,14 @@ fn parse_args() -> Args { } } + if accounts.is_empty() { + accounts.push("azalea".to_string()); + } + Args { - owner: owner_username.unwrap_or_else(|| "admin".to_string()), - name: bot_username.unwrap_or_else(|| "azalea".to_string()), - address: address.unwrap_or_else(|| "localhost".to_string()), + owner_username, + accounts, + server, pathfinder_debug_particles, } } diff --git a/azalea/src/swarm/mod.rs b/azalea/src/swarm/mod.rs index 5ae8f9a1..d0219406 100644 --- a/azalea/src/swarm/mod.rs +++ b/azalea/src/swarm/mod.rs @@ -415,7 +415,7 @@ where // SwarmBuilder (self) isn't Send so we have to take all the things we need out // of it - let mut swarm_clone = swarm.clone(); + let swarm_clone = swarm.clone(); let join_delay = self.join_delay; let accounts = self.accounts.clone(); let states = self.states.clone(); @@ -602,7 +602,7 @@ impl Swarm { /// /// Returns an `Err` if the bot could not do a handshake successfully. pub async fn add_with_opts( - &mut self, + &self, account: &Account, state: S, join_opts: &JoinOpts, @@ -663,7 +663,7 @@ impl Swarm { /// This does exponential backoff (though very limited), starting at 5 /// seconds and doubling up to 15 seconds. pub async fn add_and_retry_forever( - &mut self, + &self, account: &Account, state: S, ) -> Client { @@ -674,7 +674,7 @@ impl Swarm { /// Same as [`Self::add_and_retry_forever`], but allow passing custom join /// options. pub async fn add_and_retry_forever_with_opts( - &mut self, + &self, account: &Account, state: S, opts: &JoinOpts,