Fixes

7 files changed, 202 insertions, 174 deletions

diff --git a/azalea-protocol/src/lib.rs b/azalea-protocol/src/lib.rs
index 3573894c..d7f75f00 100755
--- a/azalea-protocol/src/lib.rs
+++ b/azalea-protocol/src/lib.rs
@@ -1,6 +1,7 @@
 //! This lib is responsible for parsing Minecraft packets.
 
 #![feature(min_specialization)]
+#![feature(arbitrary_enum_discriminant)]
 
 use std::net::IpAddr;
 use std::str::FromStr;
diff --git a/azalea-protocol/src/mc_buf/read.rs b/azalea-protocol/src/mc_buf/read.rs
index 350c0998..7cb0bb09 100644
--- a/azalea-protocol/src/mc_buf/read.rs
+++ b/azalea-protocol/src/mc_buf/read.rs
@@ -2,7 +2,8 @@ use super::{UnsizedByteArray, MAX_STRING_LENGTH};
 use azalea_chat::component::Component;
 use azalea_core::{
     difficulty::Difficulty, game_type::GameType, resource_location::ResourceLocation,
-    serializable_uuid::SerializableUuid, BlockPos, ChunkSectionPos, Direction, Slot, SlotData,
+    serializable_uuid::SerializableUuid, BlockPos, ChunkSectionPos, Direction, GlobalPos, Slot,
+    SlotData,
 };
 use azalea_crypto::SaltSignaturePair;
 use byteorder::{ReadBytesExt, BE};
@@ -481,6 +482,15 @@ impl McBufReadable for BlockPos {
     }
 }
 
+impl McBufReadable for GlobalPos {
+    fn read_into(buf: &mut impl Read) -> Result<Self, String> {
+        Ok(GlobalPos {
+            pos: BlockPos::read_into(buf)?,
+            dimension: ResourceLocation::read_into(buf)?,
+        })
+    }
+}
+
 impl McBufReadable for Direction {
     fn read_into(buf: &mut impl Read) -> Result<Self, String> {
         match buf.read_varint()? {
diff --git a/azalea-protocol/src/mc_buf/write.rs b/azalea-protocol/src/mc_buf/write.rs
index 95c39bd2..80ffaecd 100644
--- a/azalea-protocol/src/mc_buf/write.rs
+++ b/azalea-protocol/src/mc_buf/write.rs
@@ -2,7 +2,7 @@ use super::{UnsizedByteArray, MAX_STRING_LENGTH};
 use azalea_chat::component::Component;
 use azalea_core::{
     difficulty::Difficulty, game_type::GameType, resource_location::ResourceLocation,
-    serializable_uuid::SerializableUuid, BlockPos, ChunkSectionPos, Direction, Slot,
+    serializable_uuid::SerializableUuid, BlockPos, ChunkSectionPos, Direction, GlobalPos, Slot,
 };
 use azalea_crypto::SaltSignaturePair;
 use byteorder::{BigEndian, WriteBytesExt};
@@ -193,28 +193,24 @@ impl McBufWritable for Vec<u8> {
     }
 }
 
-// string
 impl McBufWritable for String {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         buf.write_utf(self)
     }
 }
 
-// ResourceLocation
 impl McBufWritable for ResourceLocation {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         buf.write_resource_location(self)
     }
 }
 
-// u32
 impl McBufWritable for u32 {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         i16::write_into(&(*self as i16), buf)
     }
 }
 
-// u32 varint
 impl McBufVarWritable for u32 {
     fn var_write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         i32::var_write_into(&(*self as i32), buf)
@@ -244,21 +240,18 @@ impl McBufVarWritable for u64 {
     }
 }
 
-// u16
 impl McBufWritable for u16 {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         i16::write_into(&(*self as i16), buf)
     }
 }
 
-// u16 varint
 impl McBufVarWritable for u16 {
     fn var_write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         i32::var_write_into(&(*self as i32), buf)
     }
 }
 
-// Vec<T> varint
 impl<T: McBufVarWritable> McBufVarWritable for Vec<T> {
     fn var_write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         u32::var_write_into(&(self.len() as u32), buf)?;
@@ -269,77 +262,66 @@ impl<T: McBufVarWritable> McBufVarWritable for Vec<T> {
     }
 }
 
-// u8
 impl McBufWritable for u8 {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         buf.write_byte(*self)
     }
 }
 
-// i16
 impl McBufWritable for i16 {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         Writable::write_short(buf, *self)
     }
 }
 
-// i64
 impl McBufWritable for i64 {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         Writable::write_long(buf, *self)
     }
 }
 
