1
2
Fork 0
mirror of https://github.com/mat-1/azalea.git synced 2025-08-02 14:26:04 +00:00

close tcp connection on bot disconnect and add swarms to testbot cli

This commit is contained in:
mat 2024-12-25 07:27:09 +00:00
parent 04eaa5c3d0
commit 8f0d0d9280
4 changed files with 51 additions and 37 deletions

View file

@ -48,8 +48,12 @@ pub fn remove_components_from_disconnected_players(
commands commands
.entity(*entity) .entity(*entity)
.remove::<JoinedClientBundle>() .remove::<JoinedClientBundle>()
// this makes it close the tcp connection
.remove::<RawConnection>()
// swarm detects when this tx gets dropped to fire SwarmEvent::Disconnect // swarm detects when this tx gets dropped to fire SwarmEvent::Disconnect
.remove::<LocalPlayerEvents>(); .remove::<LocalPlayerEvents>();
// note that we don't remove the client from the ECS, so if they decide
// to reconnect they'll keep their state
} }
} }

View file

@ -230,7 +230,7 @@ pub async fn read_packet<P: ProtocolPacket + Debug, R>(
cipher: &mut Option<Aes128CfbDec>, cipher: &mut Option<Aes128CfbDec>,
) -> Result<P, Box<ReadPacketError>> ) -> Result<P, Box<ReadPacketError>>
where 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 raw_packet = read_raw_packet(stream, buffer, compression_threshold, cipher).await?;
let packet = deserialize_packet(&mut Cursor::new(&raw_packet))?; let packet = deserialize_packet(&mut Cursor::new(&raw_packet))?;
@ -265,7 +265,7 @@ pub async fn read_raw_packet<R>(
cipher: &mut Option<Aes128CfbDec>, cipher: &mut Option<Aes128CfbDec>,
) -> Result<Box<[u8]>, Box<ReadPacketError>> ) -> Result<Box<[u8]>, Box<ReadPacketError>>
where where
R: AsyncRead + std::marker::Unpin + std::marker::Send + std::marker::Sync, R: AsyncRead + Unpin + Send + Sync,
{ {
loop { loop {
if let Some(buf) = read_raw_packet_from_buffer::<R>(buffer, compression_threshold)? { if let Some(buf) = read_raw_packet_from_buffer::<R>(buffer, compression_threshold)? {
@ -284,7 +284,7 @@ pub fn try_read_raw_packet<R>(
cipher: &mut Option<Aes128CfbDec>, cipher: &mut Option<Aes128CfbDec>,
) -> Result<Option<Box<[u8]>>, Box<ReadPacketError>> ) -> Result<Option<Box<[u8]>>, Box<ReadPacketError>>
where where
R: AsyncRead + std::marker::Unpin + std::marker::Send + std::marker::Sync, R: AsyncRead + Unpin + Send + Sync,
{ {
loop { loop {
if let Some(buf) = read_raw_packet_from_buffer::<R>(buffer, compression_threshold)? { if let Some(buf) = read_raw_packet_from_buffer::<R>(buffer, compression_threshold)? {
@ -355,7 +355,7 @@ pub fn read_raw_packet_from_buffer<R>(
compression_threshold: Option<u32>, compression_threshold: Option<u32>,
) -> Result<Option<Box<[u8]>>, Box<ReadPacketError>> ) -> Result<Option<Box<[u8]>>, Box<ReadPacketError>>
where 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 { let Some(mut buf) = frame_splitter(buffer).map_err(ReadPacketError::from)? else {
// no full packet yet :( // no full packet yet :(

View file

@ -45,22 +45,26 @@ async fn main() {
thread::spawn(deadlock_detection_thread); thread::spawn(deadlock_detection_thread);
let account = if args.name.contains('@') { let join_address = args.server.clone();
Account::microsoft(&args.name).await.unwrap()
let mut builder = SwarmBuilder::new()
.set_handler(handle)
.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 { } else {
Account::offline(&args.name) Account::offline(&username_or_email)
}; };
let mut commands = CommandDispatcher::new(); let mut commands = CommandDispatcher::new();
register_commands(&mut commands); register_commands(&mut commands);
let join_address = args.address.clone(); builder = builder.add_account_with_state(account, State::new(args.clone(), commands));
}
let builder = SwarmBuilder::new();
builder builder
.set_handler(handle)
.set_swarm_handler(swarm_handle)
.add_account_with_state(account, State::new(args, commands))
.join_delay(Duration::from_millis(100)) .join_delay(Duration::from_millis(100))
.start(join_address) .start(join_address)
.await .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 { let (Some(username), content) = chat.split_sender_and_content() else {
return Ok(()); return Ok(());
}; };
if username != state.args.owner { if username != state.args.owner_username {
return Ok(()); return Ok(());
} }
@ -211,29 +215,31 @@ async fn swarm_handle(
#[derive(Debug, Clone, Default)] #[derive(Debug, Clone, Default)]
pub struct Args { pub struct Args {
pub owner: String, pub owner_username: String,
pub name: String, pub accounts: Vec<String>,
pub address: String, pub server: String,
pub pathfinder_debug_particles: bool, pub pathfinder_debug_particles: bool,
} }
fn parse_args() -> Args { fn parse_args() -> Args {
let mut owner_username = None; let mut owner_username = "admin".to_string();
let mut bot_username = None; let mut accounts = Vec::new();
let mut address = None; let mut server = "localhost".to_string();
let mut pathfinder_debug_particles = false; let mut pathfinder_debug_particles = false;
let mut args = env::args().skip(1); let mut args = env::args().skip(1);
while let Some(arg) = args.next() { while let Some(arg) = args.next() {
match arg.as_str() { match arg.as_str() {
"--owner" | "-O" => { "--owner" | "-O" => {
owner_username = args.next(); owner_username = args.next().expect("Missing owner username");
} }
"--name" | "-N" => { "--account" | "-A" => {
bot_username = args.next(); 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" | "-P" => {
pathfinder_debug_particles = true; pathfinder_debug_particles = true;
@ -245,10 +251,14 @@ fn parse_args() -> Args {
} }
} }
if accounts.is_empty() {
accounts.push("azalea".to_string());
}
Args { Args {
owner: owner_username.unwrap_or_else(|| "admin".to_string()), owner_username,
name: bot_username.unwrap_or_else(|| "azalea".to_string()), accounts,
address: address.unwrap_or_else(|| "localhost".to_string()), server,
pathfinder_debug_particles, pathfinder_debug_particles,
} }
} }

View file

@ -415,7 +415,7 @@ where
// SwarmBuilder (self) isn't Send so we have to take all the things we need out // SwarmBuilder (self) isn't Send so we have to take all the things we need out
// of it // of it
let mut swarm_clone = swarm.clone(); let swarm_clone = swarm.clone();
let join_delay = self.join_delay; let join_delay = self.join_delay;
let accounts = self.accounts.clone(); let accounts = self.accounts.clone();
let states = self.states.clone(); let states = self.states.clone();
@ -602,7 +602,7 @@ impl Swarm {
/// ///
/// Returns an `Err` if the bot could not do a handshake successfully. /// Returns an `Err` if the bot could not do a handshake successfully.
pub async fn add_with_opts<S: Component + Clone>( pub async fn add_with_opts<S: Component + Clone>(
&mut self, &self,
account: &Account, account: &Account,
state: S, state: S,
join_opts: &JoinOpts, join_opts: &JoinOpts,
@ -663,7 +663,7 @@ impl Swarm {
/// This does exponential backoff (though very limited), starting at 5 /// This does exponential backoff (though very limited), starting at 5
/// seconds and doubling up to 15 seconds. /// seconds and doubling up to 15 seconds.
pub async fn add_and_retry_forever<S: Component + Clone>( pub async fn add_and_retry_forever<S: Component + Clone>(
&mut self, &self,
account: &Account, account: &Account,
state: S, state: S,
) -> Client { ) -> Client {
@ -674,7 +674,7 @@ impl Swarm {
/// Same as [`Self::add_and_retry_forever`], but allow passing custom join /// Same as [`Self::add_and_retry_forever`], but allow passing custom join
/// options. /// options.
pub async fn add_and_retry_forever_with_opts<S: Component + Clone>( pub async fn add_and_retry_forever_with_opts<S: Component + Clone>(
&mut self, &self,
account: &Account, account: &Account,
state: S, state: S,
opts: &JoinOpts, opts: &JoinOpts,