ms-ndr 0.1.2

//! NDR (Network Data Representation) marshaling — the DCE/RPC Transfer Syntax, little-endian.
//!
//! The foundation every MS-RPCE interface rides on. Implements the pieces the
//! SAMR/LSAT/EPM/DRSUAPI/WMI calls need: aligned primitives, conformant & varying
//! arrays, unique (referent) pointers, and conformant-varying UTF-16 strings
//! (the `RPC_UNICODE_STRING` payload shape).
//!
//! Alignment is relative to the start of the stub buffer, matching the NDR rule that an
//! aggregate aligns to its largest member. This crate is transport- and interface-agnostic:
//! it only turns Rust values into NDR bytes and back.

/// The only way marshaling fails: the decoder ran off the end of the stub buffer.
#[derive(Debug, thiserror::Error)]
pub enum NdrError {
    #[error("NDR buffer underrun: need {need} bytes at offset {pos}")]
    Underrun { need: usize, pos: usize },
}

pub type Result<T> = std::result::Result<T, NdrError>;

/// Serializer. Alignment is relative to the start of the stub buffer.
#[derive(Default)]
pub struct NdrEncoder {
    buf: Vec<u8>,
    next_referent: u32,
}

impl NdrEncoder {
    pub fn new() -> Self {
        NdrEncoder {
            buf: Vec::new(),
            next_referent: 0x0002_0000,
        }
    }

    /// Pad with zero bytes up to an `a`-byte boundary. Public so callers can align a union
    /// arm or struct to its natural alignment (NDR aligns an aggregate to its largest member).
    pub fn align(&mut self, a: usize) {
        if a <= 1 {
            return;
        }
        while self.buf.len() % a != 0 {
            self.buf.push(0);
        }
    }

    pub fn u8(&mut self, v: u8) {
        self.buf.push(v);
    }
    pub fn u16(&mut self, v: u16) {
        self.align(2);
        self.buf.extend_from_slice(&v.to_le_bytes());
    }
    pub fn u32(&mut self, v: u32) {
        self.align(4);
        self.buf.extend_from_slice(&v.to_le_bytes());
    }
    pub fn u64(&mut self, v: u64) {
        self.align(8);
        self.buf.extend_from_slice(&v.to_le_bytes());
    }
    pub fn bytes(&mut self, b: &[u8]) {
        self.buf.extend_from_slice(b);
    }
    /// A UUID / GUID field (DCE wire layout, 16 bytes, aligned to 4).
    pub fn uuid(&mut self, b: &[u8; 16]) {
        self.align(4);
        self.buf.extend_from_slice(b);
    }

    /// A non-null unique pointer: emit a fresh referent id. Caller emits the pointee next.
    pub fn referent(&mut self) -> u32 {
        let id = self.next_referent;
        self.next_referent += 4;
        self.u32(id);
        id
    }
    pub fn null_ptr(&mut self) {
        self.u32(0);
    }

    /// A conformant+varying wide string: max_count, offset(0), actual_count, then UTF-16LE
    /// including the terminating NUL. Matches the `[size_is][length_is]` wchar payload.
    pub fn conformant_varying_wstr(&mut self, s: &str) {
        let mut units: Vec<u16> = s.encode_utf16().collect();
        units.push(0);
        let n = units.len() as u32;
        self.u32(n); // max_count
        self.u32(0); // offset
        self.u32(n); // actual_count
        for u in units {
            self.u16(u);
        }
    }

    pub fn into_bytes(self) -> Vec<u8> {
        self.buf
    }
    pub fn len(&self) -> usize {
        self.buf.len()
    }
    pub fn is_empty(&self) -> bool {
        self.buf.is_empty()
    }
}

/// Deserializer over a stub buffer.
pub struct NdrDecoder<'a> {
    buf: &'a [u8],
    pos: usize,
}

impl<'a> NdrDecoder<'a> {
    pub fn new(buf: &'a [u8]) -> Self {
        NdrDecoder { buf, pos: 0 }
    }

