refactor: Server code refactor

Author: Andcool-Systems

src/api/mod.rs

@@ -1,11 +1,12 @@
 pub mod endpoint_handler;
 
 use actix_web::{dev::Server, App, HttpServer};
+use anyhow::Result;
 use std::net::SocketAddrV4;
 
 /// Build HTTP server
-pub fn build_http_server(addr: SocketAddrV4) -> Server {
+pub fn build_http_server(addr: SocketAddrV4) -> Result<Server> {
     let app = || App::new().service(endpoint_handler::code);
 
-    HttpServer::new(app).bind(addr).unwrap().run()
+    Ok(HttpServer::new(app).bind(addr)?.run())
 }

src/config/mod.rs

@@ -1,7 +1,7 @@
-pub mod server_icon;
 pub mod types;
 
 use anyhow::Result;
+use base64_light::base64_encode_bytes;
 use std::sync::OnceLock;
 use tokio::fs;
 use tracing::warn;
@@ -22,7 +22,7 @@ pub async fn load(path: &str) -> Result<()> {
     };
 
     //let mut config: types::Config = toml::from_str(&file);
-    match server_icon::load(&config.server.status.icon_path).await {
+    match load_server_icon(&config.server.status.icon_path).await {
         Ok(base64) => config.image = Some(format!("data:image/png;base64,{}", base64)),
         Err(e) => warn!("Error loading server icon: {}", e),
     }
@@ -35,3 +35,11 @@ pub async fn load(path: &str) -> Result<()> {
 pub async fn get_config() -> &'static types::Config {
     CONFIG.get().expect("Config didn't loaded")
 }
+
+/**
+Load server icon from file
+*/
+pub async fn load_server_icon(file_path: &str) -> anyhow::Result<String> {
+    let file_content = fs::read(file_path).await?;
+    Ok(base64_encode_bytes(&file_content))
+}

src/config/server_icon.rs

@@ -1,17 +1,20 @@
-use crate::client::Session;
+use crate::server::MinecraftServer;
+use aes::cipher::AsyncStreamCipher;
 use bytes::BytesMut;
 
-pub fn encrypt_packet(data: &mut BytesMut, session: &mut Session) {
-    use aes::cipher::{generic_array::GenericArray, BlockEncryptMut, BlockSizeUser};
-    use aes::Aes128;
+impl MinecraftServer {
+    /**
+    Encrypt minecraft packet in place
 
-    if session.cipher.is_none() {
-        return;
-    }
+    If session cipher is not set — skip the encryption
+    */
+    pub fn encrypt_packet(&mut self, data: &mut BytesMut) {
+        if self.session.cipher.is_none() {
+            return;
+        }
 
-    let mut cipher = session.cipher.clone().unwrap();
-    for chunk in data.chunks_mut(cfb8::Encryptor::<Aes128>::block_size()) {
-        let gen_arr = GenericArray::from_mut_slice(chunk);
-        cipher.encrypt_block_mut(gen_arr);
+        if let Some(cipher) = &mut self.session.cipher {
+            cipher.clone().encrypt(data);
+        }
     }
 }

src/encryption.rs

@@ -1,46 +1,40 @@
-use std::sync::Arc;
-
-use crate::{client::Session, packets::encryption_response::EncryptionResponsePacket};
+use crate::{packets::encryption_response::EncryptionResponsePacket, server::MinecraftServer};
 use aes::{cipher::KeyIvInit, Aes128};
-use anyhow::Result;
-use bytes::BytesMut;
+use anyhow::{Error, Result};
+use rsa::Pkcs1v15Encrypt;
 
-use rsa::{Pkcs1v15Encrypt, RsaPrivateKey};
+impl MinecraftServer {
+    /**
+    Handle encryption response
 
-pub fn handle_encryption(
-    session: &mut Session,
-    buffer: &mut BytesMut,
-    keys: Arc<rsa::RsaPrivateKey>,
-) -> Result<()> {
-    let response = EncryptionResponsePacket::parse(buffer)?;
+    Check verify tokens and set up the cipher
+    */
+    pub fn handle_encryption(&mut self) -> Result<()> {
+        let response = EncryptionResponsePacket::parse(&mut self.buffer)?;
 
-    // Decrypt client's keys
-    let decrypted_secret = decrypt(&keys, &response.shared_secret)?;
-    let decrypted_verify = decrypt(&keys, &response.verify_token)?;
+        // Decrypt client's keys
+        let decrypted_secret = self
+            .keys
+            .decrypt(Pkcs1v15Encrypt, &response.shared_secret)?;
+        let decrypted_verify = self.keys.decrypt(Pkcs1v15Encrypt, &response.verify_token)?;
 
-    // Check tokens equality
-    if decrypted_verify
-        .iter()
-        .zip(&session.verify_token)
-        .filter(|&(a, b)| a == b)
-        .count()
-        != decrypted_verify.len()
-    {
-        panic!("Verify tokens didn't match!");
-    }
+        // Check tokens equality
+        if decrypted_verify
+            .iter()
+            .zip(&self.session.verify_token)
+            .filter(|&(a, b)| a == b)
+            .count()
+            != decrypted_verify.len()
+        {
+            return Err(Error::msg("Verify tokens didn't match!"));
+        }
 
-    // Set up client cipher
-    session.secret = Some(decrypted_secret.clone());
-    session.cipher = Some(
-        cfb8::Encryptor::<Aes128>::new_from_slices(
+        // Set up client cipher
+        self.session.secret = Some(decrypted_secret.clone());
+        self.session.cipher = Some(cfb8::Encryptor::<Aes128>::new_from_slices(
             &decrypted_secret.clone(),
             &decrypted_secret.clone(),
-        )
-        .unwrap(),
-    );
-    Ok(())
-}
-
-fn decrypt(private_key: &RsaPrivateKey, data: &Vec<u8>) -> Result<Vec<u8>, rsa::Error> {
-    return private_key.decrypt(Pkcs1v15Encrypt, data);
+        )?);
+        Ok(())
+    }
 }

src/handlers/encryption_response.rs

@@ -1,33 +1,35 @@
 use anyhow::{Error, Result};
 
 use crate::{
-    client::{MinecraftClient, NextStateEnum},
     packets::handshake::HandshakePacket,
-    responses::disconnect::send_disconnect,
+    server::{MinecraftServer, NextStateEnum},
 };
 
-pub async fn handle_handshake(client: &mut MinecraftClient) -> Result<()> {
-    let handshake = HandshakePacket::parse(&mut client.buffer)?;
+impl MinecraftServer {
+    /**
+    Handle handshake packet
 
-    client.session.next_state = match handshake.next_state {
-        1 => NextStateEnum::Status,
-        2 => NextStateEnum::Login,
-        _ => NextStateEnum::Unknown,
-    };
+    Also check client's connection host
+    */
+    pub async fn handle_handshake(&mut self) -> Result<()> {
+        let handshake = HandshakePacket::parse(&mut self.buffer)?;
 
-    // Check client's connecting hostname
-    client.session.proto_ver = Some(handshake.proto_ver);
-    if let Some(server_ip) = &client.config.server.server_ip {
-        if server_ip.ne(&handshake.server_addr) {
-            send_disconnect(
-                &mut client.stream,
-                &mut client.session,
-                client.config.messages.using_proxy.clone(),
-            )
-            .await?;
-            return Err(Error::msg("Client using a proxy!"));
+        self.session.next_state = match handshake.next_state {
+            1 => NextStateEnum::Status,
+            2 => NextStateEnum::Login,
+            _ => NextStateEnum::Unknown,
+        };
+
+        // Check client's connecting hostname
+        self.session.proto_ver = Some(handshake.proto_ver);
+        if let Some(server_ip) = &self.config.server.server_ip {
+            if server_ip.ne(&handshake.server_addr) {
+                self.send_disconnect(self.config.messages.using_proxy.clone())
+                    .await?;
+                return Err(Error::msg("Client using a proxy!"));
+            }
         }
-    }
 
-    Ok(())
+        Ok(())
+    }
 }

src/handlers/handshake.rs

@@ -1,9 +1,13 @@
-use crate::{client::Session, packets::login_start::LoginStartPacket};
-use bytes::BytesMut;
+use crate::{packets::login_start::LoginStartPacket, server::MinecraftServer};
 
-pub fn handle_login_start(session: &mut Session, buff: &mut BytesMut) -> anyhow::Result<()> {
-    let packet = LoginStartPacket::parse(buff)?;
-    session.nickname = Some(packet.name);
-    session.uuid = packet.uuid;
-    Ok(())
+impl MinecraftServer {
+    /**
+    Handle login start packet from client
+    */
+    pub fn handle_login_start(&mut self) -> anyhow::Result<()> {
+        let packet = LoginStartPacket::parse(&mut self.buffer)?;
+        self.session.nickname = Some(packet.name);
+        self.session.uuid = packet.uuid;
+        Ok(())
+    }
 }

src/handlers/login_start.rs

@@ -1,6 +1,5 @@
 mod api;
 mod byte_buf_utils;
-mod client;
 mod config;
 mod encryption;
 mod generators;
@@ -9,6 +8,7 @@ mod map;
 mod mojang;
 mod packets;
 mod responses;
+mod server;
 
 use std::{
     net::{Ipv4Addr, SocketAddrV4},
@@ -18,8 +18,8 @@ use std::{
 };
 
 use api::build_http_server;
-use client::MinecraftClient;
 use generators::generate_key_pair;
+use server::MinecraftServer;
 use tokio::{net::TcpListener, sync::Notify};
 use tokio::{signal, time::timeout};
 use tracing::{error, info};
@@ -54,7 +54,8 @@ async fn main() -> anyhow::Result<()> {
     let http = build_http_server(SocketAddrV4::new(
         Ipv4Addr::from_str(&config.api.addr)?,
         config.api.port,
-    ));
+    ))?;
+
     tokio::spawn(async move { http.await });
     info!("API started on port {}", config.api.port);
 
@@ -92,7 +93,7 @@ async fn main() -> anyhow::Result<()> {
             result = listener.accept() => {
                 match result {
                     Ok((stream, addr)) => {
-                        let mut client = MinecraftClient::new(stream, keys.clone()).await;
+                        let mut client = MinecraftServer::new(stream, keys.clone()).await?;
                         info!("New connection from: {}", addr);
 
                         tokio::spawn(async move {

src/main.rs

@@ -1,6 +1,6 @@
 pub mod response;
-use crate::{client::Session, config};
-use anyhow::Result;
+use crate::{config, server::Session};
+use anyhow::{Context, Result};
 use num_bigint::BigInt;
 use rsa::pkcs8::EncodePublicKey;
 use sha1::{Digest, Sha1};
@@ -13,7 +13,13 @@ pub async fn join(
 ) -> Result<Option<response::MojangResponse>> {
     let mut digest = Sha1::new();
     digest.update(session.server_id.as_bytes());
-    digest.update(session.secret.clone().unwrap().as_slice());
+    digest.update(
+        session
+            .secret
+            .clone()
+            .context("Session secret not set")?
+            .as_slice(),
+    );
     digest.update(keys.to_public_key().to_public_key_der()?.as_bytes());
     let hash = BigInt::from_signed_bytes_be(&digest.finalize()).to_str_radix(16);
 
@@ -25,7 +31,13 @@ pub async fn join(
             .server
             .config
             .auth_url
-            .replace("{{NAME}}", &session.nickname.clone().unwrap())
+            .replace(
+                "{{NAME}}",
+                &session
+                    .nickname
+                    .clone()
+                    .context("Client's nickname not set")?,
+            )
             .replace("{{HASH}}", &hash),
     )
     .await?;

src/mojang/mod.rs

@@ -8,11 +8,31 @@ pub struct DisconnectPacket {
 }
 
 impl DisconnectPacket {
+    /**
+    Create disconnect packet with provided reason
+    */
+    pub fn with_reason(reason: String) -> Self {
+        Self { reason }
+    }
+
+    /**
+    Create empty disconnect packet
+    */
+    #[allow(dead_code)]
+    pub fn new() -> Self {
+        Self {
+            reason: String::new(),
+        }
+    }
+
+    /**
+    Build packet for sending over network
+    */
     pub fn build(&self) -> Result<BytesMut> {
         let mut buffer = BytesMut::new();
 
         write_varint(&mut buffer, 0x00); // Packet id
-        write_utf8(&mut buffer, &self.reason)?; // Server id
+        write_utf8(&mut buffer, &format!(r#"{{"text":"{}"}}"#, self.reason))?; // Disconnect reason
 
         Ok(add_size(&buffer))
     }