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    /// Could be a `[Box<[u8]>]` but that requires a realloc `if cap != len`
56    pub payload: Vec<u8>,
57}
58
59/// An error that can occur when handling packets.
60#[derive(Error, Debug)]
61pub enum PacketError {
62    #[error("failed to decode packet ID")]
63    /// Failed to decode the packet ID.
64    DecodeID,
65    #[error("packet length {0} exceeds maximum length")]
66    /// The packet length exceeds the maximum length.
67    TooLong(usize),
68    #[error("packet length is out of bounds")]
69    /// The packet length is out of bounds.
70    OutOfBounds,
71    #[error("malformed packet length VarInt: {0}")]
72    /// The packet length `VarInt` is malformed.
73    MalformedLength(String),
74    #[error("malformed packet value: {0}")]
75    /// A value in the packet is malformed.
76    MalformedValue(String),
77    #[error("failed to decompress packet: {0}")]
78    /// Failed to decompress the packet.
79    DecompressionFailed(String),
80    #[error("failed to compress packet: {0}")]
81    /// Failed to compress the packet.
82    CompressionFailed(String),
83    #[error("packet is uncompressed but greater than the threshold")]
84    /// The packet is uncompressed but greater than the threshold.
85    NotCompressed,
86    #[error("failed to decrypt packet: {0}")]
87    /// Failed to decrypt the packet.
88    DecryptionFailed(String),
89    #[error("failed to encrypt packet: {0}")]
90    /// Failed to encrypt the packet.
91    EncryptionFailed(String),
92    #[error("the connection has closed")]
93    /// The connection has closed.
94    ConnectionClosed,
95    #[error("{0}")]
96    /// An error occurred when sending a packet.
97    SendError(String),
98    #[error("Error: {0}")]
99    /// An other error occurred.
100    Other(String),
101    #[error("Invalid protocol: {0}")]
102    /// The protocol is invalid.
103    InvalidProtocol(String),
104}
105
106impl From<io::Error> for PacketError {
107    fn from(value: io::Error) -> Self {
108        //Todo! Define & Handle all cases
109        Self::MalformedValue(value.to_string())
110    }
111}
112
113///NOTE: This makes lots of small writes; make sure there is a buffer somewhere down the line
114pub struct StreamEncryptor<W: AsyncWrite + Unpin> {
115    cipher: Aes128Cfb8Enc,
116    write: W,
117    last_unwritten_encrypted_byte: Option<u8>,
118}
119
120impl<W: AsyncWrite + Unpin> StreamEncryptor<W> {
121    /// Creates a new `StreamEncryptor`.
122    pub fn new(cipher: Aes128Cfb8Enc, stream: W) -> Self {
123        debug_assert_eq!(Aes128Cfb8Enc::block_size(), 1);
124        Self {
125            cipher,
126            write: stream,
127            last_unwritten_encrypted_byte: None,
128        }
129    }
130}
131
132impl<W: AsyncWrite + Unpin> AsyncWrite for StreamEncryptor<W> {
133    #[expect(
134        clippy::unwrap_used,
135        reason = "CFB8 block size is one byte, so each chunk fits the cipher block type"
136    )]
137    fn poll_write(
138        self: Pin<&mut Self>,
139        cx: &mut Context<'_>,
140        buf: &[u8],
141    ) -> Poll<io::Result<usize>> {
142        let ref_self = self.get_mut();
143        let cipher = &mut ref_self.cipher;
144
145        let mut total_written = 0;
146        // Decrypt the raw data, note that our block size is 1 byte, so this is always safe
147        for block in buf.chunks(Aes128Cfb8Enc::block_size()) {
148            let mut out = [0u8];
149
150            if let Some(out_to_use) = ref_self.last_unwritten_encrypted_byte {
151                // This assumes that this `poll_write` is called on the same stream of bytes which I
152                // think is a fair assumption, since thats an invariant for the TCP stream anyway.
153
154                // This should never panic
155                out[0] = out_to_use;
156            } else {
157                // This is a stream cipher, so this value must be used
158                let out_block: &mut Array<u8, _> = (&mut out).into();
159                cipher.encrypt_block_b2b(block.try_into().unwrap(), out_block);
160            }
161
162            let write = Pin::new(&mut ref_self.write);
163            match write.poll_write(cx, &out) {
164                Poll::Pending => {
165                    ref_self.last_unwritten_encrypted_byte = Some(out[0]);
166                    if total_written == 0 {
167                        //If we didn't write anything, return pending
168                        return Poll::Pending;
169                    }
170                    // Otherwise, we actually did write something
171                    return Poll::Ready(Ok(total_written));
172                }
173                Poll::Ready(result) => {
174                    ref_self.last_unwritten_encrypted_byte = None;
175                    match result {
176                        Ok(written) => total_written += written,
177                        Err(err) => return Poll::Ready(Err(err)),
178                    }
179                }
180            }
181        }
182
183        Poll::Ready(Ok(total_written))
184    }
185
186    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
187        let ref_self = self.get_mut();
188        let write = Pin::new(&mut ref_self.write);
189        write.poll_flush(cx)
190    }
191
192    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
193        let ref_self = self.get_mut();
194        let write = Pin::new(&mut ref_self.write);
195        write.poll_shutdown(cx)
196    }
197}
198
199/// A stream that decrypts data.
200pub struct StreamDecryptor<R: AsyncRead + Unpin> {
201    cipher: Aes128Cfb8Dec,
202    read: R,
203}
204
205impl<R: AsyncRead + Unpin> StreamDecryptor<R> {
206    /// Creates a new `StreamDecryptor`.
207    pub const fn new(cipher: Aes128Cfb8Dec, stream: R) -> Self {
208        Self {
209            cipher,
210            read: stream,
211        }
212    }
213}
214
215impl<R: AsyncRead + Unpin> AsyncRead for StreamDecryptor<R> {
216    #[expect(
217        clippy::unwrap_used,
218        reason = "CFB8 block size is one byte, so each chunk fits the cipher block type"
219    )]
220    fn poll_read(
221        self: Pin<&mut Self>,
222        cx: &mut Context<'_>,
223        buf: &mut ReadBuf<'_>,
224    ) -> Poll<io::Result<()>> {
225        let ref_self = self.get_mut();
226        let read = Pin::new(&mut ref_self.read);
227        let cipher = &mut ref_self.cipher;
228
229        // Get the starting position
230        let original_fill = buf.filled().len();
231        // Read the raw data
232        let internal_poll = read.poll_read(cx, buf);
233
234        if matches!(internal_poll, Poll::Ready(Ok(()))) {
235            // Decrypt the raw data in-place, note that our block size is 1 byte, so this is always safe
236            for block in buf.filled_mut()[original_fill..].chunks_mut(Aes128Cfb8Dec::block_size()) {
237                cipher.decrypt_block(block.try_into().unwrap());
238            }
239        }
240
241        internal_poll
242    }
243}
steel_protocol/utils.rs

steel_protocol/
utils.rs
