azalea_client/plugins/
chat_signing.rs

1use std::time::{Duration, Instant};
2
3use azalea_auth::certs::{Certificates, FetchCertificatesError};
4use azalea_protocol::packets::game::{
5    ServerboundChatSessionUpdate,
6    s_chat_session_update::{ProfilePublicKeyData, RemoteChatSessionData},
7};
8use bevy_app::prelude::*;
9use bevy_ecs::prelude::*;
10use bevy_tasks::{IoTaskPool, Task, futures_lite::future};
11use chrono::Utc;
12use tracing::{debug, error};
13use uuid::Uuid;
14
15use super::{chat, login::IsAuthenticated, packet::game::SendGamePacketEvent};
16use crate::{InGameState, account::Account};
17
18pub struct ChatSigningPlugin;
19impl Plugin for ChatSigningPlugin {
20    fn build(&self, app: &mut App) {
21        app.add_systems(
22            Update,
23            (
24                request_certs_if_needed,
25                poll_request_certs_task,
26                handle_queued_certs_to_send,
27            )
28                .chain()
29                .before(chat::handler::handle_send_chat_kind_event),
30        );
31    }
32}
33
34#[derive(Component)]
35pub struct RequestCertsTask(pub Task<Result<Certificates, FetchCertificatesError>>);
36
37/// A component that makes us have to wait until the given time to refresh the
38/// certs.
39///
40/// This is used to avoid spamming requests if requesting certs fails. Usually,
41/// we just check [`Certificates::expires_at`].
42#[derive(Component, Debug)]
43pub struct OnlyRefreshCertsAfter {
44    pub refresh_at: Instant,
45}
46/// A component that's present when this client has sent its certificates to the
47/// server.
48///
49/// This should be removed if you want to re-send the certs.
50///
51/// If you want to get the client's actual certificates, you can get that from
52/// the `certs` in the [`Account`] component.
53#[derive(Component)]
54pub struct ChatSigningSession {
55    pub session_id: Uuid,
56    pub messages_sent: u32,
57}
58
59pub fn poll_request_certs_task(
60    mut commands: Commands,
61    mut query: Query<(Entity, &mut RequestCertsTask, &Account)>,
62) {
63    for (entity, mut auth_task, account) in query.iter_mut() {
64        if let Some(poll_res) = future::block_on(future::poll_once(&mut auth_task.0)) {
65            debug!("Finished requesting certs");
66            commands.entity(entity).remove::<RequestCertsTask>();
67
68            match poll_res {
69                Ok(certs) => {
70                    commands.entity(entity).insert(QueuedCertsToSend {
71                        certs: certs.clone(),
72                    });
73                    account.set_certs(certs);
74                }
75                Err(err) => {
76                    error!("Error requesting certs: {err:?}. Retrying in an hour.");
77
78                    commands.entity(entity).insert(OnlyRefreshCertsAfter {
79                        refresh_at: Instant::now() + Duration::from_secs(60 * 60),
80                    });
81                }
82            }
83        }
84    }
85}
86
87#[allow(clippy::type_complexity)]
88pub fn request_certs_if_needed(
89    mut commands: Commands,
90    mut query: Query<
91        (
92            Entity,
93            &Account,
94            Option<&OnlyRefreshCertsAfter>,
95            Option<&ChatSigningSession>,
96        ),
97        (
98            Without<RequestCertsTask>,
99            With<InGameState>,
100            With<IsAuthenticated>,
101        ),
102    >,
103) {
104    for (entity, account, only_refresh_certs_after, chat_signing_session) in query.iter_mut() {
105        if let Some(only_refresh_certs_after) = only_refresh_certs_after
106            && only_refresh_certs_after.refresh_at > Instant::now()
107        {
108            continue;
109        }
110
111        let certs = account.certs();
112        let should_refresh = if let Some(certs) = certs {
113            // certs were already requested and we're waiting for them to refresh
114
115            // but maybe they weren't sent yet, in which case we still want to send the
116            // certs
117            if chat_signing_session.is_none() {
118                true
119            } else {
120                Utc::now() > certs.expires_at
121            }
122        } else {
123            true
124        };
125
126        if should_refresh && let Some(access_token) = account.access_token() {
127            let task_pool = IoTaskPool::get();
128
129            debug!("Started task to fetch certs");
130            let task = task_pool.spawn(async_compat::Compat::new(async move {
131                azalea_auth::certs::fetch_certificates(&access_token).await
132            }));
133            commands
134                .entity(entity)
135                .insert(RequestCertsTask(task))
136                .remove::<OnlyRefreshCertsAfter>();
137        }
138    }
139}
140
141/// A component that's present on players that should send their chat signing
142/// certificates as soon as possible.
143///
144/// This is removed when the certificates get sent.
145#[derive(Component)]
146pub struct QueuedCertsToSend {
147    pub certs: Certificates,
148}
149
150pub fn handle_queued_certs_to_send(
151    mut commands: Commands,
152    query: Query<(Entity, &QueuedCertsToSend), With<IsAuthenticated>>,
153) {
154    for (entity, queued_certs) in &query {
155        let certs = &queued_certs.certs;
156
157        let session_id = Uuid::new_v4();
158
159        let chat_session = RemoteChatSessionData {
160            session_id,
161            profile_public_key: ProfilePublicKeyData {
162                expires_at: certs.expires_at.timestamp_millis() as u64,
163                key: certs.public_key_der.clone(),
164                key_signature: certs.signature_v2.clone(),
165            },
166        };
167
168        debug!("Sending chat signing certs to server");
169
170        commands.trigger(SendGamePacketEvent::new(
171            entity,
172            ServerboundChatSessionUpdate { chat_session },
173        ));
174        commands
175            .entity(entity)
176            .remove::<QueuedCertsToSend>()
177            .insert(ChatSigningSession {
178                session_id,
179                messages_sent: 0,
180            });
181    }
182}