borer-core 0.5.9

use anyhow::Context;
use bytes::{BufMut, BytesMut};
use rand::{Rng, RngCore};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};

/// Trojan padding frame used to disguise connection metadata.
#[derive(Debug)]
pub struct Padding {
    len: u16,
}

impl Padding {
    /// Create a padding frame with a random length in `start..end`.
    pub fn new(start: u16, end: u16) -> Self {
        let mut rng = rand::rng();
        let len = rng.random_range(start..end);
        Self { len }
    }

    /// Read a padding frame from the stream.
    pub async fn read_from<R>(stream: &mut R) -> anyhow::Result<Self>
    where
        R: AsyncRead + Unpin,
    {
        let len = stream.read_u16().await?;
        let mut buf = vec![0; len as usize];
        stream.read_exact(&mut buf).await?;
        Ok(Padding { len })
    }

    /// Write this padding frame to the stream.
    pub async fn write_to<W>(&self, w: &mut W) -> anyhow::Result<()>
    where
        W: AsyncWrite + Unpin,
    {
        let mut buf = BytesMut::with_capacity(self.serialized_len());
        self.write_to_buf(&mut buf);
        w.write_all(&buf)
            .await
            .context("padding Write buf failed")?;

        Ok(())
    }

    pub fn write_to_buf<B: BufMut>(&self, buf: &mut B) {
        let rand_buf = self.rand_buf();
        buf.put_u16(self.len);
        buf.put_slice(&rand_buf);
    }

    pub fn serialized_len(&self) -> usize {
        2 + self.len as usize
    }

    pub fn rand_buf(&self) -> Vec<u8> {
        let mut rng = rand::rng();
        let mut buf = vec![0; self.len as usize];
        rng.fill_bytes(&mut buf);
        buf
    }
}

impl Default for Padding {
    fn default() -> Self {
        Self::new(256, 2048)
    }
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;

    use bytes::{Buf, BytesMut};

    use super::Padding;

    #[test]
    fn serialized_len_matches_written_buffer_length() {
        let padding = Padding::new(8, 9);
        let mut buf = BytesMut::new();

        padding.write_to_buf(&mut buf);

        assert_eq!(buf.len(), padding.serialized_len());
        assert_eq!(buf.get_u16() as usize + 2, padding.serialized_len());
    }

    #[tokio::test]
    async fn read_from_round_trips_written_bytes() {
        let padding = Padding::new(8, 9);
        let mut raw = BytesMut::new();
        padding.write_to_buf(&mut raw);
        let expected_len = padding.serialized_len();

        let mut cursor = Cursor::new(raw.to_vec());
        let actual = Padding::read_from(&mut cursor).await.unwrap();

        assert_eq!(actual.serialized_len(), expected_len);
    }
}