-// u64
 impl McBufWritable for u64 {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         i64::write_into(&(*self as i64), buf)
     }
 }
 
-// bool
 impl McBufWritable for bool {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         buf.write_boolean(*self)
     }
 }
 
-// i8
 impl McBufWritable for i8 {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         buf.write_byte(*self as u8)
     }
 }
 
-// f32
 impl McBufWritable for f32 {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         buf.write_float(*self)
     }
 }
 
-// f64
 impl McBufWritable for f64 {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         buf.write_double(*self)
     }
 }
 
-// GameType
 impl McBufWritable for GameType {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         u8::write_into(&self.to_id(), buf)
     }
 }
 
-// Option<GameType>
 impl McBufWritable for Option<GameType> {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         buf.write_byte(GameType::to_optional_id(self) as u8)
     }
 }
 
-// Option<String>
 impl<T: McBufWritable> McBufWritable for Option<T> {
     default fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         if let Some(s) = self {
@@ -352,21 +334,18 @@ impl<T: McBufWritable> McBufWritable for Option<T> {
     }
 }
 
-// azalea_nbt::Tag
 impl McBufWritable for azalea_nbt::Tag {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         buf.write_nbt(self)
     }
 }
 
-// Difficulty
 impl McBufWritable for Difficulty {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         u8::write_into(&self.id(), buf)
     }
 }
 
-// Component
 impl McBufWritable for Component {
     // async fn read_into(buf: &mut impl Read) -> Result<Self, String>
     // where
@@ -384,7 +363,6 @@ impl McBufWritable for Component {
     }
 }
 
-// Slot
 impl McBufWritable for Slot {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         match self {
@@ -400,7 +378,6 @@ impl McBufWritable for Slot {
     }
 }
 
-// Slot
 impl McBufWritable for Uuid {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         buf.write_uuid(self)?;
@@ -409,7 +386,6 @@ impl McBufWritable for Uuid {
     }
 }
 
-// BlockPos
 impl McBufWritable for BlockPos {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         buf.write_long(
@@ -420,14 +396,21 @@ impl McBufWritable for BlockPos {
     }
 }
 
-// Direction
+impl McBufWritable for GlobalPos {
+    fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
+        BlockPos::write_into(&self.pos, buf)?;
+        ResourceLocation::write_into(&self.dimension, buf)?;
+
+        Ok(())
+    }
+}
+
 impl McBufWritable for Direction {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         buf.write_varint(*self as i32)
     }
 }
 
