AES serde

Cargo.lock

@@ -17,6 +17,17 @@ version = "2.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "512761e0bb2578dd7380c6baaa0f4ce03e84f95e960231d1dec8bf4d7d6e2627"
 
+[[package]]
+name = "aes"
+version = "0.8.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0"
+dependencies = [
+ "cfg-if",
+ "cipher",
+ "cpufeatures",
+]
+
 [[package]]
 name = "anstream"
 version = "0.6.18"
@@ -129,12 +140,31 @@ version = "1.10.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a"
 
+[[package]]
+name = "cfb8"
+version = "0.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "014c0a0e1ad0dae6a86c082db2f9bd7fe8c2c734227047d0d8b4d4a3a094a1e1"
+dependencies = [
+ "cipher",
+]
+
 [[package]]
 name = "cfg-if"
 version = "1.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
+[[package]]
+name = "cipher"
+version = "0.4.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad"
+dependencies = [
+ "crypto-common",
+ "inout",
+]
+
 [[package]]
 name = "clap"
 version = "4.5.37"
@@ -185,12 +215,15 @@ checksum = "5b63caa9aa9397e2d9480a9b13673856c78d8ac123288526c37d7839f2a86990"
 name = "composition"
 version = "0.1.0"
 dependencies = [
+ "aes",
  "async-trait",
  "base64",
+ "cfb8",
  "clap",
  "const_format",
  "der",
  "futures",
+ "generic-array",
  "nom",
  "once_cell",
  "rand",
@@ -233,6 +266,15 @@ dependencies = [
  "unicode-xid",
 ]
 
+[[package]]
+name = "cpufeatures"
+version = "0.2.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "crossbeam-channel"
 version = "0.5.15"
@@ -456,6 +498,15 @@ dependencies = [
  "hashbrown",
 ]
 
+[[package]]
+name = "inout"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01"
+dependencies = [
+ "generic-array",
+]
+
 [[package]]
 name = "is_terminal_polyfill"
 version = "1.70.1"

Cargo.toml

@@ -41,3 +41,6 @@ uuid = { version = "1.13.1", features = ["v4"] }
 rsa = "0.9.8"
 rand = { version = "0.8.5", features = ["std"] }
 der = { version = "0.7.10", features = ["alloc", "derive"] }
+aes = "0.8.4"
+cfb8 = { version = "0.8.1", features = ["alloc"] }
+generic-array = "0.14.7"

src/config.rs

@@ -1,5 +1,6 @@
 use clap::Arg;
 use once_cell::sync::OnceCell;
+use rsa::{RsaPrivateKey, RsaPublicKey};
 use serde::{Deserialize, Serialize};
 use std::io::{Read, Write};
 use std::{fs::File, path::Path, path::PathBuf};
@@ -32,7 +33,7 @@ pub fn read_file(path: &Path) -> std::io::Result<Vec<u8>> {
 }
 
 /// The global configuration.
-#[derive(Debug, Default, Deserialize, Serialize)]
+#[derive(Debug, Deserialize, Serialize)]
 #[serde(rename_all = "kebab-case")]
 #[serde(default)]
 pub struct Config {
@@ -42,6 +43,10 @@ pub struct Config {
     pub server: ServerConfig,
     #[cfg(feature = "proxy")]
     pub proxy: ProxyConfig,
+    /// RSA key pair used for encryption and decryption.
+    /// Generated on each startup and not saved to disk.
+    #[serde(skip)]
+    pub rsa_key_pair: (RsaPublicKey, RsaPrivateKey),
 }
 impl Config {
     pub fn get_formatted_version(subcommand: Subcommand) -> String {
@@ -126,6 +131,21 @@ impl Config {
         std::process::exit(1);
     }
 }
+impl Default for Config {
+    fn default() -> Self {
+        let rsa_key_pair = RsaPrivateKey::new(&mut rand::thread_rng(), 1024)
+            .map(|key| (key.to_public_key(), key))
+            .expect("Failed to generate RSA key pair");
+        Config {
+            global: GlobalConfig::default(),
+            #[cfg(feature = "server")]
+            server: ServerConfig::default(),
+            #[cfg(feature = "proxy")]
+            proxy: ProxyConfig::default(),
+            rsa_key_pair,
+        }
+    }
+}
 
 /// The global configuration.
 #[derive(Debug, Deserialize, Serialize)]

src/net/codec.rs

@@ -1,5 +1,6 @@
 use super::error::Error;
 use crate::protocol::{
+    encryption::*,
     packets::{Packet, PacketDirection},
     parsing::Parsable,
     types::VarInt,
@@ -11,16 +12,18 @@ use tokio_util::{
 };
 use tracing::trace;
 
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Debug)]
 pub struct PacketCodec {
     pub client_state: ClientState,
     pub packet_direction: PacketDirection,
+    pub aes_cipher: Option<(Aes128Cfb8Encryptor, Aes128Cfb8Decryptor, usize)>,
 }
 impl PacketCodec {
     pub fn new(client_state: ClientState, packet_direction: PacketDirection) -> PacketCodec {
         PacketCodec {
             client_state,
             packet_direction,
+            aes_cipher: None,
         }
     }
 }
@@ -29,6 +32,7 @@ impl Default for PacketCodec {
         PacketCodec {
             client_state: ClientState::Handshake,
             packet_direction: PacketDirection::Serverbound,
+            aes_cipher: None,
         }
     }
 }
@@ -37,10 +41,38 @@ impl Decoder for PacketCodec {
     type Error = Error;
 
     fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
+        // Bytes from [0..encryption_start] are decrypted.
+        // Bytes from [encryption_start..] are decrypted.
+        if let Some((_, ref mut aes_decryptor, encryption_start)) = self.aes_cipher.as_mut() {
+            // We have to do a bunch of stupid type fuckery to get CFB8 working
+            // because for some reason decrypt() consumes self?!
+            let encrypted_src = src.split_off(*encryption_start);
+            // Convert BytesMut to Vec<GenericArray<u8, UInt<UTerm, B1>>>
+            let mut encrypted_src = encrypted_src
+                .into_iter()
+                .map(|b| *GenericCFB8BlockArray::from_slice(&[b]))
+                .collect::<Vec<_>>();
+            // Decrypt the bytes in place.
+            aes_decryptor.decrypt_blocks_mut(encrypted_src.as_mut_slice());
+            // Convert Vec<GenericArray<u8, UInt<UTerm, B1>>> to Vec<u8>
+            let encrypted_src = encrypted_src
+                .into_iter()
+                .flat_map(|b| b.to_vec())
+                .collect::<Vec<u8>>();
+
+            // Append the decrypted bytes back to the source and move the encryption start position.
+            *encryption_start += encrypted_src.len();
+            src.extend_from_slice(&encrypted_src);
+        }
+
         match Packet::parse(self.client_state, self.packet_direction, src) {
             Ok((rest, packet)) => {
                 let bytes_consumed = src.len() - rest.len();
                 src.advance(bytes_consumed);
+                if let Some((_, _, encryption_start)) = &mut self.aes_cipher {
+                    // Adjust the encryption start position if we are using AES encryption.
+                    *encryption_start -= bytes_consumed;
+                }
 
                 if let Some(next_state) = packet.state_change() {
                     self.client_state = next_state;
@@ -74,12 +106,35 @@ impl Encoder<Packet> for PacketCodec {
     type Error = Error;
 
     fn encode(&mut self, item: Packet, dst: &mut BytesMut) -> Result<(), Self::Error> {
-        let mut out = vec![];
+        let mut body = vec![];
         let (packet_id, packet_body) = item.serialize();
-        out.extend(packet_id.serialize().to_vec());
+        body.extend(packet_id.serialize().to_vec());
-        out.extend(packet_body);
+        body.extend(packet_body);
-        let packet_len = VarInt::from(out.len());
+        // TODO: Packet compression on `body`.
-        dst.extend(packet_len.serialize());
+        let packet_len = VarInt::from(body.len()).serialize();
+        let mut out = Vec::with_capacity(packet_len.len() + body.len());
+        out.extend(packet_len);
+        out.extend(body);
+
+        if let Some((ref mut aes_encryptor, _, _)) = &mut self.aes_cipher {
+            // We have to do a bunch of stupid type fuckery to get CFB8 working
+            // because for some reason decrypt() consumes self?!
+            // Convert Vec<u8> to Vec<GenericArray<u8, UInt<UTerm, B1>>>
+            let mut encrypted_out = out
+                .into_iter()
+                .map(|b| *GenericCFB8BlockArray::from_slice(&[b]))
+                .collect::<Vec<_>>();
+            // Decrypt the bytes in place.
+            aes_encryptor.encrypt_blocks_mut(encrypted_out.as_mut_slice());
+            // Convert Vec<GenericArray<u8, UInt<UTerm, B1>>> to Vec<u8>
+            let encrypted_out = encrypted_out
+                .into_iter()
+                .flat_map(|b| b.to_vec())
+                .collect::<Vec<u8>>();
+
+            out = encrypted_out;
+        }
+
         dst.extend(out);
         Ok(())
     }

src/net/connection.rs

@@ -1,10 +1,14 @@
 use super::{codec::PacketCodec, error::Error};
 use crate::protocol::{
+    encryption::*,
     packets::{self, Packet, PacketDirection},
     types::Chat,
     ClientState,
 };
 use futures::{stream::StreamExt, SinkExt};
+use rand::rngs::StdRng;
+use rand::Rng;
+use rand::SeedableRng;
 use std::{
     collections::HashMap,
     time::{Duration, Instant},
@@ -196,9 +200,11 @@ impl Connection {
             last_sent_data_time: Instant::now(),
         }
     }
+    /// Make a Connection from a `TcpStream`, acting as a client talking to a server.
     pub fn new_client(id: u128, stream: TcpStream) -> Self {
         Self::new(id, PacketDirection::Serverbound, stream)
     }
+    /// Make a Connection from a `TcpStream`, acting as a server talking to a client.
     pub fn new_server(id: u128, stream: TcpStream) -> Self {
         Self::new(id, PacketDirection::Clientbound, stream)
     }
@@ -218,8 +224,7 @@ impl Connection {
         self.last_sent_data_time.elapsed()
     }
     pub async fn read_packet(&mut self) -> Option<Result<Packet, Error>> {
-        self.last_received_data_time = Instant::now();
+        let packet = self.stream.next().await.map(|packet| {
-        self.stream.next().await.map(|packet| {
             packet.map_err(|mut e| {
                 // Set the codec error to something more descriptive.
                 if e.to_string() == "bytes remaining on stream" {
@@ -228,10 +233,50 @@ impl Connection {
                 trace!("Error reading packet from connection {}: {:?}", self.id, e);
                 e
             })
-        })
+        });
+
+        if let Some(Ok(ref packet)) = packet {
+            trace!("Received packet from connection {}: {:?}", self.id, packet);
+            self.last_received_data_time = Instant::now();
+
+            if let Packet::EncryptionRequest(packet) = packet {
+                // Extract the public key from the packet.
+                let public_key = rsa::RsaPublicKey::parse(&packet.public_key)
+                    .expect("Failed to parse RSA public key from packet")
+                    .1;
+
+                // Generate a shared secret.
+                let mut rng = StdRng::from_entropy();
+                let shared_secret: [u8; 16] = rng.gen();
+
+                // Create the AES stream cipher and initialize it with the shared secret.
+                let encryptor =
+                    Aes128Cfb8Encryptor::new((&shared_secret).into(), (&shared_secret).into());
+                let decryptor =
+                    Aes128Cfb8Decryptor::new((&shared_secret).into(), (&shared_secret).into());
+
+                // Send the encryption response packet.
+                self.send_packet(packets::login::serverbound::EncryptionResponse {
+                    shared_secret: public_key
+                        .encrypt(&mut rng, rsa::Pkcs1v15Encrypt, &shared_secret[..])
+                        .expect("Failed to encrypt shared secret"),
+                    verify_token: public_key
+                        .encrypt(&mut rng, rsa::Pkcs1v15Encrypt, &packet.verify_token[..])
+                        .expect("Failed to encrypt shared secret"),
+                })
+                .await
+                .expect("Failed to send encryption response");
+
+                // Enable encryption on the connection.
+                self.stream.codec_mut().aes_cipher = Some((encryptor, decryptor, 0));
+            }
+        }
+
+        packet
     }
     pub async fn send_packet<P: Into<Packet>>(&mut self, packet: P) -> Result<(), Error> {
         let packet: Packet = packet.into();
+        trace!("Sending packet to connection {}: {:?}", self.id, packet);
         self.stream.send(packet).await.inspect_err(|e| {
             trace!("Error sending packet to connection {}: {:?}", self.id, e);
         })

src/protocol/encryption.rs

@@ -4,7 +4,15 @@ use der::{
 };
 
 pub use crate::protocol::parsing::Parsable;
-pub use rsa::RsaPublicKey;
+pub use aes::cipher::{BlockDecryptMut, BlockEncryptMut, KeyIvInit};
+pub use generic_array::{
+    typenum::{UInt, UTerm, B1},
+    GenericArray,
+};
+pub use rsa::{RsaPrivateKey, RsaPublicKey};
+pub type Aes128Cfb8Encryptor = cfb8::Encryptor<aes::Aes128>;
+pub type Aes128Cfb8Decryptor = cfb8::Decryptor<aes::Aes128>;
+pub type GenericCFB8BlockArray = GenericArray<u8, UInt<UTerm, B1>>;
 
 impl Parsable for RsaPublicKey {
     fn parse(data: &[u8]) -> nom::IResult<&[u8], Self> {

src/proxy/mod.rs

@@ -28,7 +28,7 @@ impl Proxy {
             .map_err(Error::Io)?;
         Ok(Connection::new_server(0, upstream))
     }
-    pub fn rewrite_packet(packet: Packet) -> Packet {
+    pub fn rewrite_packet(packet: Packet) -> Option<Packet> {
         match packet {
             Packet::StatusResponse(mut status) => {
                 let new_description = ProxyConfig::default().version.clone();
@@ -38,29 +38,10 @@ impl Proxy {
                     .unwrap()
                     .get_mut("description")
                     .unwrap() = serde_json::Value::String(new_description);
-                Packet::StatusResponse(status)
+                Some(Packet::StatusResponse(status))
             }
-            Packet::EncryptionRequest(mut p) => {
+            Packet::EncryptionRequest(_) => None,
-                trace!("Rewriting encryption request packet: {:?}", p);
+            p => Some(p),
-                use crate::protocol::parsing::Parsable;
-
-                // Decode the upstream public key from the packet.
-                let upstream_public_key = rsa::RsaPublicKey::parse(&p.public_key)
-                    .expect("Failed to parse RSA public key from packet");
-                trace!("server public key: {:?}", upstream_public_key);
-
-                // Make our own private and public rsa keys and send those instead.
-                use rand::SeedableRng;
-                let mut rng = rand::rngs::StdRng::from_entropy();
-                let private_key = rsa::RsaPrivateKey::new(&mut rng, 1024)
-                    .expect("Failed to generate RSA private key");
-                let public_key = private_key.to_public_key();
-                trace!("local public key: {:?}", public_key);
-                p.public_key = public_key.serialize();
-
-                Packet::EncryptionRequest(p)
-            }
-            p => p,
         }
     }
 }
@@ -119,9 +100,10 @@ impl App for Proxy {
                 if let Some(packet) = packet {
                     match packet {
                         Ok(packet) => {
-                            trace!("Got packet from client: {:?}", packet);
                             let next_state = packet.state_change();
-                            self.upstream.send_packet(Proxy::rewrite_packet(packet)).await.map_err(Error::Network)?;
+                            if let Some(packet) = Proxy::rewrite_packet(packet) {
+                                self.upstream.send_packet(packet).await.map_err(Error::Network)?;
+                            }
                             if let Some(next_state) = next_state {
                                 *self.upstream.client_state_mut() = next_state;
                             }
@@ -151,9 +133,10 @@ impl App for Proxy {
                 if let Some(packet) = packet {
                     match packet {
                         Ok(packet) => {
-                            trace!("Got packet from upstream: {:?}", packet);
                             let next_state = packet.state_change();
-                            client.send_packet(Proxy::rewrite_packet(packet)).await.map_err(Error::Network)?;
+                            if let Some(packet) = Proxy::rewrite_packet(packet) {
+                                client.send_packet(packet).await.map_err(Error::Network)?;
+                            }
                             if let Some(next_state) = next_state {
                                 *client.client_state_mut() = next_state;
                             }