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::{Compression, bufread::ZlibEncoder};
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    let network_packet = encode_to_network_packet(raw_packet, compression_threshold, cipher);
58    stream.write_all(&network_packet).await
59}
60
61pub fn encode_to_network_packet(
62    raw_packet: &[u8],
63    compression_threshold: Option<u32>,
64    cipher: &mut Option<Aes128CfbEnc>,
65) -> Vec<u8> {
66    trace!("Writing raw packet: {raw_packet:?}");
67    let mut raw_packet = raw_packet.to_vec();
68    if let Some(threshold) = compression_threshold {
69        raw_packet = compression_encoder(&raw_packet, threshold).unwrap();
70    }
71    raw_packet = frame_prepender(raw_packet).unwrap();
72    // if we were given a cipher, encrypt the packet
73    if let Some(cipher) = cipher {
74        azalea_crypto::encrypt_packet(cipher, &mut raw_packet);
75    }
76    raw_packet
77}
78
79pub fn compression_encoder(
80    data: &[u8],
81    compression_threshold: u32,
82) -> Result<Vec<u8>, PacketCompressError> {
83    let n = data.len();
84    // if it's less than the compression threshold, don't compress
85    if n < compression_threshold as usize {
86        let mut buf = Vec::new();
87        0_u32.azalea_write_var(&mut buf)?;
88        io::Write::write_all(&mut buf, data)?;
89        Ok(buf)
90    } else {
91        // otherwise, compress
92        let mut deflater = ZlibEncoder::new(data, Compression::default());
93
94        // write deflated data to buf
95        let mut compressed_data = Vec::new();
96        deflater.read_to_end(&mut compressed_data)?;
97
98        // prepend the length
99        let mut len_prepended_compressed_data = Vec::new();
100        (data.len() as u32).azalea_write_var(&mut len_prepended_compressed_data)?;
101        len_prepended_compressed_data.append(&mut compressed_data);
102
103        Ok(len_prepended_compressed_data)
104    }
105}
106
107/// Prepend the length of the packet to it.
108fn frame_prepender(mut data: Vec<u8>) -> Result<Vec<u8>, io::Error> {
109    let mut buf = Vec::new();
110    (data.len() as u32).azalea_write_var(&mut buf)?;
111    buf.append(&mut data);
112    Ok(buf)
113}
114
115#[derive(Error, Debug)]
116pub enum PacketEncodeError {
117    #[error("{0}")]
118    Io(#[from] io::Error),
119    #[error("Packet too big (is {actual} bytes, should be less than {maximum}): {packet_string}")]
120    TooBig {
121        actual: usize,
122        maximum: usize,
123        packet_string: String,
124    },
125}
126
127#[derive(Error, Debug)]
128pub enum PacketCompressError {
129    #[error("{0}")]
130    Io(#[from] io::Error),
131}
azalea_protocol/write.rs

azalea_protocol/
write.rs
