use crate::{
Error,
Result,
};
use bytes::{
Buf,
BufMut,
BytesMut,
};
pub fn read_varint(buffer: &mut &[u8]) -> Result<u64> {
let mut result: u64 = 0;
let mut shift = 0;
loop {
if buffer.is_empty() {
return Err(Error::Protocol(
"Unexpected end of buffer reading varint".to_string(),
));
}
let byte = buffer[0];
buffer.advance(1);
result |= ((byte & 0x7F) as u64) << shift;
if byte & 0x80 == 0 {
break;
}
shift += 7;
if shift >= 64 {
return Err(Error::Protocol("Varint overflow".to_string()));
}
}
Ok(result)
}
pub fn write_varint(buffer: &mut BytesMut, mut value: u64) {
loop {
let mut byte = (value & 0x7F) as u8;
value >>= 7;
if value != 0 {
byte |= 0x80;
}
buffer.put_u8(byte);
if value == 0 {
break;
}
}
}
pub fn read_string(buffer: &mut &[u8]) -> Result<String> {
let len = read_varint(buffer)? as usize;
if buffer.len() < len {
return Err(Error::Protocol(format!(
"Not enough data for string: need {}, have {}",
len,
buffer.len()
)));
}
let string_data = &buffer[..len];
let s = String::from_utf8(string_data.to_vec()).map_err(|e| {
Error::Protocol(format!("Invalid UTF-8 in string: {}", e))
})?;
buffer.advance(len);
Ok(s)
}
pub fn write_string(buffer: &mut BytesMut, s: &str) {
write_varint(buffer, s.len() as u64);
buffer.put_slice(s.as_bytes());
}
pub fn write_varint_to_vec(buf: &mut Vec<u8>, mut value: u64) {
loop {
let mut byte = (value & 0x7F) as u8;
value >>= 7;
if value != 0 {
byte |= 0x80;
}
buf.push(byte);
if value == 0 {
break;
}
}
}
#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
mod tests {
use super::*;
#[test]
fn test_varint_roundtrip() {
let test_cases =
vec![0u64, 1, 127, 128, 255, 256, 65535, 65536, u64::MAX];
for value in test_cases {
let mut buf = BytesMut::new();
write_varint(&mut buf, value);
let mut slice = &buf[..];
let decoded = read_varint(&mut slice).unwrap();
assert_eq!(
value, decoded,
"Varint roundtrip failed for {}",
value
);
assert!(slice.is_empty(), "Buffer should be fully consumed");
}
}
#[test]
fn test_string_roundtrip() {
let test_strings =
vec!["", "hello", "мир", "🦀", "test\nwith\nnewlines"];
for s in test_strings {
let mut buf = BytesMut::new();
write_string(&mut buf, s);
let mut slice = &buf[..];
let decoded = read_string(&mut slice).unwrap();
assert_eq!(s, decoded, "String roundtrip failed for '{}'", s);
assert!(slice.is_empty(), "Buffer should be fully consumed");
}
}
#[test]
fn test_varint_overflow() {
let mut buf = BytesMut::new();
for _ in 0..10 {
buf.put_u8(0xFF); }
let mut slice = &buf[..];
let result = read_varint(&mut slice);
assert!(result.is_err());
}
#[test]
fn test_string_truncated() {
let mut buf = BytesMut::new();
write_varint(&mut buf, 100); buf.put_slice(b"only10");
let mut slice = &buf[..];
let result = read_string(&mut slice);
assert!(result.is_err());
}
#[test]
fn test_varint_to_vec() {
let mut buf = Vec::new();
write_varint_to_vec(&mut buf, 300);
let mut slice = &buf[..];
let decoded = read_varint(&mut slice).unwrap();
assert_eq!(decoded, 300);
}
}