azalea_client/plugins/
login.rs

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