azalea_protocol/
write.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
//! Write packets to a stream.

use std::{fmt::Debug, io::Read};

use azalea_buf::McBufVarWritable;
use azalea_crypto::Aes128CfbEnc;
use flate2::{bufread::ZlibEncoder, Compression};
use thiserror::Error;
use tokio::io::{AsyncWrite, AsyncWriteExt};
use tracing::trace;

use crate::{packets::ProtocolPacket, read::MAXIMUM_UNCOMPRESSED_LENGTH};

/// Prepend the length of the packet to it.
fn frame_prepender(mut data: Vec<u8>) -> Result<Vec<u8>, std::io::Error> {
    let mut buf = Vec::new();
    (data.len() as u32).var_write_into(&mut buf)?;
    buf.append(&mut data);
    Ok(buf)
}

#[derive(Error, Debug)]
pub enum PacketEncodeError {
    #[error("{0}")]
    Io(#[from] std::io::Error),
    #[error("Packet too big (is {actual} bytes, should be less than {maximum}): {packet_string}")]
    TooBig {
        actual: usize,
        maximum: usize,
        packet_string: String,
    },
}

pub fn serialize_packet<P: ProtocolPacket + Debug>(
    packet: &P,
) -> Result<Vec<u8>, PacketEncodeError> {
    let mut buf = Vec::new();
    packet.id().var_write_into(&mut buf)?;
    packet.write(&mut buf)?;
    if buf.len() > MAXIMUM_UNCOMPRESSED_LENGTH as usize {
        return Err(PacketEncodeError::TooBig {
            actual: buf.len(),
            maximum: MAXIMUM_UNCOMPRESSED_LENGTH as usize,
            packet_string: format!("{packet:?}"),
        });
    }
    Ok(buf)
}

#[derive(Error, Debug)]
pub enum PacketCompressError {
    #[error("{0}")]
    Io(#[from] std::io::Error),
}

pub fn compression_encoder(
    data: &[u8],
    compression_threshold: u32,
) -> Result<Vec<u8>, PacketCompressError> {
    let n = data.len();
    // if it's less than the compression threshold, don't compress
    if n < compression_threshold as usize {
        let mut buf = Vec::new();
        0.var_write_into(&mut buf)?;
        std::io::Write::write_all(&mut buf, data)?;
        Ok(buf)
    } else {
        // otherwise, compress
        let mut deflater = ZlibEncoder::new(data, Compression::default());
        // write deflated data to buf
        let mut compressed_data = Vec::new();
        deflater.read_to_end(&mut compressed_data)?;

        // prepend the length
        let mut len_prepended_compressed_data = Vec::new();
        (data.len() as u32).var_write_into(&mut len_prepended_compressed_data)?;
        len_prepended_compressed_data.append(&mut compressed_data);

        Ok(len_prepended_compressed_data)
    }
}

pub async fn write_packet<P, W>(
    packet: &P,
    stream: &mut W,
    compression_threshold: Option<u32>,
    cipher: &mut Option<Aes128CfbEnc>,
) -> std::io::Result<()>
where
    P: ProtocolPacket + Debug,
    W: AsyncWrite + Unpin + Send,
{
    trace!("Sending packet: {packet:?}");
    let raw_packet = serialize_packet(packet).unwrap();
    write_raw_packet(&raw_packet, stream, compression_threshold, cipher).await
}

pub async fn write_raw_packet<W>(
    raw_packet: &[u8],
    stream: &mut W,
    compression_threshold: Option<u32>,
    cipher: &mut Option<Aes128CfbEnc>,
) -> std::io::Result<()>
where
    W: AsyncWrite + Unpin + Send,
{
    trace!("Writing raw packet: {raw_packet:?}");
    let mut raw_packet = raw_packet.to_vec();
    if let Some(threshold) = compression_threshold {
        raw_packet = compression_encoder(&raw_packet, threshold).unwrap();
    }
    raw_packet = frame_prepender(raw_packet).unwrap();
    // if we were given a cipher, encrypt the packet
    if let Some(cipher) = cipher {
        azalea_crypto::encrypt_packet(cipher, &mut raw_packet);
    }
    stream.write_all(&raw_packet).await
}