zero_mysql/protocol/

primitive.rs

use crate::nightly::{cold_path, unlikely};

use crate::error::{Error, Result, eyre};
use zerocopy::FromBytes;
use zerocopy::byteorder::little_endian::{U16 as U16LE, U32 as U32LE, U64 as U64LE};

/// Read 1-byte integer
pub fn read_int_1(data: &[u8]) -> Result<(u8, &[u8])> {
    let (&byte, rest) = data
        .split_first()
        .ok_or_else(|| Error::LibraryBug(eyre!("read_int_1: empty buffer")))?;
    Ok((byte, rest))
}

/// Read 2-byte little-endian integer
pub fn read_int_2(data: &[u8]) -> Result<(u16, &[u8])> {
    if unlikely(data.len() < 2) {
        return Err(Error::LibraryBug(eyre!(
            "read_int_2: buffer too short: {} < 2",
            data.len()
        )));
    }
    let value = U16LE::ref_from_bytes(&data[..2])?.get();
    Ok((value, &data[2..]))
}

/// Read 3-byte little-endian integer
pub fn read_int_3(data: &[u8]) -> Result<(u32, &[u8])> {
    let (chunk, rest) = data.split_first_chunk::<3>().ok_or_else(|| {
        Error::LibraryBug(eyre!("read_int_3: buffer too short: {} < 3", data.len()))
    })?;
    let value = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], 0]);
    Ok((value, rest))
}

/// Read 4-byte little-endian integer
pub fn read_int_4(data: &[u8]) -> Result<(u32, &[u8])> {
    if unlikely(data.len() < 4) {
        return Err(Error::LibraryBug(eyre!(
            "read_int_4: buffer too short: {} < 4",
            data.len()
        )));
    }
    let value = U32LE::ref_from_bytes(&data[..4])?.get();
    Ok((value, &data[4..]))
}

/// Read 6-byte little-endian integer
pub fn read_int_6(data: &[u8]) -> Result<(u64, &[u8])> {
    let (chunk, rest) = data.split_first_chunk::<6>().ok_or_else(|| {
        Error::LibraryBug(eyre!("read_int_6: buffer too short: {} < 6", data.len()))
    })?;
    let value = u64::from_le_bytes([
        chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], 0, 0,
    ]);
    Ok((value, rest))
}

/// Read 8-byte little-endian integer
pub fn read_int_8(data: &[u8]) -> Result<(u64, &[u8])> {
    if unlikely(data.len() < 8) {
        return Err(Error::LibraryBug(eyre!(
            "read_int_8: buffer too short: {} < 8",
            data.len()
        )));
    }
    let value = U64LE::ref_from_bytes(&data[..8])?.get();
    Ok((value, &data[8..]))
}

/// Read length-encoded integer
pub fn read_int_lenenc(data: &[u8]) -> Result<(u64, &[u8])> {
    match data.first() {
        Some(0xFC) => {
            let (val, rest) = read_int_2(&data[1..])?;
            Ok((val as u64, rest))
        }
        Some(0xFD) => {
            let (val, rest) = read_int_3(&data[1..])?;
            Ok((val as u64, rest))
        }
        Some(0xFE) => {
            let (val, rest) = read_int_8(&data[1..])?;
            Ok((val, rest))
        }
        Some(val) => Ok((*val as u64, &data[1..])),
        None => {
            cold_path();
            Err(Error::LibraryBug(eyre!("read_int_lenenc: empty buffer")))
        }
    }
}

/// Read fixed-length string
pub fn read_string_fix(data: &[u8], len: usize) -> Result<(&[u8], &[u8])> {
    if unlikely(data.len() < len) {
        return Err(Error::LibraryBug(eyre!(
            "read_string_fix: buffer too short: {} < {}",
            data.len(),
            len
        )));
    }
    Ok((&data[..len], &data[len..]))
}

/// Read null-terminated string
/// TODO: use memchr
pub fn read_string_null(data: &[u8]) -> Result<(&[u8], &[u8])> {
    for (i, &byte) in data.iter().enumerate() {
        if byte == 0 {
            return Ok((&data[..i], &data[i + 1..]));
        }
    }
    Err(Error::LibraryBug(eyre!(
        "read_string_null: no null terminator found"
    )))
}

/// Read length-encoded string
pub fn read_string_lenenc(data: &[u8]) -> Result<(&[u8], &[u8])> {
    let (len, rest) = read_int_lenenc(data)?;
    read_string_fix(rest, len as usize)
}

/// Write 1-byte integer
#[inline]
pub fn write_int_1(out: &mut Vec<u8>, value: u8) {
    out.push(value);
}

/// Write 2-byte little-endian integer
#[inline]
pub fn write_int_2(out: &mut Vec<u8>, value: u16) {
    out.extend_from_slice(&value.to_le_bytes());
}

/// Write 3-byte little-endian integer
#[inline]
pub fn write_int_3(out: &mut Vec<u8>, value: u32) {
    out.extend_from_slice(&value.to_le_bytes()[..3]);
}

/// Write 4-byte little-endian integer
#[inline]
pub fn write_int_4(out: &mut Vec<u8>, value: u32) {
    out.extend_from_slice(&value.to_le_bytes());
}

/// Write 8-byte little-endian integer
#[inline]
pub fn write_int_8(out: &mut Vec<u8>, value: u64) {
    out.extend_from_slice(&value.to_le_bytes());
}

/// Write length-encoded integer
pub fn write_int_lenenc(out: &mut Vec<u8>, value: u64) {
    if value < 251 {
        out.push(value as u8);
    } else if value < (1 << 16) {
        out.push(0xfc);
        write_int_2(out, value as u16);
    } else if value < (1 << 24) {
        out.push(0xfd);
        write_int_3(out, value as u32);
    } else {
        out.push(0xfe);
        write_int_8(out, value);
    }
}

/// Write fixed-length bytes
#[inline]
pub fn write_bytes_fix(out: &mut Vec<u8>, data: &[u8]) {
    out.extend_from_slice(data);
}

/// Write null-terminated string
#[inline]
pub fn write_string_null(out: &mut Vec<u8>, bytes: &[u8]) {
    out.extend_from_slice(bytes);
    out.push(0);
}

/// Write length-encoded string
#[inline]
pub fn write_string_lenenc(out: &mut Vec<u8>, s: &str) {
    write_int_lenenc(out, s.len() as u64);
    out.extend_from_slice(s.as_bytes());
}

/// Write length-encoded bytes
#[inline]
pub fn write_bytes_lenenc(out: &mut Vec<u8>, data: &[u8]) {
    write_int_lenenc(out, data.len() as u64);
    out.extend_from_slice(data);
}