#[derive(Debug, Clone, Copy)]
pub struct GrpcLimits {
pub max_depth: u32,
pub max_fields: u32,
pub max_leaves: usize,
}
impl Default for GrpcLimits {
fn default() -> Self {
Self { max_depth: 16, max_fields: 4096, max_leaves: 1024 }
}
}
#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct GrpcExtract {
pub messages: u32,
pub total_payload_bytes: u64,
pub max_message_len: u32,
pub fields: u32,
pub max_depth: u32,
pub depth_exceeded: bool,
pub fields_exceeded: bool,
pub compressed: bool,
pub malformed: bool,
pub leaves: Vec<String>,
}
fn read_varint(b: &[u8], i: &mut usize) -> Option<u64> {
let mut result: u64 = 0;
for k in 0..10 {
let byte = *b.get(*i)?;
*i += 1;
result |= ((byte & 0x7f) as u64) << (7 * k);
if byte & 0x80 == 0 {
return Some(result);
}
}
None }
pub fn grpc_extract(body: &[u8], limits: GrpcLimits) -> GrpcExtract {
let mut out = GrpcExtract::default();
let mut pos = 0;
while pos + 5 <= body.len() {
let flag = body[pos];
let len = u32::from_be_bytes([body[pos + 1], body[pos + 2], body[pos + 3], body[pos + 4]]) as usize;
let msg_start = pos + 5;
let Some(msg_end) = msg_start.checked_add(len) else {
out.malformed = true;
return out;
};
if msg_end > body.len() {
out.malformed = true; return out;
}
out.messages = out.messages.saturating_add(1);
out.total_payload_bytes = out.total_payload_bytes.saturating_add(len as u64);
out.max_message_len = out.max_message_len.max(len as u32);
let msg = &body[msg_start..msg_end];
match flag {
0 => parse_message(msg, 1, &limits, &mut out),
1 => out.compressed = true, _ => out.malformed = true, }
pos = msg_end;
}
if pos != body.len() {
out.malformed = true;
}
out
}
fn parse_message(buf: &[u8], depth: u32, limits: &GrpcLimits, out: &mut GrpcExtract) {
out.max_depth = out.max_depth.max(depth);
let mut i = 0;
while i < buf.len() {
let Some(tag) = read_varint(buf, &mut i) else {
out.malformed = true;
return;
};
out.fields = out.fields.saturating_add(1);
if out.fields > limits.max_fields {
out.fields_exceeded = true;
return;
}
match (tag & 0x7) as u8 {
0 => {
if read_varint(buf, &mut i).is_none() {
out.malformed = true;
return;
}
}
1 => {
if i + 8 > buf.len() {
out.malformed = true;
return;
}
i += 8;
}
5 => {
if i + 4 > buf.len() {
out.malformed = true;
return;
}
i += 4;
}
2 => {
let Some(len) = read_varint(buf, &mut i) else {
out.malformed = true;
return;
};
let end = match i.checked_add(len as usize) {
Some(e) if e <= buf.len() => e,
_ => {
out.malformed = true;
return;
}
};
let field = &buf[i..end];
i = end;
match std::str::from_utf8(field) {
Ok(s) if !s.is_empty() => {
if out.leaves.len() < limits.max_leaves {
out.leaves.push(s.to_string());
}
}
Ok(_) => {} Err(_) => {
if depth >= limits.max_depth {
out.depth_exceeded = true;
} else {
parse_message(field, depth + 1, limits, out);
}
}
}
}
_ => {
out.malformed = true;
return;
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn varint(mut v: u64, out: &mut Vec<u8>) {
loop {
let mut byte = (v & 0x7f) as u8;
v >>= 7;
if v != 0 {
byte |= 0x80;
}
out.push(byte);
if v == 0 {
break;
}
}
}
fn len_field(field: u64, data: &[u8]) -> Vec<u8> {
let mut out = Vec::new();
varint((field << 3) | 2, &mut out);
varint(data.len() as u64, &mut out);
out.extend_from_slice(data);
out
}
fn varint_field(field: u64, value: u64) -> Vec<u8> {
let mut out = Vec::new();
varint(field << 3, &mut out); varint(value, &mut out);
out
}
fn frame(msg: &[u8]) -> Vec<u8> {
let mut out = vec![0u8];
out.extend_from_slice(&(msg.len() as u32).to_be_bytes());
out.extend_from_slice(msg);
out
}
#[test]
fn extracts_string_leaf_from_unary_frame() {
let sqli = "1 UNION SELECT a,b FROM users--";
let body = frame(&len_field(1, sqli.as_bytes()));
let ex = grpc_extract(&body, GrpcLimits::default());
assert_eq!(ex.messages, 1);
assert!(!ex.compressed && !ex.malformed);
assert!(ex.leaves.iter().any(|l| l == sqli), "SQLi leaf must be extracted: {:?}", ex.leaves);
}
#[test]
fn recurses_into_non_utf8_submessage() {
let mut sub = varint_field(2, 300);
sub.extend_from_slice(&len_field(1, b"inner-secret"));
let body = frame(&len_field(3, &sub));
let ex = grpc_extract(&body, GrpcLimits::default());
assert!(ex.max_depth >= 2, "must recurse into the sub-message (depth {})", ex.max_depth);
assert!(ex.leaves.iter().any(|l| l == "inner-secret"), "nested leaf: {:?}", ex.leaves);
assert!(!ex.depth_exceeded);
}
fn nested_messages(depth: u32) -> Vec<u8> {
let mut inner = varint_field(2, 300); inner.extend_from_slice(&len_field(15, b"benign-leaf"));
for _ in 0..depth {
inner = {
let mut wrap = varint_field(2, 300);
wrap.extend_from_slice(&len_field(1, &inner));
wrap
};
}
frame(&inner)
}
#[test]
fn benign_deep_but_within_cap_is_not_a_false_reject() {
let limits = GrpcLimits { max_depth: 16, ..GrpcLimits::default() };
let ex = grpc_extract(&nested_messages(8), limits);
assert!(!ex.depth_exceeded, "legitimate nesting must not trip the depth cap");
assert!(!ex.malformed);
assert!(ex.leaves.iter().any(|l| l == "benign-leaf"));
}
#[test]
fn depth_bomb_beyond_cap_is_flagged() {
let limits = GrpcLimits { max_depth: 8, ..GrpcLimits::default() };
let ex = grpc_extract(&nested_messages(40), limits);
assert!(ex.depth_exceeded, "a depth-bomb past the cap must be flagged");
assert!(ex.max_depth <= 8, "recursion must stop at the cap (got {})", ex.max_depth);
}
#[test]
fn compressed_frame_is_flagged_not_parsed() {
let mut body = vec![1u8];
let payload = len_field(1, b"would-be-secret");
body.extend_from_slice(&(payload.len() as u32).to_be_bytes());
body.extend_from_slice(&payload);
let ex = grpc_extract(&body, GrpcLimits::default());
assert!(ex.compressed);
assert!(ex.leaves.is_empty(), "a compressed frame must not be parsed here");
}
#[test]
fn field_cap_is_flagged() {
let mut msg = Vec::new();
for n in 1..=10u64 {
msg.extend_from_slice(&varint_field(n, 1));
}
let ex = grpc_extract(&frame(&msg), GrpcLimits { max_fields: 5, ..GrpcLimits::default() });
assert!(ex.fields_exceeded);
}
#[test]
fn multi_frame_body_counts_all() {
let mut body = frame(&len_field(1, b"first"));
body.extend_from_slice(&frame(&len_field(1, b"second")));
let ex = grpc_extract(&body, GrpcLimits::default());
assert_eq!(ex.messages, 2);
assert!(ex.leaves.iter().any(|l| l == "first"));
assert!(ex.leaves.iter().any(|l| l == "second"));
}
#[test]
fn truncated_and_malformed_inputs_do_not_panic() {
let limits = GrpcLimits::default();
let _ = grpc_extract(&[], limits);
let _ = grpc_extract(&[0], limits); let _ = grpc_extract(&[0, 0, 0, 0, 100], limits); let _ = grpc_extract(&[0, 0, 0, 0, 3, 0xff, 0xff, 0xff], limits); let _ = grpc_extract(&[0, 0, 0, 0, 2, 0x08], limits); let _ = grpc_extract(&frame(&[0x0a, 0x7f]), limits);
}
#[test]
fn empty_string_field_yields_no_leaf() {
let ex = grpc_extract(&frame(&len_field(1, b"")), GrpcLimits::default());
assert!(ex.leaves.is_empty());
assert_eq!(ex.messages, 1);
assert!(!ex.malformed);
}
}