azalea_protocol/
write.rs

1//! Write packets to a stream.
2
3use std::{
4    fmt::Debug,
5    io::{self, Read},
6};
7
8use azalea_buf::AzaleaWriteVar;
9use azalea_crypto::Aes128CfbEnc;
10use flate2::{bufread::ZlibEncoder, Compression};
11use thiserror::Error;
12use tokio::io::{AsyncWrite, AsyncWriteExt};
13use tracing::trace;
14
15use crate::{packets::ProtocolPacket, read::MAXIMUM_UNCOMPRESSED_LENGTH};
16
17pub async fn write_packet<P, W>(
18    packet: &P,
19    stream: &mut W,
20    compression_threshold: Option<u32>,
21    cipher: &mut Option<Aes128CfbEnc>,
22) -> io::Result<()>
23where
24    P: ProtocolPacket + Debug,
25    W: AsyncWrite + Unpin + Send,
26{
27    trace!("Sending packet: {packet:?}");
28    let raw_packet = serialize_packet(packet).unwrap();
29    write_raw_packet(&raw_packet, stream, compression_threshold, cipher).await
30}
31
32pub fn serialize_packet<P: ProtocolPacket + Debug>(
33    packet: &P,
34) -> Result<Box<[u8]>, PacketEncodeError> {
35    let mut buf = Vec::new();
36    packet.id().azalea_write_var(&mut buf)?;
37    packet.write(&mut buf)?;
38    if buf.len() > MAXIMUM_UNCOMPRESSED_LENGTH as usize {
39        return Err(PacketEncodeError::TooBig {
40            actual: buf.len(),
41            maximum: MAXIMUM_UNCOMPRESSED_LENGTH as usize,
42            packet_string: format!("{packet:?}"),
43        });
44    }
45    Ok(buf.into_boxed_slice())
46}
47
48pub async fn write_raw_packet<W>(
49    raw_packet: &[u8],
50    stream: &mut W,
51    compression_threshold: Option<u32>,
52    cipher: &mut Option<Aes128CfbEnc>,
53) -> io::Result<()>
54where
55    W: AsyncWrite + Unpin + Send,
56{
57    trace!("Writing raw packet: {raw_packet:?}");
58    let mut raw_packet = raw_packet.to_vec();
59    if let Some(threshold) = compression_threshold {
60        raw_packet = compression_encoder(&raw_packet, threshold).unwrap();
61    }
62    raw_packet = frame_prepender(raw_packet).unwrap();
63    // if we were given a cipher, encrypt the packet
64    if let Some(cipher) = cipher {
65        azalea_crypto::encrypt_packet(cipher, &mut raw_packet);
66    }
67    stream.write_all(&raw_packet).await
68}
69
70pub fn compression_encoder(
71    data: &[u8],
72    compression_threshold: u32,
73) -> Result<Vec<u8>, PacketCompressError> {
74    let n = data.len();
75    // if it's less than the compression threshold, don't compress
76    if n < compression_threshold as usize {
77        let mut buf = Vec::new();
78        0_u32.azalea_write_var(&mut buf)?;
79        io::Write::write_all(&mut buf, data)?;
80        Ok(buf)
81    } else {
82        // otherwise, compress
83        let mut deflater = ZlibEncoder::new(data, Compression::default());
84
85        // write deflated data to buf
86        let mut compressed_data = Vec::new();
87        deflater.read_to_end(&mut compressed_data)?;
88
89        // prepend the length
90        let mut len_prepended_compressed_data = Vec::new();
91        (data.len() as u32).azalea_write_var(&mut len_prepended_compressed_data)?;
92        len_prepended_compressed_data.append(&mut compressed_data);
93
94        Ok(len_prepended_compressed_data)
95    }
96}
97
98/// Prepend the length of the packet to it.
99fn frame_prepender(mut data: Vec<u8>) -> Result<Vec<u8>, io::Error> {
100    let mut buf = Vec::new();
101    (data.len() as u32).azalea_write_var(&mut buf)?;
102    buf.append(&mut data);
103    Ok(buf)
104}
105
106#[derive(Error, Debug)]
107pub enum PacketEncodeError {
108    #[error("{0}")]
109    Io(#[from] io::Error),
110    #[error("Packet too big (is {actual} bytes, should be less than {maximum}): {packet_string}")]
111    TooBig {
112        actual: usize,
113        maximum: usize,
114        packet_string: String,
115    },
116}
117
118#[derive(Error, Debug)]
119pub enum PacketCompressError {
120    #[error("{0}")]
121    Io(#[from] io::Error),
122}
azalea_protocol/write.rs

azalea_protocol/
write.rs
