azalea_client/plugins/
chat_signing.rs1use 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#[derive(Component)]
43pub struct OnlyRefreshCertsAfter {
44 pub refresh_at: Instant,
45}
46#[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 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#[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}