azalea_client/plugins/
login.rs

1#[cfg(feature = "online-mode")]
2use azalea_auth::sessionserver::ClientSessionServerError;
3use azalea_protocol::{
4    connect::Proxy,
5    packets::login::{ClientboundHello, ServerboundCustomQueryAnswer, ServerboundKey},
6};
7use bevy_app::prelude::*;
8use bevy_ecs::prelude::*;
9use bevy_tasks::{IoTaskPool, Task, futures_lite::future};
10use thiserror::Error;
11use tracing::{debug, error, trace, warn};
12
13use super::{
14    connection::RawConnection,
15    packet::login::{ReceiveCustomQueryEvent, ReceiveHelloEvent, SendLoginPacketEvent},
16};
17use crate::{Account, join::ConnectOpts};
18
19/// Some systems that run during the `login` state.
20pub struct LoginPlugin;
21impl Plugin for LoginPlugin {
22    fn build(&self, app: &mut App) {
23        app.add_observer(handle_receive_hello_event)
24            .add_systems(Update, (poll_auth_task, reply_to_custom_queries));
25    }
26}
27
28fn handle_receive_hello_event(
29    receive_hello: On<ReceiveHelloEvent>,
30    mut commands: Commands,
31    query: Query<&ConnectOpts>,
32) {
33    let task_pool = IoTaskPool::get();
34
35    let account = receive_hello.account.clone();
36    let packet = receive_hello.packet.clone();
37    let client = receive_hello.entity;
38
39    // we store the auth proxy in the ConnectOpts component to make it easily
40    // configurable. that component should've definitely been inserted by now, but
41    // if it somehow wasn't then we should let the user know.
42    let connect_opts = if let Ok(opts) = query.get(client) {
43        opts.sessionserver_proxy.clone()
44    } else {
45        warn!("ConnectOpts component missing on a client ({client}) that got ReceiveHelloEvent");
46        None
47    };
48
49    let task = task_pool.spawn(auth_with_account(account, packet, connect_opts));
50    commands.entity(client).insert(AuthTask(task));
51}
52
53/// A marker component on our clients that indicates that the server is
54/// online-mode and the client has authenticated their join with Mojang.
55#[derive(Component)]
56pub struct IsAuthenticated;
57
58pub fn poll_auth_task(
59    mut commands: Commands,
60    mut query: Query<(Entity, &mut AuthTask, &mut RawConnection)>,
61) {
62    for (entity, mut auth_task, mut raw_conn) in query.iter_mut() {
63        if let Some(poll_res) = future::block_on(future::poll_once(&mut auth_task.0)) {
64            debug!("Finished auth");
65            commands
66                .entity(entity)
67                .remove::<AuthTask>()
68                .insert(IsAuthenticated);
69            match poll_res {
70                Ok((packet, private_key)) => {
71                    // we use this instead of SendLoginPacketEvent to ensure that it's sent right
72                    // before encryption is enabled. i guess another option would be to make a
73                    // Trigger+observer for set_encryption_key; the current implementation is
74                    // simpler though.
75                    if let Err(e) = raw_conn.write(packet) {
76                        error!("Error sending key packet: {e:?}");
77                    }
78                    if let Some(net_conn) = raw_conn.net_conn() {
79                        net_conn.set_encryption_key(private_key);
80                    }
81                }
82                Err(err) => {
83                    error!("Error during authentication: {err:?}");
84                }
85            }
86        }
87    }
88}
89
90type PrivateKey = [u8; 16];
91
92#[derive(Component)]
93pub struct AuthTask(Task<Result<(ServerboundKey, PrivateKey), AuthWithAccountError>>);
94
95#[derive(Debug, Error)]
96pub enum AuthWithAccountError {
97    #[error("Failed to encrypt the challenge from the server for {0:?}")]
98    Encryption(ClientboundHello),
99    #[cfg(feature = "online-mode")]
100    #[error("{0}")]
101    SessionServer(#[from] ClientSessionServerError),
102    #[cfg(feature = "online-mode")]
103    #[error("Couldn't refresh access token: {0}")]
104    Auth(#[from] azalea_auth::AuthError),
105}
106
107pub async fn auth_with_account(
108    account: Account,
109    packet: ClientboundHello,
110    proxy: Option<Proxy>,
111) -> Result<(ServerboundKey, PrivateKey), AuthWithAccountError> {
112    let Ok(encrypt_res) = azalea_crypto::encrypt(&packet.public_key, &packet.challenge) else {
113        return Err(AuthWithAccountError::Encryption(packet));
114    };
115    let key_packet = ServerboundKey {
116        key_bytes: encrypt_res.encrypted_public_key,
117        encrypted_challenge: encrypt_res.encrypted_challenge,
118    };
119    let private_key = encrypt_res.secret_key;
120
121    #[cfg(not(feature = "online-mode"))]
122    let _ = (account, proxy);
123
124    #[cfg(feature = "online-mode")]
125    if packet.should_authenticate {
126        let Some(access_token) = &account.access_token else {
127            // offline mode account, no need to do auth
128            return Ok((key_packet, private_key));
129        };
130
131        // keep track of the number of times we tried authenticating so we can give up
132        // after too many
133        let mut attempts: usize = 1;
134
135        let proxy = proxy.map(Proxy::into);
136
137        while let Err(err) = {
138            use azalea_auth::sessionserver::{self, SessionServerJoinOpts};
139
140            let access_token = access_token.lock().clone();
141
142            let uuid = &account
143                .uuid
144                .expect("Uuid must be present if access token is present.");
145
146            let proxy = proxy.clone();
147
148            // this is necessary since reqwest usually depends on tokio and we're using
149            // `futures` here
150            async_compat::Compat::new(sessionserver::join(SessionServerJoinOpts {
151                access_token: &access_token,
152                public_key: &packet.public_key,
153                private_key: &private_key,
154                uuid,
155                server_id: &packet.server_id,
156                proxy,
157            }))
158            .await
159        } {
160            if attempts >= 2 {
161                // if this is the second attempt and we failed
162                // both times, give up
163                return Err(err.into());
164            }
165            if matches!(
166                err,
167                ClientSessionServerError::InvalidSession
168                    | ClientSessionServerError::ForbiddenOperation
169            ) {
170                // uh oh, we got an invalid session and have
171                // to reauthenticate now
172                async_compat::Compat::new(account.refresh()).await?;
173            } else {
174                return Err(err.into());
175            }
176            attempts += 1;
177        }
178    }
179
180    Ok((key_packet, private_key))
181}
182
183pub fn reply_to_custom_queries(
184    mut commands: Commands,
185    mut events: MessageReader<ReceiveCustomQueryEvent>,
186) {
187    for event in events.read() {
188        trace!("Maybe replying to custom query: {event:?}");
189        if event.disabled {
190            continue;
191        }
192
193        commands.trigger(SendLoginPacketEvent::new(
194            event.entity,
195            ServerboundCustomQueryAnswer {
196                transaction_id: event.packet.transaction_id,
197                data: None,
198            },
199        ));
200    }
201}
azalea_client/plugins/login.rs

azalea_client/plugins/
login.rs
