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