add chat signing

9 files changed, 259 insertions, 22 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8fc6363e..af76fc77 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,7 @@ is breaking anyways, semantic versioning is not followed.
 - Add auto-reconnecting which is enabled by default.
 - The pathfinder no longer avoids slabs, stairs, and dirt path blocks.
 - Non-standard legacy hex colors like `§#ff0000` are now supported in azalea-chat.
+- Chat signing.
 
 ### Changed
 
diff --git a/Cargo.lock b/Cargo.lock
index 362a99f6..3588f2ad 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -360,6 +360,7 @@ dependencies = [
  "bevy_log",
  "bevy_tasks",
  "bevy_time",
+ "chrono",
  "derive_more 2.0.1",
  "minecraft_folder_path",
  "parking_lot",
diff --git a/azalea-client/Cargo.toml b/azalea-client/Cargo.toml
index 88bf02bc..c9adbe5a 100644
--- a/azalea-client/Cargo.toml
+++ b/azalea-client/Cargo.toml
@@ -25,6 +25,7 @@ bevy_ecs.workspace = true
 bevy_log = { workspace = true, optional = true }
 bevy_tasks.workspace = true
 bevy_time.workspace = true
+chrono = { workspace = true, features = ["now"] }
 derive_more = { workspace = true, features = ["deref", "deref_mut"] }
 minecraft_folder_path.workspace = true
 parking_lot.workspace = true
diff --git a/azalea-client/src/account.rs b/azalea-client/src/account.rs
index a4b35e81..9b2c2350 100644
--- a/azalea-client/src/account.rs
+++ b/azalea-client/src/account.rs
@@ -53,7 +53,7 @@ pub struct Account {
     ///
     /// This is set when you call [`Self::request_certs`], but you only
     /// need to if the servers you're joining require it.
-    pub certs: Option<Certificates>,
+    pub certs: Arc<Mutex<Option<Certificates>>>,
 }
 
 /// The parameters that were passed for creating the associated [`Account`].
@@ -82,7 +82,7 @@ impl Account {
             account_opts: AccountOpts::Offline {
                 username: username.to_string(),
             },
-            certs: None,
+            certs: Arc::new(Mutex::new(None)),
         }
     }
 
@@ -127,7 +127,7 @@ impl Account {
                 email: email.to_string(),
             },
             // we don't do chat signing by default unless the user asks for it
-            certs: None,
+            certs: Arc::new(Mutex::new(None)),
         })
     }
 
@@ -194,7 +194,7 @@ impl Account {
             account_opts: AccountOpts::MicrosoftWithAccessToken {
                 msa: Arc::new(Mutex::new(msa)),
             },
-            certs: None,
+            certs: Arc::new(Mutex::new(None)),
         })
     }
     /// Refresh the access_token for this account to be valid again.
@@ -260,7 +260,7 @@ impl Account {
             .lock()
             .clone();
         let certs = azalea_auth::certs::fetch_certificates(&access_token).await?;
-        self.certs = Some(certs);
+        *self.certs.lock() = Some(certs);
 
         Ok(())
     }
diff --git a/azalea-client/src/plugins/chat/handler.rs b/azalea-client/src/plugins/chat/handler.rs
index 356d3fba..a289eb14 100644
--- a/azalea-client/src/plugins/chat/handler.rs
+++ b/azalea-client/src/plugins/chat/handler.rs
@@ -1,5 +1,6 @@
 use std::time::{SystemTime, UNIX_EPOCH};
 
