codetics-ocpp-ws 0.1.0

// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 Kevin Vavelin, Codetics (https://codetics.fr)

//! RFC 6455 §5 base framing.

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OpCode {
    Continuation,
    Text,
    Binary,
    Close,
    Ping,
    Pong,
}

impl OpCode {
    pub fn from_u8(b: u8) -> Option<OpCode> {
        Some(match b {
            0x0 => OpCode::Continuation,
            0x1 => OpCode::Text,
            0x2 => OpCode::Binary,
            0x8 => OpCode::Close,
            0x9 => OpCode::Ping,
            0xA => OpCode::Pong,
            _ => return None,
        })
    }

    pub fn as_u8(self) -> u8 {
        match self {
            OpCode::Continuation => 0x0,
            OpCode::Text => 0x1,
            OpCode::Binary => 0x2,
            OpCode::Close => 0x8,
            OpCode::Ping => 0x9,
            OpCode::Pong => 0xA,
        }
    }

    pub fn is_control(self) -> bool {
        matches!(self, OpCode::Close | OpCode::Ping | OpCode::Pong)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Header {
    pub fin: bool,
    pub rsv1: bool,
    pub rsv2: bool,
    pub rsv3: bool,
    pub opcode: u8,
    pub mask: Option<[u8; 4]>,
    pub payload_len: u64,
    pub header_len: usize,
}

/// Parses a frame header. Returns `Ok(None)` when more bytes are needed.
pub fn parse_header(buf: &[u8]) -> Result<Option<Header>, &'static str> {
    if buf.len() < 2 {
        return Ok(None);
    }
    let b0 = buf[0];
    let b1 = buf[1];
    let masked = b1 & 0x80 != 0;
    let len7 = (b1 & 0x7F) as u64;
    let mut pos = 2;
    let payload_len = match len7 {
        126 => {
            if buf.len() < pos + 2 {
                return Ok(None);
            }
            let l = u16::from_be_bytes([buf[2], buf[3]]) as u64;
            pos += 2;
            if l < 126 {
                return Err("non-minimal payload length encoding");
            }
            l
        }
        127 => {
            if buf.len() < pos + 8 {
                return Ok(None);
            }
            let l = u64::from_be_bytes(buf[2..10].try_into().expect("8 bytes"));
            pos += 8;
            if l >> 63 != 0 {
                return Err("payload length has most significant bit set");
            }
            if l <= 0xFFFF {
                return Err("non-minimal payload length encoding");
            }
            l
        }
        l => l,
    };
    let mask = if masked {
        if buf.len() < pos + 4 {
            return Ok(None);
        }
        let m = [buf[pos], buf[pos + 1], buf[pos + 2], buf[pos + 3]];
        pos += 4;
        Some(m)
    } else {
        None
    };
    Ok(Some(Header {
        fin: b0 & 0x80 != 0,
        rsv1: b0 & 0x40 != 0,
        rsv2: b0 & 0x20 != 0,
        rsv3: b0 & 0x10 != 0,
        opcode: b0 & 0x0F,
        mask,
        payload_len,
        header_len: pos,
    }))
}

/// Applies (or removes) a masking key in place.
pub fn apply_mask(data: &mut [u8], mask: [u8; 4]) {
    for (i, b) in data.iter_mut().enumerate() {
        *b ^= mask[i & 3];
    }
}

/// Encodes a single (final) frame into `out`.
pub fn encode(
    out: &mut Vec<u8>,
    opcode: OpCode,
    rsv1: bool,
    payload: &[u8],
    mask: Option<[u8; 4]>,
) {
    out.push(0x80 | if rsv1 { 0x40 } else { 0 } | opcode.as_u8());
    let mask_bit = if mask.is_some() { 0x80 } else { 0 };
    match payload.len() {
        n if n < 126 => out.push(mask_bit | n as u8),
        n if n <= 0xFFFF => {
            out.push(mask_bit | 126);
            out.extend_from_slice(&(n as u16).to_be_bytes());
        }
        n => {
            out.push(mask_bit | 127);
            out.extend_from_slice(&(n as u64).to_be_bytes());
        }
    }
    match mask {
        Some(m) => {
            out.extend_from_slice(&m);
            let start = out.len();
            out.extend_from_slice(payload);
            apply_mask(&mut out[start..], m);
        }
        None => out.extend_from_slice(payload),
    }
}

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

    #[test]
    fn rfc6455_examples() {
        // §5.7: single-frame unmasked text "Hello"
        let h = parse_header(&[0x81, 0x05, 0x48, 0x65, 0x6c, 0x6c, 0x6f])
            .unwrap()
            .unwrap();
        assert!(
            h.fin && h.mask.is_none() && h.opcode == 1 && h.payload_len == 5 && h.header_len == 2
        );
        // §5.7: single-frame masked text "Hello"
        let mut f = [
            0x81, 0x85, 0x37, 0xfa, 0x21, 0x3d, 0x7f, 0x9f, 0x4d, 0x51, 0x58,
        ];
        let h = parse_header(&f).unwrap().unwrap();
        let m = h.mask.unwrap();
        apply_mask(&mut f[6..], m);
        assert_eq!(&f[6..], b"Hello");
        // 256-byte binary
        let mut out = Vec::new();
        encode(&mut out, OpCode::Binary, false, &[0u8; 256], None);
        assert_eq!(&out[..4], &[0x82, 0x7E, 0x01, 0x00]);
        // 64 KiB binary
        out.clear();
        encode(&mut out, OpCode::Binary, false, &vec![0u8; 65536], None);
        assert_eq!(&out[..10], &[0x82, 0x7F, 0, 0, 0, 0, 0, 1, 0, 0]);
        assert_eq!(parse_header(&out).unwrap().unwrap().payload_len, 65536);
        // incomplete
        assert_eq!(parse_header(&[0x81]).unwrap(), None);
        assert_eq!(parse_header(&[0x82, 0x7F, 0]).unwrap(), None);
    }
}