pub const MAX_RESP_DEPTH: usize = 32;
pub const MAX_RESP_ELEMENTS: usize = 1 << 20;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RespLimits {
pub max_depth: usize,
pub max_elements: usize,
pub max_bulk_len: usize,
}
impl RespLimits {
pub const fn with_max_bulk_len(max_bulk_len: usize) -> Self {
Self {
max_depth: MAX_RESP_DEPTH,
max_elements: MAX_RESP_ELEMENTS,
max_bulk_len,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RespLength {
Aggregate,
Bulk,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RespScan {
Complete { consumed: usize },
Incomplete,
TooDeep { limit: usize },
TooLong(RespLength),
Malformed,
}
pub fn scan_resp2_frame(data: &[u8], limits: &RespLimits) -> RespScan {
let max_depth = limits.max_depth.min(MAX_RESP_DEPTH);
let mut remaining = [0usize; MAX_RESP_DEPTH];
let mut depth = 0usize;
let mut values = 0usize;
let mut i = 0usize;
loop {
let Some(&kind) = data.get(i) else {
return RespScan::Incomplete;
};
i += 1;
values += 1;
if values > limits.max_elements {
return RespScan::TooLong(RespLength::Aggregate);
}
match kind {
b'+' | b'-' => match find_crlf(data, i) {
Some(end) => i = end + 2,
None => return RespScan::Incomplete,
},
b':' => {
let Some(end) = find_crlf(data, i) else {
return RespScan::Incomplete;
};
if parse_i64(&data[i..end]).is_none() {
return RespScan::Malformed;
}
i = end + 2;
}
b'$' => {
let Some(end) = find_crlf(data, i) else {
return RespScan::Incomplete;
};
let Some(len) = parse_isize(&data[i..end]) else {
return RespScan::Malformed;
};
i = end + 2;
if len != -1 {
let Ok(len) = usize::try_from(len) else {
return RespScan::Malformed;
};
if len > limits.max_bulk_len {
return RespScan::TooLong(RespLength::Bulk);
}
let Some(next) = i.checked_add(len).and_then(|v| v.checked_add(2)) else {
return RespScan::TooLong(RespLength::Bulk);
};
if next > data.len() {
return RespScan::Incomplete;
}
i = next;
}
}
b'*' => {
let Some(end) = find_crlf(data, i) else {
return RespScan::Incomplete;
};
let Some(len) = parse_isize(&data[i..end]) else {
return RespScan::Malformed;
};
i = end + 2;
if len != -1 {
let Ok(len) = usize::try_from(len) else {
return RespScan::Malformed;
};
if len > limits.max_elements.saturating_sub(values) {
return RespScan::TooLong(RespLength::Aggregate);
}
if depth + 1 > max_depth {
return RespScan::TooDeep { limit: max_depth };
}
if len > 0 {
remaining[depth] = len;
depth += 1;
continue;
}
}
}
_ => return RespScan::Malformed,
}
loop {
if depth == 0 {
return RespScan::Complete { consumed: i };
}
remaining[depth - 1] -= 1;
if remaining[depth - 1] > 0 {
break;
}
depth -= 1;
}
}
}
fn find_crlf(data: &[u8], from: usize) -> Option<usize> {
data.get(from..)?
.windows(2)
.position(|w| w == b"\r\n")
.map(|p| from + p)
}
fn parse_isize(line: &[u8]) -> Option<isize> {
std::str::from_utf8(line).ok()?.parse::<isize>().ok()
}
fn parse_i64(line: &[u8]) -> Option<i64> {
std::str::from_utf8(line).ok()?.parse::<i64>().ok()
}