-// ChunkSectionPos
 impl McBufWritable for ChunkSectionPos {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
         let long = (((self.x & 0x3FFFFF) as i64) << 42)
diff --git a/azalea-protocol/src/packets/game/clientbound_declare_commands_packet.rs b/azalea-protocol/src/packets/game/clientbound_declare_commands_packet.rs
index 6743c3af..01f633f4 100755
--- a/azalea-protocol/src/packets/game/clientbound_declare_commands_packet.rs
+++ b/azalea-protocol/src/packets/game/clientbound_declare_commands_packet.rs
@@ -1,6 +1,7 @@
 use super::GamePacket;
 use crate::mc_buf::{McBufReadable, McBufWritable, Readable, Writable};
 use azalea_core::resource_location::ResourceLocation;
+use packet_macros::McBuf;
 use std::{
     hash::Hash,
     io::{Read, Write},
@@ -110,6 +111,58 @@ impl McBufWritable for BrigadierString {
     }
 }
 
+#[derive(Debug, Clone, McBuf)]
+pub enum BrigadierParserType {
+    Bool = 0,
+    Double,
+    Float,
+    Integer,
+    Long,
+    String,
+    Entity,
+    GameProfile,
+    BlockPos,
+    ColumnPos,
+    Vec3,
+    Vec2,
+    BlockState,
+    BlockPredicate,
+    ItemStack,
+    ItemPredicate,
+    Color,
+    Component,
+    Message,
+    Nbt,
+    NbtPath,
+    Objective,
+    ObjectiveCriteira,
+    Operation,
+    Particle,
+    Rotation,
+    Angle,
+    ScoreboardSlot,
+    ScoreHolder,
+    Swizzle,
+    Team,
+    ItemSlot,
+    ResourceLocation,
+    MobEffect,
+    Function,
+    EntityAnchor,
+    Range,
+    IntRange,
+    FloatRange,
+    ItemEnchantment,
+    EntitySummon,
+    Dimension,
+    Uuid,
+    NbtTag,
+    NbtCompoundTag,
+    Time,
+    ResourceOrTag,
+    Resource,
+}
+
 #[derive(Debug, Clone)]
 pub enum BrigadierParser {
     Bool,
@@ -164,119 +217,84 @@ pub enum BrigadierParser {
 
 impl McBufReadable for BrigadierParser {
     fn read_into(buf: &mut impl Read) -> Result<Self, String> {
-        let parser = buf.read_resource_location()?;
+        let parser_type = BrigadierParserType::read_into(buf)?;
 
-        if parser == ResourceLocation::new("brigadier:bool")? {
-            Ok(BrigadierParser::Bool)
-        } else if parser == ResourceLocation::new("brigadier:double")? {
-            Ok(BrigadierParser::Double(BrigadierNumber::read_into(buf)?))
-        } else if parser == ResourceLocation::new("brigadier:float")? {
-            Ok(BrigadierParser::Float(BrigadierNumber::read_into(buf)?))
-        } else if parser == ResourceLocation::new("brigadier:integer")? {
-            Ok(BrigadierParser::Integer(BrigadierNumber::read_into(buf)?))
-        } else if parser == ResourceLocation::new("brigadier:long")? {
-            Ok(BrigadierParser::Long(BrigadierNumber::read_into(buf)?))
-        } else if parser == ResourceLocation::new("brigadier:string")? {
-            Ok(BrigadierParser::String(BrigadierString::read_into(buf)?))
-        } else if parser == ResourceLocation::new("minecraft:entity")? {
-            let flags = buf.read_byte()?;
-            Ok(BrigadierParser::Entity {
-                single: flags & 0x01 != 0,
-                players_only: flags & 0x02 != 0,
-            })
-        } else if parser == ResourceLocation::new("minecraft:game_profile")? {
-            Ok(BrigadierParser::GameProfile)
-        } else if parser == ResourceLocation::new("minecraft:block_pos")? {
-            Ok(BrigadierParser::BlockPos)
-        } else if parser == ResourceLocation::new("minecraft:column_pos")? {
-            Ok(BrigadierParser::ColumnPos)
-        } else if parser == ResourceLocation::new("minecraft:vec3")? {
-            Ok(BrigadierParser::Vec3)
-        } else if parser == ResourceLocation::new("minecraft:vec2")? {
-            Ok(BrigadierParser::Vec2)
-        } else if parser == ResourceLocation::new("minecraft:block_state")? {
-            Ok(BrigadierParser::BlockState)
-        } else if parser == ResourceLocation::new("minecraft:block_predicate")? {
-            Ok(BrigadierParser::BlockPredicate)
-        } else if parser == ResourceLocation::new("minecraft:item_stack")? {
-            Ok(BrigadierParser::ItemStack)
-        } else if parser == ResourceLocation::new("minecraft:item_predicate")? {
-            Ok(BrigadierParser::ItemPredicate)
-        } else if parser == ResourceLocation::new("minecraft:color")? {
-            Ok(BrigadierParser::Color)
-        } else if parser == ResourceLocation::new("minecraft:component")? {
-            Ok(BrigadierParser::Component)
-        } else if parser == ResourceLocation::new("minecraft:message")? {
-            Ok(BrigadierParser::Message)
-        } else if parser == ResourceLocation::new("minecraft:nbt")? {
-            Ok(BrigadierParser::Nbt)
-        } else if parser == ResourceLocation::new("minecraft:nbt_path")? {
-            Ok(BrigadierParser::NbtPath)
-        } else if parser == ResourceLocation::new("minecraft:objective")? {
-            Ok(BrigadierParser::Objective)
-        } else if parser == ResourceLocation::new("minecraft:objective_criteria")? {
-            Ok(BrigadierParser::ObjectiveCriteira)
-        } else if parser == ResourceLocation::new("minecraft:operation")? {
-            Ok(BrigadierParser::Operation)
-        } else if parser == ResourceLocation::new("minecraft:particle")? {
-            Ok(BrigadierParser::Particle)
-        } else if parser == ResourceLocation::new("minecraft:rotation")? {
-            Ok(BrigadierParser::Rotation)
-        } else if parser == ResourceLocation::new("minecraft:angle")? {
-            Ok(BrigadierParser::Angle)
-        } else if parser == ResourceLocation::new("minecraft:scoreboard_slot")? {
-            Ok(BrigadierParser::ScoreboardSlot)
-        } else if parser == ResourceLocation::new("minecraft:score_holder")? {
-            let flags = buf.read_byte()?;
-            Ok(BrigadierParser::ScoreHolder {
-                allows_multiple: flags & 0x01 != 0,
-            })
-        } else if parser == ResourceLocation::new("minecraft:swizzle")? {
-            Ok(BrigadierParser::Swizzle)
-        } else if parser == ResourceLocation::new("minecraft:team")? {
-            Ok(BrigadierParser::Team)
-        } else if parser == ResourceLocation::new("minecraft:item_slot")? {
-            Ok(BrigadierParser::ItemSlot)
-        } else if parser == ResourceLocation::new("minecraft:resource_location")? {
-            Ok(BrigadierParser::ResourceLocation)
-        } else if parser == ResourceLocation::new("minecraft:mob_effect")? {
-            Ok(BrigadierParser::MobEffect)
-        } else if parser == ResourceLocation::new("minecraft:function")? {
-            Ok(BrigadierParser::Function)
-        } else if parser == ResourceLocation::new("minecraft:entity_anchor")? {
-            Ok(BrigadierParser::EntityAnchor)
-        } else if parser == ResourceLocation::new("minecraft:range")? {
-            Ok(BrigadierParser::Range {
+        match parser_type {
+            BrigadierParserType::Bool => Ok(BrigadierParser::Bool),
+            BrigadierParserType::Double => {
+                Ok(BrigadierParser::Double(BrigadierNumber::read_into(buf)?))
+            }
+            BrigadierParserType::Float => {
+                Ok(BrigadierParser::Float(BrigadierNumber::read_into(buf)?))
+            }
+            BrigadierParserType::Integer => {
+                Ok(BrigadierParser::Integer(BrigadierNumber::read_into(buf)?))
+            }
+            BrigadierParserType::Long => {
+                Ok(BrigadierParser::Long(BrigadierNumber::read_into(buf)?))
+            }
+            BrigadierParserType::String => {
+                Ok(BrigadierParser::String(BrigadierString::read_into(buf)?))
+            }
+            BrigadierParserType::Entity {
+                single,
+                players_only,
+            } => Ok(BrigadierParser::Entity {
+                single,
+                players_only,
+            }),
+            BrigadierParserType::GameProfile => Ok(BrigadierParser::GameProfile),
+            BrigadierParserType::BlockPos => Ok(BrigadierParser::BlockPos),
+            BrigadierParserType::ColumnPos => Ok(BrigadierParser::ColumnPos),
+            BrigadierParserType::Vec3 => Ok(BrigadierParser::Vec3),
+            BrigadierParserType::Vec2 => Ok(BrigadierParser::Vec2),
+            BrigadierParserType::BlockState => Ok(BrigadierParser::BlockState),
+            BrigadierParserType::BlockPredicate => Ok(BrigadierParser::BlockPredicate),
+            BrigadierParserType::ItemStack => Ok(BrigadierParser::ItemStack),
+            BrigadierParserType::ItemPredicate => Ok(BrigadierParser::ItemPredicate),
+            BrigadierParserType::Color => Ok(BrigadierParser::Color),
+            BrigadierParserType::Component => Ok(BrigadierParser::Component),
+            BrigadierParserType::Message => Ok(BrigadierParser::Message),
+            BrigadierParserType::Nbt => Ok(BrigadierParser::Nbt),
+            BrigadierParserType::NbtPath => Ok(BrigadierParser::NbtPath),
+            BrigadierParserType::Objective => Ok(BrigadierParser::Objective),
+            BrigadierParserType::ObjectiveCriteira => Ok(BrigadierParser::ObjectiveCriteira),
+            BrigadierParserType::Operation => Ok(BrigadierParser::Operation),
+            BrigadierParserType::Particle => Ok(BrigadierParser::Particle),
+            BrigadierParserType::Rotation => Ok(BrigadierParser::Rotation),
+            BrigadierParserType::Angle => Ok(BrigadierParser::Angle),
+            BrigadierParserType::ScoreboardSlot => Ok(BrigadierParser::ScoreboardSlot),
+            BrigadierParserType::ScoreHolder => {
+                let flags = buf.read_byte()?;
+                Ok(BrigadierParser::ScoreHolder {
+                    allows_multiple: flags & 0x01 != 0,
+                })
+            }
+            BrigadierParserType::Swizzle => Ok(BrigadierParser::Swizzle),
+            BrigadierParserType::Team => Ok(BrigadierParser::Team),
+            BrigadierParserType::ItemSlot => Ok(BrigadierParser::ItemSlot),
+            BrigadierParserType::ResourceLocation => Ok(BrigadierParser::ResourceLocation),
+            BrigadierParserType::MobEffect => Ok(BrigadierParser::MobEffect),
+            BrigadierParserType::Function => Ok(BrigadierParser::Function),
+            BrigadierParserType::EntityAnchor => Ok(BrigadierParser::EntityAnchor),
+            BrigadierParserType::Range => Ok(BrigadierParser::Range {
                 decimals_allowed: buf.read_boolean()?,
-            })
-        } else if parser == ResourceLocation::new("minecraft:int_range")? {
-            Ok(BrigadierParser::IntRange)
-        } else if parser == ResourceLocation::new("minecraft:float_range")? {
-            Ok(BrigadierParser::FloatRange)
-        } else if parser == ResourceLocation::new("minecraft:item_enchantment")? {
-            Ok(BrigadierParser::ItemEnchantment)
-        } else if parser == ResourceLocation::new("minecraft:entity_summon")? {
-            Ok(BrigadierParser::EntitySummon)
-        } else if parser == ResourceLocation::new("minecraft:dimension")? {
-            Ok(BrigadierParser::Dimension)
-        } else if parser == ResourceLocation::new("minecraft:uuid")? {
-            Ok(BrigadierParser::Uuid)
-        } else if parser == ResourceLocation::new("minecraft:nbt_tag")? {
-            Ok(BrigadierParser::NbtTag)
-        } else if parser == ResourceLocation::new("minecraft:nbt_compound_tag")? {
-            Ok(BrigadierParser::NbtCompoundTag)
-        } else if parser == ResourceLocation::new("minecraft:time")? {
-            Ok(BrigadierParser::Time)
-        } else if parser == ResourceLocation::new("minecraft:resource_or_tag")? {
-            Ok(BrigadierParser::ResourceOrTag {
+            }),
+            BrigadierParserType::IntRange => Ok(BrigadierParser::IntRange),
+            BrigadierParserType::FloatRange => Ok(BrigadierParser::FloatRange),
+            BrigadierParserType::ItemEnchantment => Ok(BrigadierParser::ItemEnchantment),
+            BrigadierParserType::EntitySummon => Ok(BrigadierParser::EntitySummon),
+            BrigadierParserType::Dimension => Ok(BrigadierParser::Dimension),
+            BrigadierParserType::Uuid => Ok(BrigadierParser::Uuid),
+            BrigadierParserType::NbtTag => Ok(BrigadierParser::NbtTag),
+            BrigadierParserType::NbtCompoundTag => Ok(BrigadierParser::NbtCompoundTag),
+            BrigadierParserType::Time => Ok(BrigadierParser::Time),
+            BrigadierParserType::ResourceOrTag => Ok(BrigadierParser::ResourceOrTag {
                 registry_key: buf.read_resource_location()?,
-            })
-        } else if parser == ResourceLocation::new("minecraft:resource")? {
-            Ok(BrigadierParser::Resource {
+            }),
+            BrigadierParserType::Resource => Ok(BrigadierParser::Resource {
                 registry_key: buf.read_resource_location()?,
-            })
-        } else {
-            panic!("Unknown Brigadier parser: {}", parser)
+            }),
         }
     }
 }
diff --git a/azalea-protocol/src/packets/game/clientbound_login_packet.rs b/azalea-protocol/src/packets/game/clientbound_login_packet.rs
index b4a1b8d4..6ddc6b5a 100755
--- a/azalea-protocol/src/packets/game/clientbound_login_packet.rs
+++ b/azalea-protocol/src/packets/game/clientbound_login_packet.rs
@@ -1,4 +1,4 @@
-use azalea_core::{game_type::GameType, resource_location::ResourceLocation};
+use azalea_core::{game_type::GameType, resource_location::ResourceLocation, GlobalPos};
 use packet_macros::{GamePacket, McBuf};
 
 #[derive(Clone, Debug, McBuf, GamePacket)]
@@ -9,7 +9,7 @@ pub struct ClientboundLoginPacket {
     pub previous_game_type: Option<GameType>,
     pub levels: Vec<ResourceLocation>,
     pub registry_holder: azalea_nbt::Tag,
-    pub dimension_type: azalea_nbt::Tag,
+    pub dimension_type: ResourceLocation,
     pub dimension: ResourceLocation,
     pub seed: i64,
     #[var]
@@ -22,4 +22,5 @@ pub struct ClientboundLoginPacket {
     pub show_death_screen: bool,
     pub is_debug: bool,
     pub is_flat: bool,
+    pub last_death_location: Option<GlobalPos>,
 }
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(())
+    }
 }