SoulStream/crates/soulseek-proto/src/wire.rs

//! Low-level message building and reading, mirroring Soulseek.NET's
//! `MessageBuilder` / `MessageReader`.

use bytes::{Buf, BufMut, Bytes, BytesMut};
use flate2::read::ZlibDecoder;
use std::io::Read;

use crate::{ProtoError, Result};

/// Builds one outgoing message. `finish_*` prepends the length header.
pub struct Writer {
    buf: BytesMut,
}

impl Default for Writer {
    fn default() -> Self {
        Self::new()
    }
}

impl Writer {
    pub fn new() -> Self {
        Writer {
            buf: BytesMut::with_capacity(64),
        }
    }

    pub fn u8(mut self, v: u8) -> Self {
        self.buf.put_u8(v);
        self
    }

    pub fn u32(mut self, v: u32) -> Self {
        self.buf.put_u32_le(v);
        self
    }

    pub fn u64(mut self, v: u64) -> Self {
        self.buf.put_u64_le(v);
        self
    }

    pub fn bool(self, v: bool) -> Self {
        self.u8(v as u8)
    }

    pub fn string(mut self, s: &str) -> Self {
        self.buf.put_u32_le(s.len() as u32);
        self.buf.put_slice(s.as_bytes());
        self
    }

    /// Frame with a u32 code: `[len][code:u32][payload]`.
    pub fn finish_u32(self, code: u32) -> Vec<u8> {
        let mut out = Vec::with_capacity(self.buf.len() + 8);
        out.extend_from_slice(&((self.buf.len() as u32 + 4).to_le_bytes()));
        out.extend_from_slice(&code.to_le_bytes());
        out.extend_from_slice(&self.buf);
        out
    }

    /// Frame with a u32 code and a zlib-compressed payload:
    /// `[len][code:u32][zlib(payload)]` (search responses, browse responses).
    pub fn finish_u32_compressed(self, code: u32) -> Vec<u8> {
        use flate2::write::ZlibEncoder;
        use flate2::Compression;
        use std::io::Write;

        let mut enc = ZlibEncoder::new(Vec::new(), Compression::default());
        enc.write_all(&self.buf).expect("zlib write");
        let compressed = enc.finish().expect("zlib finish");

        let mut out = Vec::with_capacity(compressed.len() + 8);
        out.extend_from_slice(&((compressed.len() as u32 + 4).to_le_bytes()));
        out.extend_from_slice(&code.to_le_bytes());
        out.extend_from_slice(&compressed);
        out
    }

    /// Frame with a single-byte code: `[len][code:u8][payload]`
    /// (PeerInit / PierceFirewall only).
    pub fn finish_u8(self, code: u8) -> Vec<u8> {
        let mut out = Vec::with_capacity(self.buf.len() + 5);
        out.extend_from_slice(&((self.buf.len() as u32 + 1).to_le_bytes()));
        out.push(code);
        out.extend_from_slice(&self.buf);
        out
    }
}

/// Reads one message payload (the bytes after the code).
pub struct Reader {
    buf: Bytes,
}

impl Reader {
    pub fn new(payload: Bytes) -> Self {
        Reader { buf: payload }
    }

    pub fn remaining(&self) -> usize {
        self.buf.remaining()
    }

    pub fn has_more(&self) -> bool {
        self.buf.has_remaining()
    }

    /// Take all remaining bytes.
    pub fn rest(&mut self) -> Bytes {
        let n = self.buf.remaining();
        self.buf.copy_to_bytes(n)
    }

    pub fn u8(&mut self) -> Result<u8> {
        if self.buf.remaining() < 1 {
            return Err(ProtoError::UnexpectedEof);
        }
        Ok(self.buf.get_u8())
    }

    pub fn bool(&mut self) -> Result<bool> {
        Ok(self.u8()? != 0)
    }

    pub fn u32(&mut self) -> Result<u32> {
        if self.buf.remaining() < 4 {
            return Err(ProtoError::UnexpectedEof);
        }
        Ok(self.buf.get_u32_le())
    }

    pub fn u64(&mut self) -> Result<u64> {
        if self.buf.remaining() < 8 {
            return Err(ProtoError::UnexpectedEof);
        }
        Ok(self.buf.get_u64_le())
    }

    /// IPv4 address; wire order is reversed relative to display order.
    pub fn ipv4(&mut self) -> Result<std::net::Ipv4Addr> {
        if self.buf.remaining() < 4 {
            return Err(ProtoError::UnexpectedEof);
        }
        let mut b = [0u8; 4];
        self.buf.copy_to_slice(&mut b);
        Ok(std::net::Ipv4Addr::new(b[3], b[2], b[1], b[0]))
    }

    /// UTF-8 with Latin-1 fallback, matching what real-world peers send.
    pub fn string(&mut self) -> Result<String> {
        let len = self.u32()? as usize;
        if self.buf.remaining() < len {
            return Err(ProtoError::UnexpectedEof);
        }
        let bytes = self.buf.copy_to_bytes(len);
        Ok(match std::str::from_utf8(&bytes) {
            Ok(s) => s.to_owned(),
            Err(_) => bytes.iter().map(|&b| b as char).collect(),
        })
    }

    /// A length-prefixed raw byte string (a user picture, for one): the same
    /// framing as `string`, without the UTF-8 interpretation.
    pub fn bytes(&mut self) -> Result<Bytes> {
        let len = self.u32()? as usize;
        if self.buf.remaining() < len {
            return Err(ProtoError::UnexpectedEof);
        }
        Ok(self.buf.copy_to_bytes(len))
    }

    /// Replace the remaining payload with its zlib-decompressed form.
    pub fn decompress(&mut self) -> Result<()> {
        let mut out = Vec::new();
        ZlibDecoder::new(&self.buf[..]).read_to_end(&mut out)?;
        if out.len() > crate::MAX_MESSAGE_SIZE {
            return Err(ProtoError::TooLarge(out.len()));
        }
        self.buf = Bytes::from(out);
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn roundtrip_primitives() {
        let msg = Writer::new().u32(7).string("héllo").u64(9).finish_u32(26);
        // len = 4 (code) + 4 + 4+6 + 8 = 26
        assert_eq!(&msg[0..4], &26u32.to_le_bytes());
        assert_eq!(&msg[4..8], &26u32.to_le_bytes());

        let mut r = Reader::new(Bytes::copy_from_slice(&msg[8..]));
        assert_eq!(r.u32().unwrap(), 7);
        assert_eq!(r.string().unwrap(), "héllo");
        assert_eq!(r.u64().unwrap(), 9);
        assert!(!r.has_more());
    }

    #[test]
    fn u8_code_frame() {
        let msg = Writer::new().u32(123).finish_u8(0);
        assert_eq!(&msg[0..4], &5u32.to_le_bytes());
        assert_eq!(msg[4], 0);
    }

    #[test]
    fn latin1_fallback() {
        let mut raw = Vec::new();
        raw.extend_from_slice(&2u32.to_le_bytes());
        raw.extend_from_slice(&[0xE9, 0x21]); // "é!" in Latin-1, invalid UTF-8
        let mut r = Reader::new(Bytes::from(raw));
        assert_eq!(r.string().unwrap(), "é!");
    }
}