+use azalea_crypto::SignChatMessageOptions;
 use azalea_protocol::packets::{
     Packet,
     game::{ServerboundChat, ServerboundChatCommand, s_chat::LastSeenMessagesUpdate},
@@ -7,7 +8,7 @@ use azalea_protocol::packets::{
 use bevy_ecs::prelude::*;
 
 use super::ChatKind;
-use crate::packet::game::SendPacketEvent;
+use crate::{Account, chat_signing::ChatSigningSession, packet::game::SendPacketEvent};
 
 /// Send a chat packet to the server of a specific kind (chat message or
 /// command). Usually you just want [`SendChatEvent`] instead.
@@ -30,6 +31,7 @@ pub struct SendChatKindEvent {
 pub fn handle_send_chat_kind_event(
     mut events: EventReader<SendChatKindEvent>,
     mut commands: Commands,
+    mut query: Query<(&Account, &mut ChatSigningSession)>,
 ) {
     for event in events.read() {
         let content = event
@@ -38,22 +40,43 @@ pub fn handle_send_chat_kind_event(
             .filter(|c| !matches!(c, '\x00'..='\x1F' | '\x7F' | '§'))
             .take(256)
             .collect::<String>();
+
+        let timestamp = SystemTime::now();
+
         let packet = match event.kind {
-            ChatKind::Message => ServerboundChat {
-                message: content,
-                timestamp: SystemTime::now()
-                    .duration_since(UNIX_EPOCH)
-                    .expect("Time shouldn't be before epoch")
-                    .as_millis()
-                    .try_into()
-                    .expect("Instant should fit into a u64"),
-                salt: azalea_crypto::make_salt(),
-                signature: None,
-                last_seen_messages: LastSeenMessagesUpdate::default(),
+            ChatKind::Message => {
+                let salt = azalea_crypto::make_salt();
+
+                let signature = if let Ok((account, mut chat_session)) = query.get_mut(event.entity)
+                {
+                    Some(create_signature(
+                        account,
+                        &mut chat_session,
+                        salt,
+                        timestamp,
+                        &content,
+                    ))
+                } else {
+                    None
+                };
+
+                ServerboundChat {
+                    message: content,
+                    timestamp: timestamp
+                        .duration_since(UNIX_EPOCH)
+                        .expect("Time shouldn't be before epoch")
+                        .as_millis()
+                        .try_into()
+                        .expect("Instant should fit into a u64"),
+                    salt,
+                    signature,
+                    // TODO: implement last_seen_messages
+                    last_seen_messages: LastSeenMessagesUpdate::default(),
+                }
             }
             .into_variant(),
             ChatKind::Command => {
-                // TODO: chat signing
+                // TODO: commands that require chat signing
                 ServerboundChatCommand { command: content }.into_variant()
             }
         };
@@ -61,3 +84,28 @@ pub fn handle_send_chat_kind_event(
         commands.trigger(SendPacketEvent::new(event.entity, packet));
     }
 }
+
+pub fn create_signature(
+    account: &Account,
+    chat_session: &mut ChatSigningSession,
+    salt: u64,
+    timestamp: SystemTime,
+    message: &str,
+) -> azalea_crypto::MessageSignature {
+    let certs = account.certs.lock();
+    let certs = certs.as_ref().expect("certs shouldn't be set back to None");
+
+    let signature = azalea_crypto::sign_chat_message(&SignChatMessageOptions {
+        account_uuid: account.uuid.expect("account must have a uuid"),
+        chat_session_uuid: chat_session.session_id,
+        message_index: chat_session.messages_sent,
+        salt,
+        timestamp,
+        message: message.to_owned(),
+        private_key: certs.private_key.clone(),
+    });
+
+    chat_session.messages_sent += 1;
+
+    signature
+}
diff --git a/azalea-client/src/plugins/chat_signing.rs b/azalea-client/src/plugins/chat_signing.rs
new file mode 100644
index 00000000..03dd1a21
--- /dev/null
+++ b/azalea-client/src/plugins/chat_signing.rs
@@ -0,0 +1,182 @@
+use std::time::{Duration, Instant};
+
+use azalea_auth::certs::{Certificates, FetchCertificatesError};
+use azalea_protocol::packets::game::{
+    ServerboundChatSessionUpdate,
+    s_chat_session_update::{ProfilePublicKeyData, RemoteChatSessionData},
+};
+use bevy_app::prelude::*;
+use bevy_ecs::prelude::*;
+use bevy_tasks::{IoTaskPool, Task, futures_lite::future};
+use chrono::Utc;
+use tracing::{debug, error};
+use uuid::Uuid;
+
+use super::{chat, packet::game::SendPacketEvent};
+use crate::{Account, InGameState};
+
+pub struct ChatSigningPlugin;
+impl Plugin for ChatSigningPlugin {
+    fn build(&self, app: &mut App) {
+        app.add_systems(
+            Update,
+            (
+                request_certs_if_needed,
+                poll_request_certs_task,
+                handle_queued_certs_to_send,
+            )
+                .chain()
+                .before(chat::handler::handle_send_chat_kind_event),
+        );
+    }
+}
+
+#[derive(Component)]
+pub struct RequestCertsTask(pub Task<Result<Certificates, FetchCertificatesError>>);
+
+/// A component that makes us have to wait until the given time to refresh the
+/// certs.
+///
+/// This is used to avoid spamming requests if requesting certs fails. Usually,
+/// we just check [`Certificates::expires_at`].
+#[derive(Component)]
+pub struct OnlyRefreshCertsAfter {
+    pub refresh_at: Instant,
+}
+/// A component that's present when that this client has sent its certificates
+/// to the server.
+///
+/// This should be removed if you want to re-send the certs.
+///
+/// If you want to get the client's actual certificates, you can get that from
+/// the `certs` in the [`Account`] component.
+#[derive(Component)]
+pub struct ChatSigningSession {
+    pub session_id: Uuid,
+    pub messages_sent: u32,
+}
+
+pub fn poll_request_certs_task(
+    mut commands: Commands,
+    mut query: Query<(Entity, &mut RequestCertsTask, &Account)>,
+) {
+    for (entity, mut auth_task, account) in query.iter_mut() {
+        if let Some(poll_res) = future::block_on(future::poll_once(&mut auth_task.0)) {
+            debug!("Finished requesting certs");
+            commands.entity(entity).remove::<RequestCertsTask>();
+
+            match poll_res {
+                Ok(certs) => {
+                    commands.entity(entity).insert(QueuedCertsToSend {
+                        certs: certs.clone(),
+                    });
+                    *account.certs.lock() = Some(certs);
+                }
+                Err(err) => {
+                    error!("Error requesting certs: {err:?}. Retrying in an hour.");
+
+                    commands.entity(entity).insert(OnlyRefreshCertsAfter {
+                        refresh_at: Instant::now() + Duration::from_secs(60 * 60),
+                    });
+                }
+            }
+        }
+    }
+}
+
+#[allow(clippy::type_complexity)]
+pub fn request_certs_if_needed(
+    mut commands: Commands,
+    mut query: Query<
+        (
+            Entity,
+            &Account,
+            Option<&OnlyRefreshCertsAfter>,
+            Option<&ChatSigningSession>,
+        ),
+        (Without<RequestCertsTask>, With<InGameState>),
+    >,
+) {
+    for (entity, account, only_refresh_certs_after, chat_signing_session) in query.iter_mut() {
+        if let Some(only_refresh_certs_after) = only_refresh_certs_after {
+            if only_refresh_certs_after.refresh_at > Instant::now() {
+                continue;
+            }
+        }
+
+        let certs = account.certs.lock();
+        let should_refresh = if let Some(certs) = &*certs {
+            // certs were already requested and we're waiting for them to refresh
+
+            // but maybe they weren't sent yet, in which case we still want to send the
+            // certs
+            if chat_signing_session.is_none() {
+                true
+            } else {
+                Utc::now() > certs.expires_at
+            }
+        } else {
+            true
+        };
+        drop(certs);
+
+        if should_refresh {
+            if let Some(access_token) = &account.access_token {
+                let task_pool = IoTaskPool::get();
+
+                let access_token = access_token.lock().clone();
+                debug!("Started task to fetch certs");
+                let task = task_pool.spawn(async_compat::Compat::new(async move {
+                    azalea_auth::certs::fetch_certificates(&access_token).await
+                }));
+                commands
+                    .entity(entity)
+                    .insert(RequestCertsTask(task))
+                    .remove::<OnlyRefreshCertsAfter>();
+            }
+        }
+    }
+}
+
+/// A component that's present on players that should send their chat signing
+/// certificates as soon as possible.
+///
+/// This is removed when the certificates get sent.
+#[derive(Component)]
+pub struct QueuedCertsToSend {
+    pub certs: Certificates,
+}
+
+pub fn handle_queued_certs_to_send(
+    mut commands: Commands,
+    query: Query<(Entity, &QueuedCertsToSend)>,
+) {
+    for (entity, queued_certs) in &query {
+        let certs = &queued_certs.certs;
+
+        let session_id = Uuid::new_v4();
+
+        let chat_session = RemoteChatSessionData {
+            session_id,
+            profile_public_key: ProfilePublicKeyData {
+                expires_at: certs.expires_at.timestamp_millis() as u64,
+                key: certs.public_key_der.clone(),
+                key_signature: certs.signature_v2.clone(),
+            },
+        };
+
+        debug!("Sending chat signing certs to server");
+
+        commands.trigger(SendPacketEvent::new(
+            entity,
+            ServerboundChatSessionUpdate { chat_session },
+        ));
+        commands
+            .entity(entity)
+            .remove::<QueuedCertsToSend>()
+            .insert(ChatSigningSession {
+                session_id,
+                messages_sent: 0,
+            });
+    }
+}
diff --git a/azalea-client/src/plugins/disconnect.rs b/azalea-client/src/plugins/disconnect.rs
index 987007c2..3e27eba3 100644
--- a/azalea-client/src/plugins/disconnect.rs
+++ b/azalea-client/src/plugins/disconnect.rs
@@ -7,7 +7,7 @@ use bevy_ecs::prelude::*;
 use derive_more::Deref;
 use tracing::info;
 
-use crate::{InstanceHolder, client::JoinedClientBundle, connection::RawConnection};
+use crate::{InstanceHolder, chat_signing, client::JoinedClientBundle, connection::RawConnection};
 
 pub struct DisconnectPlugin;
 impl Plugin for DisconnectPlugin {
@@ -70,7 +70,9 @@ pub fn remove_components_from_disconnected_players(
             // this makes it close the tcp connection
             .remove::<RawConnection>()
             // this makes it not send DisconnectEvent again
-            .remove::<IsConnectionAlive>();
+            .remove::<IsConnectionAlive>()
+            // resend our chat signing certs next time
+            .remove::<chat_signing::ChatSigningSession>();
         // note that we don't remove the client from the ECS, so if they decide
         // to reconnect they'll keep their state
 
diff --git a/azalea-client/src/plugins/mod.rs b/azalea-client/src/plugins/mod.rs
index f657b9e9..dea2fdb8 100644
--- a/azalea-client/src/plugins/mod.rs
+++ b/azalea-client/src/plugins/mod.rs
@@ -4,6 +4,7 @@ pub mod attack;
 pub mod auto_reconnect;
 pub mod brand;
 pub mod chat;
+pub mod chat_signing;
 pub mod chunks;
 pub mod connection;
 pub mod disconnect;
@@ -53,7 +54,8 @@ impl PluginGroup for DefaultPlugins {
             .add(connection::ConnectionPlugin)
             .add(login::LoginPlugin)
             .add(join::JoinPlugin)
-            .add(auto_reconnect::AutoReconnectPlugin);
+            .add(auto_reconnect::AutoReconnectPlugin)
+            .add(chat_signing::ChatSigningPlugin);
         #[cfg(feature = "log")]
         {
             group = group.add(bevy_log::LogPlugin::default());
diff --git a/azalea-protocol/src/connect.rs b/azalea-protocol/src/connect.rs
index 8d2711a7..bd8cf115 100644
--- a/azalea-protocol/src/connect.rs
+++ b/azalea-protocol/src/connect.rs
@@ -278,7 +278,7 @@ impl Display for Proxy {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "socks5://")?;
         if let Some(auth) = &self.auth {
-            write!(f, "{}@", auth)?;
+            write!(f, "{auth}@")?;
         }
         write!(f, "{}", self.addr)
     }