Make Velocity implementation more robust

Make own Error enum, No more unwraps :D

pumpkin-protocol/src/lib.rs

@@ -138,12 +138,12 @@ pub enum PacketError {
     EncodeData,
     #[error("failed to write encoded packet")]
     EncodeFailedWrite,
-    #[error("failed to write into decoder")]
-    FailedWrite,
+    #[error("failed to write into decoder: {0}")]
+    FailedWrite(String),
     #[error("failed to flush decoder")]
     FailedFinish,
-    #[error("failed to write encoded packet to connection")]
-    ConnectionWrite,
+    #[error("failed to write encoded packet to connection: {0}")]
+    ConnectionWrite(String),
     #[error("packet exceeds maximum length")]
     TooLong,
     #[error("packet length is out of bounds")]

pumpkin-protocol/src/packet_decoder.rs

@@ -63,7 +63,8 @@ impl PacketDecoder {
                 // TODO: use libdeflater or zune-inflate?
                 let mut z = ZlibDecoder::new(&mut self.decompress_buf[..]);
 
-                z.write_all(r).map_err(|_| PacketError::FailedWrite)?;
+                z.write_all(r)
+                    .map_err(|e| PacketError::FailedWrite(e.to_string()))?;
                 z.finish().map_err(|_| PacketError::FailedFinish)?;
 
                 let total_packet_len = VarInt(packet_len).written_size() + packet_len as usize;

pumpkin/src/client/client_packet.rs

@@ -21,8 +21,8 @@ use crate::{
     client::authentication::{self, validate_textures, GameProfile},
     entity::player::{ChatMode, Hand},
     proxy::{
-        bungeecord::bungeecord_login,
-        velocity::{receive_plugin_response, velocity_login},
+        bungeecord,
+        velocity::{self, velocity_login},
     },
     server::{Server, CURRENT_MC_VERSION},
 };
@@ -88,7 +88,7 @@ impl Client {
             if proxy.velocity.enabled {
                 velocity_login(self);
             } else if proxy.bungeecord.enabled {
-                match bungeecord_login(self, login_start.name) {
+                match bungeecord::bungeecord_login(self, login_start.name) {
                     Ok((_ip, profile)) => {
                         // self.address.lock() = ip;
                         self.finish_login(&profile);
@@ -136,20 +136,22 @@ impl Client {
         }
 
         if let Some(profile) = gameprofile.as_ref() {
+            if ADVANCED_CONFIG.packet_compression.enabled {
+                self.enable_compression();
+            }
             self.finish_login(profile);
         } else {
             self.kick("No Game profile");
         }
     }
 
-    fn finish_login(&self, profile: &GameProfile) {
-        // enable compression
-        if ADVANCED_CONFIG.packet_compression.enabled {
-            let compression = ADVANCED_CONFIG.packet_compression.compression_info.clone();
-            self.send_packet(&CSetCompression::new(compression.threshold.into()));
-            self.set_compression(Some(compression));
-        }
+    fn enable_compression(&self) {
+        let compression = ADVANCED_CONFIG.packet_compression.compression_info.clone();
+        self.send_packet(&CSetCompression::new(compression.threshold.into()));
+        self.set_compression(Some(compression));
+    }
 
+    fn finish_login(&self, profile: &GameProfile) {
         let packet = CLoginSuccess::new(&profile.id, &profile.name, &profile.properties, false);
         self.send_packet(&packet);
     }
@@ -173,7 +175,7 @@ impl Client {
                     .allow_banned_players
                 {
                     if !actions.is_empty() {
-                        self.kick("Your account can't join");
+                        return Err(AuthError::Banned);
                     }
                 } else {
                     for allowed in &ADVANCED_CONFIG
@@ -182,7 +184,7 @@ impl Client {
                         .allowed_actions
                     {
                         if !actions.contains(allowed) {
-                            self.kick("Your account can't join");
+                            return Err(AuthError::DisallowedAction);
                         }
                     }
                 }
@@ -198,7 +200,22 @@ impl Client {
     }
 
     pub fn handle_plugin_response(&self, plugin_response: SLoginPluginResponse) {
-        receive_plugin_response(self, &ADVANCED_CONFIG.proxy.velocity, plugin_response);
+        let velocity_config = &ADVANCED_CONFIG.proxy.velocity;
+        if velocity_config.enabled {
+            let mut address = self.address.lock();
+            match velocity::receive_velocity_plugin_response(
+                address.port(),
+                velocity_config,
+                plugin_response,
+            ) {
+                Ok((profile, new_address)) => {
+                    self.finish_login(&profile);
+                    *self.gameprofile.lock() = Some(profile);
+                    *address = new_address
+                }
+                Err(error) => self.kick(&error.to_string()),
+            }
+        }
     }
 
     pub fn handle_login_acknowledged(

pumpkin/src/client/mod.rs

@@ -192,7 +192,7 @@ impl Client {
         self.connection
             .lock()
             .write_all(&enc.take())
-            .map_err(|_| PacketError::ConnectionWrite)
+            .map_err(|e| PacketError::ConnectionWrite(e.to_string()))
             .unwrap_or_else(|e| self.kick(&e.to_string()));
     }
 
@@ -204,7 +204,7 @@ impl Client {
         self.connection
             .lock()
             .write_all(&enc.take())
-            .map_err(|_| PacketError::ConnectionWrite)?;
+            .map_err(|e| PacketError::ConnectionWrite(e.to_string()))?;
         Ok(())
     }
 
@@ -226,7 +226,6 @@ impl Client {
         server: &Arc<Server>,
         packet: &mut RawPacket,
     ) -> Result<(), DeserializerError> {
-        println!("{:?}", self.connection_state.load());
         match self.connection_state.load() {
             pumpkin_protocol::ConnectionState::HandShake => self.handle_handshake_packet(packet),
             pumpkin_protocol::ConnectionState::Status => self.handle_status_packet(server, packet),

pumpkin/src/proxy/velocity.rs

@@ -4,25 +4,52 @@ use bytes::{BufMut, BytesMut};
 use hmac::{Hmac, Mac};
 use pumpkin_config::proxy::VelocityConfig;
 use pumpkin_protocol::{
-    bytebuf::ByteBuffer,
-    client::login::{CLoginPluginRequest, CLoginSuccess},
-    server::login::SLoginPluginResponse,
+    bytebuf::ByteBuffer, client::login::CLoginPluginRequest, server::login::SLoginPluginResponse,
     Property,
 };
+use rand::Rng;
 use sha2::Sha256;
+use thiserror::Error;
 
 use crate::client::{authentication::GameProfile, Client};
 
+/// Proxy implementation for Velocity <https://papermc.io/software/velocity> by PaperMC
+/// Sadly PaperMC does not care about 3th Parties providing support for Velocity, There is no documentation.
+/// I had to understand the Code logic by looking at PaperMC's Velocity implementation: <https://github.com/PaperMC/Paper/blob/master/patches/server/0731-Add-Velocity-IP-Forwarding-Support.patch>
+
 type HmacSha256 = Hmac<Sha256>;
 
-const MAX_SUPPORTED_FORWARDING_VERSION: i32 = 4;
+const MAX_SUPPORTED_FORWARDING_VERSION: u8 = 4;
 const PLAYER_INFO_CHANNEL: &str = "velocity:player_info";
 
+#[derive(Error, Debug)]
+pub enum VelocityError {
+    #[error("No response data received")]
+    NoData,
+    #[error("Unable to verify player details")]
+    FailedVerifyIntegrity,
+    #[error("Failed to read forward version")]
+    FailedReadForwardVersion,
+    #[error("Unsupported forwarding version {0}. Maximum supported version is {1}")]
+    UnsupportedForwardVersion(u8, u8),
+    #[error("Failed to read address")]
+    FailedReadAddress,
+    #[error("Failed to parse address")]
+    FailedParseAddres,
+    #[error("Failed to read game profile name")]
+    FailedReadProfileName,
+    #[error("Failed to read game profile UUID")]
+    FailedReadProfileUUID,
+    #[error("Failed to read game profile properties")]
+    FailedReadProfileProperties,
+}
+
 pub fn velocity_login(client: &Client) {
-    let velocity_message_id: i32 = 0;
+    // TODO: validate packet transaction id from plugin response with this
+    let velocity_message_id: i32 = rand::thread_rng().gen();
 
     let mut buf = BytesMut::new();
-    buf.put_u8(MAX_SUPPORTED_FORWARDING_VERSION as u8);
+    buf.put_u8(MAX_SUPPORTED_FORWARDING_VERSION);
     client.send_packet(&CLoginPluginRequest::new(
         velocity_message_id.into(),
         PLAYER_INFO_CHANNEL,
@@ -32,76 +59,79 @@ pub fn velocity_login(client: &Client) {
 
 pub fn check_integrity(data: (&[u8], &[u8]), secret: &str) -> bool {
     let (signature, data_without_signature) = data;
+    // Our fault, We can panic/expect ?
     let mut mac =
         HmacSha256::new_from_slice(secret.as_bytes()).expect("HMAC can take key of any size");
     mac.update(data_without_signature);
     mac.verify_slice(signature).is_ok()
 }
 
-pub fn receive_plugin_response(
-    client: &Client,
+fn read_game_profile(buf: &mut ByteBuffer) -> Result<GameProfile, VelocityError> {
+    let id = buf
+        .get_uuid()
+        .map_err(|_| VelocityError::FailedReadProfileUUID)?;
+
+    let name = buf
+        .get_string()
+        .map_err(|_| VelocityError::FailedReadProfileName)?;
+    let properties = buf
+        .get_list(|data| {
+            let name = data.get_string()?;
+            let value = data.get_string()?;
+            let signature = data.get_option(|data| data.get_string())?;
+
+            Ok(Property {
+                name,
+                value,
+                signature,
+            })
+        })
+        .map_err(|_| VelocityError::FailedReadProfileProperties)?;
+    Ok(GameProfile {
+        id,
+        name,
+        properties,
+        profile_actions: None,
+    })
+}
+
+pub fn receive_velocity_plugin_response(
+    port: u16,
     config: &VelocityConfig,
     response: SLoginPluginResponse,
-) {
+) -> Result<(GameProfile, SocketAddr), VelocityError> {
     dbg!("velocity response");
     if let Some(data) = response.data {
         let (signature, data_without_signature) = data.split_at(32);
 
         if !check_integrity((signature, data_without_signature), &config.secret) {
-            client.kick("Unable to verify player details");
-            return;
+            return Err(VelocityError::FailedVerifyIntegrity);
         }
         let mut buf = ByteBuffer::new(BytesMut::new());
         buf.put_slice(data_without_signature);
 
         // check velocity version
-        let version = buf.get_var_int().unwrap();
-        let version = version.0;
+        let version = buf
+            .get_var_int()
+            .map_err(|_| VelocityError::FailedReadForwardVersion)?;
+        let version = version.0 as u8;
         if version > MAX_SUPPORTED_FORWARDING_VERSION {
-            client.kick(&format!(
-                "Unsupported forwarding version {version}, Max: {MAX_SUPPORTED_FORWARDING_VERSION}"
+            return Err(VelocityError::UnsupportedForwardVersion(
+                version,
+                MAX_SUPPORTED_FORWARDING_VERSION,
             ));
-            return;
         }
-        // TODO: no unwrap
-        let addr: SocketAddr = SocketAddr::new(
-            buf.get_string().unwrap().parse::<IpAddr>().unwrap(),
-            client.address.lock().port(),
+        let addr = buf
+            .get_string()
+            .map_err(|_| VelocityError::FailedReadAddress)?;
+
+        let socket_addr: SocketAddr = SocketAddr::new(
+            addr.parse::<IpAddr>()
+                .map_err(|_| VelocityError::FailedParseAddres)?,
+            port,
         );
-
-        *client.address.lock() = addr;
-
-        let uuid = buf.get_uuid().unwrap();
-
-        let username = buf.get_string().unwrap();
-
-        // Read game profile properties
-        let properties = buf
-            .get_list(|data| {
-                let name = data.get_string()?;
-                let value = data.get_string()?;
-                let signature = data.get_option(|data| data.get_string())?;
-
-                Ok(Property {
-                    name,
-                    value,
-                    signature,
-                })
-            })
-            .unwrap();
-
-        client.send_packet(&CLoginSuccess {
-            uuid: &uuid,
-            username: &username,
-            properties: &properties,
-            strict_error_handling: false,
-        });
-
-        *client.gameprofile.lock() = Some(GameProfile {
-            id: uuid,
-            name: username,
-            properties,
-            profile_actions: None,
-        });
+        let profile = read_game_profile(&mut buf)?;
+        return Ok((profile, socket_addr));
     }
+    Err(VelocityError::NoData)
 }