Fixes

1 files changed, 45 insertions, 4 deletions

diff --git a/azalea-protocol/src/packets/login/serverbound_key_packet.rs b/azalea-protocol/src/packets/login/serverbound_key_packet.rs
index 9100823d..2c970036 100644
--- a/azalea-protocol/src/packets/login/serverbound_key_packet.rs
+++ b/azalea-protocol/src/packets/login/serverbound_key_packet.rs
@@ -1,8 +1,49 @@
+use azalea_crypto::SaltSignaturePair;
 use packet_macros::{LoginPacket, McBuf};
-use std::hash::Hash;
+use std::{
+    hash::Hash,
+    io::{Read, Write},
+};
 
-#[derive(Hash, Clone, Debug, McBuf, LoginPacket)]
+use crate::mc_buf::{McBufReadable, McBufWritable};
+
+#[derive(Clone, Debug, McBuf, LoginPacket)]
 pub struct ServerboundKeyPacket {
-    pub shared_secret: Vec<u8>,
-    pub nonce: Vec<u8>,
+    pub key_bytes: Vec<u8>,
+    pub nonce_or_salt_signature: NonceOrSaltSignature,
+}
+
+#[derive(Clone, Debug)]
+pub enum NonceOrSaltSignature {
+    Nonce(Vec<u8>),
+    SaltSignature(SaltSignaturePair),
+}
+
+impl McBufReadable for NonceOrSaltSignature {
+    fn read_into(buf: &mut impl Read) -> Result<Self, String> {
+        let is_nonce = bool::read_into(buf)?;
+        if is_nonce {
+            Ok(NonceOrSaltSignature::Nonce(Vec::<u8>::read_into(buf)?))
+        } else {
+            Ok(NonceOrSaltSignature::SaltSignature(
+                SaltSignaturePair::read_into(buf)?,
+            ))
+        }
+    }
+}
+
+impl McBufWritable for NonceOrSaltSignature {
+    fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
+        match self {
+            NonceOrSaltSignature::Nonce(nonce) => {
+                bool::write_into(&true, buf)?;
+                nonce.write_into(buf)?;
+            }
+            NonceOrSaltSignature::SaltSignature(salt_signature) => {
+                bool::write_into(&false, buf)?;
+                salt_signature.write_into(buf)?;
+            }
+        }
+        Ok(())
+    }
 }