use crate::error::{Result, SegbError};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WireType {
Varint = 0,
Bit64 = 1,
LengthDelimited = 2,
Bit32 = 5,
}
impl WireType {
fn from_u64(v: u64) -> Option<Self> {
match v {
0 => Some(Self::Varint),
1 => Some(Self::Bit64),
2 => Some(Self::LengthDelimited),
5 => Some(Self::Bit32),
_ => None,
}
}
}
#[derive(Debug, Clone)]
pub struct Field<'a> {
pub field_number: u32,
pub wire_type: WireType,
pub raw: &'a [u8],
}
fn read_varint(data: &[u8], pos: &mut usize) -> Result<u64> {
let start = *pos;
let mut result: u64 = 0;
let mut shift = 0u32;
loop {
if *pos >= data.len() {
return Err(SegbError::MalformedVarint { offset: start });
}
let byte = data[*pos];
*pos += 1;
let low_7 = u64::from(byte & 0x7F);
if shift < 64 {
result |= low_7 << shift;
}
shift += 7;
if byte & 0x80 == 0 {
return Ok(result);
}
if shift >= 70 {
return Err(SegbError::MalformedVarint { offset: start });
}
}
}
pub fn iter_fields(buf: &[u8]) -> impl Iterator<Item = Result<Field<'_>>> {
FieldIter { buf, pos: 0 }
}
struct FieldIter<'a> {
buf: &'a [u8],
pos: usize,
}
impl<'a> Iterator for FieldIter<'a> {
type Item = Result<Field<'a>>;
fn next(&mut self) -> Option<Self::Item> {
if self.pos >= self.buf.len() {
return None;
}
let tag = match read_varint(self.buf, &mut self.pos) {
Ok(v) => v,
Err(e) => return Some(Err(e)),
};
let wire_type_raw = tag & 0x07;
let field_number = (tag >> 3) as u32;
let wire_type = if let Some(wt) = WireType::from_u64(wire_type_raw) {
wt
} else {
self.pos = self.buf.len();
return None;
};
let raw: &'a [u8] = match wire_type {
WireType::Varint => {
let start = self.pos;
match read_varint(self.buf, &mut self.pos) {
Ok(_) => &self.buf[start..self.pos],
Err(e) => return Some(Err(e)),
}
}
WireType::Bit64 => {
let start = self.pos;
if self.pos + 8 > self.buf.len() {
return Some(Err(SegbError::ProtobufOverflow {
offset: start,
length: 8,
remaining: self.buf.len() - self.pos,
}));
}
self.pos += 8;
&self.buf[start..self.pos]
}
WireType::LengthDelimited => {
let len_pos = self.pos;
let length = match read_varint(self.buf, &mut self.pos) {
Ok(v) => v as usize,
Err(e) => return Some(Err(e)),
};
let start = self.pos;
let remaining = self.buf.len() - self.pos;
if length > remaining {
return Some(Err(SegbError::ProtobufOverflow {
offset: len_pos,
length,
remaining,
}));
}
self.pos += length;
&self.buf[start..self.pos]
}
WireType::Bit32 => {
let start = self.pos;
if self.pos + 4 > self.buf.len() {
return Some(Err(SegbError::ProtobufOverflow {
offset: start,
length: 4,
remaining: self.buf.len() - self.pos,
}));
}
self.pos += 4;
&self.buf[start..self.pos]
}
};
Some(Ok(Field {
field_number,
wire_type,
raw,
}))
}
}
#[inline]
pub fn as_str(raw: &[u8]) -> Option<&str> {
std::str::from_utf8(raw).ok()
}