    pub fn align(&mut self, a: usize) {
        if a <= 1 {
            return;
        }
        let remainder = self.pos % a;
        if remainder != 0 {
            self.pos = self.pos.saturating_add(a - remainder);
        }
    }
    fn take(&mut self, n: usize) -> Result<&'a [u8]> {
        let end = self.pos.checked_add(n).ok_or(NdrError::Underrun {
            need: n,
            pos: self.pos,
        })?;
        if end > self.buf.len() {
            return Err(NdrError::Underrun {
                need: n,
                pos: self.pos,
            });
        }
        let s = &self.buf[self.pos..end];
        self.pos = end;
        Ok(s)
    }

    pub fn u8(&mut self) -> Result<u8> {
        Ok(self.take(1)?[0])
    }
    pub fn u16(&mut self) -> Result<u16> {
        self.align(2);
        Ok(u16::from_le_bytes(self.take(2)?.try_into().unwrap()))
    }
    pub fn u32(&mut self) -> Result<u32> {
        self.align(4);
        Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap()))
    }
    pub fn u64(&mut self) -> Result<u64> {
        self.align(8);
        Ok(u64::from_le_bytes(self.take(8)?.try_into().unwrap()))
    }
    pub fn uuid(&mut self) -> Result<[u8; 16]> {
        self.align(4);
        Ok(self.take(16)?.try_into().unwrap())
    }

    /// Read a conformant+varying wide string; returns it without the trailing NUL.
    ///
    /// # Bounded allocation
    ///
    /// The `actual_count` field is attacker-controlled (comes off the wire).
    /// A naive `Vec::with_capacity(actual)` here would let a remote peer
    /// trigger a multi-GB allocation before we notice the buffer is short.
    /// This decoder preflights the size against the remaining buffer: if
    /// `actual * 2` bytes aren't physically available, we return
    /// `NdrError::Underrun` before touching the allocator.
    pub fn conformant_varying_wstr(&mut self) -> Result<String> {
        let max = self.u32()?;
        let offset = self.u32()?;
        let actual = self.u32()? as usize;
        if offset > max || actual > (max - offset) as usize {
            return Err(NdrError::Underrun {
                need: actual,
                pos: self.pos,
            });
        }
        // Bounded preflight: refuse to allocate for more bytes than the
        // buffer can supply. Also guards against overflow on 32-bit hosts.
        let need = actual.checked_mul(2).ok_or(NdrError::Underrun {
            need: usize::MAX,
            pos: self.pos,
        })?;
        if need > self.remaining() {
            return Err(NdrError::Underrun {
                need,
                pos: self.pos,
            });
        }
        let mut units = Vec::with_capacity(actual);
        for _ in 0..actual {
            units.push(self.u16()?);
        }
        while units.last() == Some(&0) {
            units.pop();
        }
        Ok(String::from_utf16_lossy(&units))
    }

    /// Read `n` raw bytes with no alignment (e.g. a fixed context handle or tower blob).
    pub fn read_bytes(&mut self, n: usize) -> Result<&'a [u8]> {
        self.take(n)
    }

    pub fn remaining(&self) -> usize {
        self.buf.len().saturating_sub(self.pos)
    }
    pub fn position(&self) -> usize {
        self.pos
    }
}

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

    #[test]
    fn primitive_alignment() {
        let mut e = NdrEncoder::new();
        e.u8(0x01);
        e.u32(0xAABB_CCDD); // must align to 4 → 3 pad bytes
        assert_eq!(e.into_bytes(), vec![0x01, 0, 0, 0, 0xDD, 0xCC, 0xBB, 0xAA]);
    }

    #[test]
    fn wstr_roundtrip() {
        let mut e = NdrEncoder::new();
        e.conformant_varying_wstr("ADHAMMER");
        let bytes = e.into_bytes();
        // max/offset/actual each 4 bytes = 12, then 9 u16 (8 chars + NUL) = 18 → 30 bytes.
        assert_eq!(bytes.len(), 30);
        let mut d = NdrDecoder::new(&bytes);
        assert_eq!(d.conformant_varying_wstr().unwrap(), "ADHAMMER");
        assert_eq!(d.remaining(), 0);
    }

    #[test]
    fn referent_ids_are_nonzero_and_advance() {
        let mut e = NdrEncoder::new();
        let a = e.referent();
        let b = e.referent();
        assert_ne!(a, 0);
        assert_eq!(b, a + 4);
    }

    #[test]
    fn underrun_reports_offset() {
        let bytes = [0x01, 0x02]; // only 2 bytes
        let mut d = NdrDecoder::new(&bytes);
        assert!(matches!(
            d.u32(),
            Err(NdrError::Underrun { need: 4, pos: 0 })
        ));
    }

    /// Regression: a hostile wire response with `actual_count = 0x7FFF_FFFF`
    /// used to reach `Vec::with_capacity(2_147_483_647)` — a multi-GB
    /// allocation for zero payload bytes. The bounded preflight in
    /// `conformant_varying_wstr` must reject this in microseconds, not
    /// after touching the allocator.
    #[test]
    fn wstr_rejects_hostile_actual_count() {
        let mut bytes = Vec::new();
        bytes.extend_from_slice(&0u32.to_le_bytes()); // max
        bytes.extend_from_slice(&0u32.to_le_bytes()); // offset
        bytes.extend_from_slice(&0x7FFF_FFFFu32.to_le_bytes()); // actual (hostile)
                                                                // no wchar payload follows
        let mut d = NdrDecoder::new(&bytes);
        let start = std::time::Instant::now();
        let r = d.conformant_varying_wstr();
        let elapsed_ms = start.elapsed().as_millis();
        assert!(
            matches!(r, Err(NdrError::Underrun { .. })),
            "expected Underrun, got {r:?}"
        );
        assert!(
            elapsed_ms < 50,
            "preflight must reject in <50ms, took {elapsed_ms}ms — probably allocated"
        );
    }

    /// Boundary: `actual_count` that would overflow `usize::checked_mul(2)`
    /// on 32-bit targets. On 64-bit it's just a very large `need`; both
    /// must return Underrun cleanly, never panic.
    #[test]
    fn wstr_rejects_actual_count_overflow() {
        let mut bytes = Vec::new();
        bytes.extend_from_slice(&0u32.to_le_bytes());
        bytes.extend_from_slice(&0u32.to_le_bytes());
        bytes.extend_from_slice(&0xFFFF_FFFFu32.to_le_bytes()); // max u32
        let mut d = NdrDecoder::new(&bytes);
        assert!(matches!(
            d.conformant_varying_wstr(),
            Err(NdrError::Underrun { .. })
        ));
    }

    #[test]
    fn zero_alignment_is_a_noop() {
        let mut e = NdrEncoder::new();
        e.u8(1);
        e.align(0);
        assert_eq!(e.into_bytes(), vec![1]);

        let bytes = [1u8];
        let mut d = NdrDecoder::new(&bytes);
        d.align(0);
        assert_eq!(d.u8().unwrap(), 1);
    }

    #[test]
    fn alignment_past_end_reports_underrun() {
        let bytes = [1u8];
        let mut d = NdrDecoder::new(&bytes);
        assert_eq!(d.u8().unwrap(), 1);
        d.align(4);
        assert_eq!(d.remaining(), 0);
        assert!(matches!(
            d.u32(),
            Err(NdrError::Underrun { need: 4, pos: 4 })
        ));
    }

    #[test]
    fn wstr_rejects_inconsistent_varying_bounds() {
        let mut bytes = Vec::new();
        bytes.extend_from_slice(&2u32.to_le_bytes()); // max
        bytes.extend_from_slice(&2u32.to_le_bytes()); // offset == max
        bytes.extend_from_slice(&1u32.to_le_bytes()); // actual exceeds max - offset
        bytes.extend_from_slice(&0u16.to_le_bytes());
        let mut d = NdrDecoder::new(&bytes);
        assert!(d.conformant_varying_wstr().is_err());
    }
}