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

azalea_protocol/
write.rs
