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::SendPacketEvent};
16use crate::{Account, InGameState};
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)]
43pub struct OnlyRefreshCertsAfter {
44    pub refresh_at: Instant,
45}
46/// A component that's present when that this client has sent its certificates
47/// to the 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.certs.lock() = Some(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.lock();
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        drop(certs);
126
127        if should_refresh && let Some(access_token) = &account.access_token {
128            let task_pool = IoTaskPool::get();
129
130            let access_token = access_token.lock().clone();
131            debug!("Started task to fetch certs");
132            let task = task_pool.spawn(async_compat::Compat::new(async move {
133                azalea_auth::certs::fetch_certificates(&access_token).await
134            }));
135            commands
136                .entity(entity)
137                .insert(RequestCertsTask(task))
138                .remove::<OnlyRefreshCertsAfter>();
139        }
140    }
141}
142
143/// A component that's present on players that should send their chat signing
144/// certificates as soon as possible.
145///
146/// This is removed when the certificates get sent.
147#[derive(Component)]
148pub struct QueuedCertsToSend {
149    pub certs: Certificates,
150}
151
152pub fn handle_queued_certs_to_send(
153    mut commands: Commands,
154    query: Query<(Entity, &QueuedCertsToSend), With<IsAuthenticated>>,
155) {
156    for (entity, queued_certs) in &query {
157        let certs = &queued_certs.certs;
158
159        let session_id = Uuid::new_v4();
160
161        let chat_session = RemoteChatSessionData {
162            session_id,
163            profile_public_key: ProfilePublicKeyData {
164                expires_at: certs.expires_at.timestamp_millis() as u64,
165                key: certs.public_key_der.clone(),
166                key_signature: certs.signature_v2.clone(),
167            },
168        };
169
170        debug!("Sending chat signing certs to server");
171
172        commands.trigger(SendPacketEvent::new(
173            entity,
174            ServerboundChatSessionUpdate { chat_session },
175        ));
176        commands
177            .entity(entity)
178            .remove::<QueuedCertsToSend>()
179            .insert(ChatSigningSession {
180                session_id,
181                messages_sent: 0,
182            });
183    }
184}