steel_protocol/
utils.rs

1//! # Steel Protocol Utils
2//! Utility functions and types for the protocol.
3
4use std::{
5    io,
6    pin::Pin,
7    task::{Context, Poll},
8};
9
10use aes::cipher::{Array, BlockModeDecrypt, BlockModeEncrypt, BlockSizeUser};
11use thiserror::Error;
12use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
13
14/// An AES-128 CFB-8 encryptor.
15pub type Aes128Cfb8Enc = cfb8::Encryptor<aes::Aes128>;
16/// An AES-128 CFB-8 decryptor.
17pub type Aes128Cfb8Dec = cfb8::Decryptor<aes::Aes128>;
18
19/// The maximum size of a packet.
20pub const MAX_PACKET_SIZE: usize = 2_097_152;
21/// The maximum size of a packet's data.
22pub const MAX_PACKET_DATA_SIZE: usize = 8_388_608;
23
24/// Describes the set of packets a connection understands at a given point.
25///
26/// A connection always starts out in state [`ConnectionProtocol::Handshake`]. In this state,
27/// the client sends its desired protocol using [`crate::packets::handshake::SClientIntention`]. The
28/// server then either accepts the connection and switches to the desired
29/// protocol, or it disconnects the client (for example, in case of an
30/// outdated client).
31///
32/// Each protocol has a `PacketListener` implementation tied to it for
33/// server and client respectively.
34///
35/// Every packet must correspond to exactly one protocol.
36#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
37pub enum ConnectionProtocol {
38    /// The handshake protocol. This is the initial protocol, in which the client tells the server its intention (i.e. which protocol it wants to use).
39    Handshake,
40    /// The play protocol. This is the main protocol that is used while "in game" and most normal packets reside in here.
41    Play,
42    /// The status protocol. This protocol is used when a client pings a server while on the multiplayer screen.
43    Status,
44    /// The login protocol. This is the first protocol the client switches to to join a server. It handles authentication with the mojang servers. After it is complete, the connection is switched to the PLAY protocol.
45    Login,
46    /// The configuration protocol. Used for syncing registered registries.
47    Config,
48}
49
50/// A raw packet.
51#[derive(Debug)]
52pub struct RawPacket {
53    /// The ID of the packet.
54    pub id: i32,
55    buffer: Box<[u8]>,
56    payload_start: u32,
57}
58
59impl RawPacket {
60    /// Creates a raw packet from an already-separated payload.
61    #[must_use]
62    pub fn new(id: i32, payload: Vec<u8>) -> Self {
63        Self {
64            id,
65            buffer: payload.into_boxed_slice(),
66            payload_start: 0,
67        }
68    }
69
70    #[expect(
71        clippy::cast_possible_truncation,
72        reason = "packet buffers are limited to MAX_PACKET_DATA_SIZE bytes"
73    )]
74    pub(crate) fn from_buffer(id: i32, buffer: Vec<u8>, payload_start: usize) -> Self {
75        debug_assert!(payload_start <= buffer.len());
76        debug_assert!(payload_start <= MAX_PACKET_DATA_SIZE);
77        Self {
78            id,
79            buffer: buffer.into_boxed_slice(),
80            payload_start: payload_start as u32,
81        }
82    }
83
84    /// Returns the packet payload without its packet ID.
85    #[must_use]
86    pub fn payload(&self) -> &[u8] {
87        &self.buffer[self.payload_start as usize..]
88    }
89}
90
91/// An error that can occur when handling packets.
92#[derive(Error, Debug)]
93pub enum PacketError {
94    #[error("failed to decode packet ID")]
95    /// Failed to decode the packet ID.
96    DecodeID,
97    #[error("packet length {0} exceeds maximum length")]
98    /// The packet length exceeds the maximum length.
99    TooLong(usize),
100    #[error("packet length is out of bounds")]
101    /// The packet length is out of bounds.
102    OutOfBounds,
103    #[error("malformed packet length VarInt: {0}")]
104    /// The packet length `VarInt` is malformed.
105    MalformedLength(String),
106    #[error("malformed packet value: {0}")]
107    /// A value in the packet is malformed.
108    MalformedValue(String),
109    #[error("failed to decompress packet: {0}")]
110    /// Failed to decompress the packet.
111    DecompressionFailed(String),
112    #[error("failed to compress packet: {0}")]
113    /// Failed to compress the packet.
114    CompressionFailed(String),
115    #[error("packet is uncompressed but greater than the threshold")]
116    /// The packet is uncompressed but greater than the threshold.
117    NotCompressed,
118    #[error("failed to decrypt packet: {0}")]
119    /// Failed to decrypt the packet.
120    DecryptionFailed(String),
121    #[error("failed to encrypt packet: {0}")]
122    /// Failed to encrypt the packet.
123    EncryptionFailed(String),
124    #[error("the connection has closed")]
125    /// The connection has closed.
126    ConnectionClosed,
127    #[error("{0}")]
128    /// An error occurred when sending a packet.
129    SendError(String),
130    #[error("Error: {0}")]
131    /// An other error occurred.
132    Other(String),
133    #[error("Invalid protocol: {0}")]
134    /// The protocol is invalid.
135    InvalidProtocol(String),
136}
137
138impl From<io::Error> for PacketError {
139    fn from(value: io::Error) -> Self {
140        //Todo! Define & Handle all cases
141        Self::MalformedValue(value.to_string())
142    }
143}
144
145///NOTE: This makes lots of small writes; make sure there is a buffer somewhere down the line
146pub struct StreamEncryptor<W: AsyncWrite + Unpin> {
147    cipher: Aes128Cfb8Enc,
148    write: W,
149    last_unwritten_encrypted_byte: Option<u8>,
150}
151
152impl<W: AsyncWrite + Unpin> StreamEncryptor<W> {
153    /// Creates a new `StreamEncryptor`.
154    pub fn new(cipher: Aes128Cfb8Enc, stream: W) -> Self {
155        debug_assert_eq!(Aes128Cfb8Enc::block_size(), 1);
156        Self {
157            cipher,
158            write: stream,
159            last_unwritten_encrypted_byte: None,
160        }
161    }
162}
163
164impl<W: AsyncWrite + Unpin> AsyncWrite for StreamEncryptor<W> {
165    #[expect(
166        clippy::unwrap_used,
167        reason = "CFB8 block size is one byte, so each chunk fits the cipher block type"
168    )]
169    fn poll_write(
170        self: Pin<&mut Self>,
171        cx: &mut Context<'_>,
172        buf: &[u8],
173    ) -> Poll<io::Result<usize>> {
174        let ref_self = self.get_mut();
175        let cipher = &mut ref_self.cipher;
176
177        let mut total_written = 0;
178        // Decrypt the raw data, note that our block size is 1 byte, so this is always safe
179        for block in buf.chunks(Aes128Cfb8Enc::block_size()) {
180            let mut out = [0u8];
181
182            if let Some(out_to_use) = ref_self.last_unwritten_encrypted_byte {
183                // This assumes that this `poll_write` is called on the same stream of bytes which I
184                // think is a fair assumption, since thats an invariant for the TCP stream anyway.
185
186                // This should never panic
187                out[0] = out_to_use;
188            } else {
189                // This is a stream cipher, so this value must be used
190                let out_block: &mut Array<u8, _> = (&mut out).into();
191                cipher.encrypt_block_b2b(block.try_into().unwrap(), out_block);
192            }
193
194            let write = Pin::new(&mut ref_self.write);
195            match write.poll_write(cx, &out) {
196                Poll::Pending => {
197                    ref_self.last_unwritten_encrypted_byte = Some(out[0]);
198                    if total_written == 0 {
199                        //If we didn't write anything, return pending
200                        return Poll::Pending;
201                    }
202                    // Otherwise, we actually did write something
203                    return Poll::Ready(Ok(total_written));
204                }
205                Poll::Ready(result) => {
206                    ref_self.last_unwritten_encrypted_byte = None;
207                    match result {
208                        Ok(written) => total_written += written,
209                        Err(err) => return Poll::Ready(Err(err)),
210                    }
211                }
212            }
213        }
214
215        Poll::Ready(Ok(total_written))
216    }
217
218    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
219        let ref_self = self.get_mut();
220        let write = Pin::new(&mut ref_self.write);
221        write.poll_flush(cx)
222    }
223
224    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
225        let ref_self = self.get_mut();
226        let write = Pin::new(&mut ref_self.write);
227        write.poll_shutdown(cx)
228    }
229}
230
231/// A stream that decrypts data.
232pub struct StreamDecryptor<R: AsyncRead + Unpin> {
233    cipher: Aes128Cfb8Dec,
234    read: R,
235}
236
237impl<R: AsyncRead + Unpin> StreamDecryptor<R> {
238    /// Creates a new `StreamDecryptor`.
239    pub const fn new(cipher: Aes128Cfb8Dec, stream: R) -> Self {
240        Self {
241            cipher,
242            read: stream,
243        }
244    }
245}
246
247impl<R: AsyncRead + Unpin> AsyncRead for StreamDecryptor<R> {
248    #[expect(
249        clippy::unwrap_used,
250        reason = "CFB8 block size is one byte, so each chunk fits the cipher block type"
251    )]
252    fn poll_read(
253        self: Pin<&mut Self>,
254        cx: &mut Context<'_>,
255        buf: &mut ReadBuf<'_>,
256    ) -> Poll<io::Result<()>> {
257        let ref_self = self.get_mut();
258        let read = Pin::new(&mut ref_self.read);
259        let cipher = &mut ref_self.cipher;
260
261        // Get the starting position
262        let original_fill = buf.filled().len();
263        // Read the raw data
264        let internal_poll = read.poll_read(cx, buf);
265
266        if matches!(internal_poll, Poll::Ready(Ok(()))) {
267            // Decrypt the raw data in-place, note that our block size is 1 byte, so this is always safe
268            for block in buf.filled_mut()[original_fill..].chunks_mut(Aes128Cfb8Dec::block_size()) {
269                cipher.decrypt_block(block.try_into().unwrap());
270            }
271        }
272
273        internal_poll
274    }
275}
steel_protocol/utils.rs

steel_protocol/
utils.rs
