diff options
Diffstat (limited to 'azalea')
| -rw-r--r-- | azalea/examples/testbot/main.rs | 6 | ||||
| -rw-r--r-- | azalea/src/lib.rs | 71 | ||||
| -rw-r--r-- | azalea/src/swarm/mod.rs | 124 |
3 files changed, 144 insertions, 57 deletions
diff --git a/azalea/examples/testbot/main.rs b/azalea/examples/testbot/main.rs index 86395b7e..6795e6cf 100644 --- a/azalea/examples/testbot/main.rs +++ b/azalea/examples/testbot/main.rs @@ -181,10 +181,12 @@ async fn swarm_handle( _state: SwarmState, ) -> anyhow::Result<()> { match &event { - SwarmEvent::Disconnect(account) => { + SwarmEvent::Disconnect(account, join_opts) => { println!("bot got kicked! {}", account.username); tokio::time::sleep(Duration::from_secs(5)).await; - swarm.add_and_retry_forever(account, State::default()).await; + swarm + .add_and_retry_forever_with_opts(account, State::default(), join_opts) + .await; } SwarmEvent::Chat(chat) => { if chat.message().to_string() == "The particle was not visible for anybody" { diff --git a/azalea/src/lib.rs b/azalea/src/lib.rs index 6e18ff7d..7d8b424c 100644 --- a/azalea/src/lib.rs +++ b/azalea/src/lib.rs @@ -38,6 +38,7 @@ pub use azalea_world as world; pub use bot::*; use ecs::component::Component; use futures::{future::BoxFuture, Future}; +use protocol::connect::Proxy; use protocol::{resolver::ResolverError, ServerAddress}; use swarm::SwarmBuilder; use thiserror::Error; @@ -185,30 +186,26 @@ where account: Account, address: impl TryInto<ServerAddress>, ) -> Result<!, StartError> { - self.swarm.accounts = vec![account]; + self.swarm.accounts = vec![(account, JoinOpts::default())]; if self.swarm.states.is_empty() { self.swarm.states = vec![S::default()]; } 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( + /// Do the same as [`Self::start`], but allow passing in custom join + /// options. + pub async fn start_with_opts( mut self, account: Account, address: impl TryInto<ServerAddress>, - resolved_address: SocketAddr, + opts: JoinOpts, ) -> Result<!, StartError> { - self.swarm.accounts = vec![account]; + self.swarm.accounts = vec![(account, opts.clone())]; if self.swarm.states.is_empty() { self.swarm.states = vec![S::default()]; } - self.swarm - .start_with_custom_resolved_address(address, resolved_address) - .await + self.swarm.start_with_default_opts(address, opts).await } } impl Default for ClientBuilder<NoState> { @@ -224,3 +221,55 @@ impl Default for ClientBuilder<NoState> { /// [`SwarmBuilder`]: swarm::SwarmBuilder #[derive(Component, Clone, Default)] pub struct NoState; + +/// Optional settings when adding an account to a swarm or client. +#[derive(Clone, Debug, Default)] +#[non_exhaustive] +pub struct JoinOpts { + /// The Socks5 proxy that this bot will use. + pub proxy: Option<Proxy>, + /// Override the server address that this specific bot will send in the + /// handshake packet. + pub custom_address: Option<ServerAddress>, + /// Override the socket address that this specific bot will use to connect + /// to the server. + pub custom_resolved_address: Option<SocketAddr>, +} + +impl JoinOpts { + pub fn new() -> Self { + Self::default() + } + + pub fn update(&mut self, other: &Self) { + if let Some(proxy) = other.proxy.clone() { + self.proxy = Some(proxy); + } + if let Some(custom_address) = other.custom_address.clone() { + self.custom_address = Some(custom_address); + } + if let Some(custom_resolved_address) = other.custom_resolved_address { + self.custom_resolved_address = Some(custom_resolved_address); + } + } + + /// Set the proxy that this bot will use. + #[must_use] + pub fn proxy(mut self, proxy: Proxy) -> Self { + self.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 { + self.custom_address = Some(custom_address); + self + } + /// Set the custom resolved address that this bot will use to connect to the + /// server. + #[must_use] + pub fn custom_resolved_address(mut self, custom_resolved_address: SocketAddr) -> Self { + self.custom_resolved_address = Some(custom_resolved_address); + self + } +} diff --git a/azalea/src/swarm/mod.rs b/azalea/src/swarm/mod.rs index 2be56567..a53e6fe8 100644 --- a/azalea/src/swarm/mod.rs +++ b/azalea/src/swarm/mod.rs @@ -17,7 +17,7 @@ use std::{collections::HashMap, future::Future, net::SocketAddr, sync::Arc, time use tokio::sync::mpsc; use tracing::error; -use crate::{BoxHandleFn, DefaultBotPlugins, HandleFn, NoState, StartError}; +use crate::{BoxHandleFn, DefaultBotPlugins, HandleFn, JoinOpts, NoState, StartError}; /// A swarm is a way to conveniently control many bots at once, while also /// being able to control bots at an individual level when desired. @@ -51,8 +51,8 @@ where SS: Default + Send + Sync + Clone + Resource + 'static, { pub(crate) app: App, - /// The accounts that are going to join the server. - pub(crate) accounts: Vec<Account>, + /// The accounts and proxies that are going to join the server. + pub(crate) accounts: Vec<(Account, JoinOpts)>, /// The individual bot states. This must be the same length as `accounts`, /// since each bot gets one state. pub(crate) states: Vec<S>, @@ -257,8 +257,20 @@ where /// Add an account with a custom initial state. Use just /// [`Self::add_account`] to use the Default implementation for the state. #[must_use] - pub fn add_account_with_state(mut self, account: Account, state: S) -> Self { - self.accounts.push(account); + pub fn add_account_with_state(self, account: Account, state: S) -> Self { + self.add_account_with_state_and_opts(account, state, JoinOpts::default()) + } + + /// Same as [`Self::add_account_with_state`], but allow passing in custom + /// join options. + #[must_use] + pub fn add_account_with_state_and_opts( + mut self, + account: Account, + state: S, + join_opts: JoinOpts, + ) -> Self { + self.accounts.push((account, join_opts)); self.states.push(state); self } @@ -302,21 +314,16 @@ where 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) + self.start_with_default_opts(address, JoinOpts::default()) .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( + /// Do the same as [`Self::start`], but allow passing in default join + /// options for the bots. + pub async fn start_with_default_opts( self, address: impl TryInto<ServerAddress>, - resolved_address: SocketAddr, + default_join_opts: JoinOpts, ) -> Result<!, StartError> { assert_eq!( self.accounts.len(), @@ -325,11 +332,17 @@ where ); // convert the TryInto<ServerAddress> into a ServerAddress - let address: ServerAddress = match address.try_into() { + let address = match address.try_into() { Ok(address) => address, Err(_) => return Err(StartError::InvalidAddress), }; + let address: ServerAddress = default_join_opts.custom_address.clone().unwrap_or(address); + let resolved_address: SocketAddr = match default_join_opts.custom_resolved_address { + Some(resolved_address) => resolved_address, + None => resolver::resolve_address(&address).await?, + }; + let instance_container = Arc::new(RwLock::new(InstanceContainer::default())); // we can't modify the swarm plugins after this @@ -378,24 +391,27 @@ where tokio::spawn(async move { if let Some(join_delay) = join_delay { // if there's a join delay, then join one by one - for (account, state) in accounts.iter().zip(states) { - swarm_clone.add_and_retry_forever(account, state).await; + for ((account, bot_join_opts), state) in accounts.iter().zip(states) { + let mut join_opts = default_join_opts.clone(); + join_opts.update(bot_join_opts); + swarm_clone + .add_and_retry_forever_with_opts(account, state, &join_opts) + .await; tokio::time::sleep(join_delay).await; } } else { // otherwise, join all at once let swarm_borrow = &swarm_clone; - join_all( - accounts - .iter() - .zip(states) - .map(move |(account, state)| async { - swarm_borrow - .clone() - .add_and_retry_forever(account, state) - .await; - }), - ) + join_all(accounts.iter().zip(states).map( + |((account, bot_join_opts), state)| async { + let mut join_opts = default_join_opts.clone(); + join_opts.update(bot_join_opts); + swarm_borrow + .clone() + .add_and_retry_forever_with_opts(account, state, &join_opts) + .await; + }, + )) .await; } @@ -460,9 +476,9 @@ pub enum SwarmEvent { Init, /// A bot got disconnected from the server. /// - /// You can implement an auto-reconnect by calling [`Swarm::add`] - /// with the account from this event. - Disconnect(Box<Account>), + /// You can implement an auto-reconnect by calling [`Swarm::add_with_opts`] + /// with the account and options from this event. + Disconnect(Box<Account>, JoinOpts), /// At least one bot received a chat message. Chat(ChatPacket), } @@ -544,31 +560,36 @@ impl Swarm { account: &Account, state: S, ) -> Result<Client, JoinError> { - let address = self.address.read().clone(); - let resolved_address = *self.resolved_address.read(); - - self.add_with_custom_address(account, state, address, resolved_address) + self.add_with_opts(account, state, JoinOpts::default()) .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. + /// Add a new account to the swarm, using custom options. 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>( + pub async fn add_with_opts<S: Component + Clone>( &mut self, account: &Account, state: S, - address: ServerAddress, - resolved_address: SocketAddr, + opts: JoinOpts, ) -> Result<Client, JoinError> { + let address = opts + .custom_address + .clone() + .unwrap_or_else(|| self.address.read().clone()); + let resolved_address = opts + .custom_resolved_address + .unwrap_or_else(|| *self.resolved_address.read()); + let (bot, mut rx) = Client::start_client( self.ecs_lock.clone(), account, &address, &resolved_address, + opts.proxy.clone(), self.run_schedule_sender.clone(), ) .await?; @@ -597,7 +618,7 @@ impl Swarm { .get_component::<Account>() .expect("bot is missing required Account component"); swarm_tx - .send(SwarmEvent::Disconnect(Box::new(account))) + .send(SwarmEvent::Disconnect(Box::new(account), opts)) .unwrap(); }); @@ -614,9 +635,24 @@ impl Swarm { account: &Account, state: S, ) -> Client { + self.add_and_retry_forever_with_opts(account, state, &JoinOpts::default()) + .await + } + + /// Same as [`Self::add_and_retry_forever`], but allow passing custom join + /// options. + pub async fn add_and_retry_forever_with_opts<S: Component + Clone>( + &mut self, + account: &Account, + state: S, + opts: &JoinOpts, + ) -> Client { let mut disconnects = 0; loop { - match self.add(account, state.clone()).await { + match self + .add_with_opts(account, state.clone(), opts.clone()) + .await + { Ok(bot) => return bot, Err(e) => { disconnects += 1; |