Fixes

2 files changed, 50 insertions, 35 deletions

diff --git a/azalea-protocol/src/packets/login/clientbound_hello_packet.rs b/azalea-protocol/src/packets/login/clientbound_hello_packet.rs
index f7de4c21..58d48ffe 100755
--- a/azalea-protocol/src/packets/login/clientbound_hello_packet.rs
+++ b/azalea-protocol/src/packets/login/clientbound_hello_packet.rs
@@ -1,37 +1,11 @@
-use std::{
-    hash::Hash,
-    io::{Read, Write},
-};
+use packet_macros::LoginPacket;
+use packet_macros::McBuf;
 
-use super::LoginPacket;
-use crate::mc_buf::Readable;
-
-#[derive(Hash, Clone, Debug)]
+#[derive(Clone, Debug, McBuf, LoginPacket)]
 pub struct ClientboundHelloPacket {
+    // TODO: make this len thing work
+    // #[len(20)]
     pub server_id: String,
     pub public_key: Vec<u8>,
     pub nonce: Vec<u8>,
 }
-
-impl ClientboundHelloPacket {
-    pub fn get(self) -> LoginPacket {
-        LoginPacket::ClientboundHelloPacket(self)
-    }
-
-    pub fn write(&self, _buf: &mut impl Write) -> Result<(), std::io::Error> {
-        panic!("ClientboundHelloPacket::write not implemented")
-    }
-
-    pub fn read(buf: &mut impl Read) -> Result<LoginPacket, String> {
-        let server_id = buf.read_utf_with_len(20)?;
-        let public_key = buf.read_byte_array()?;
-        let nonce = buf.read_byte_array()?;
-
-        Ok(ClientboundHelloPacket {
-            server_id,
-            public_key,
-            nonce,
-        }
-        .get())
-    }
-}
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(())
+    }
 }