azalea_client/plugins/
login.rs

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

azalea_client/plugins/
login.rs
