filip/Stalwart
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

IMAP: Argument buffer parser checked push

Browse Source
This commit is contained in:
mdecimus
2025-09-30 12:34:59 +02:00
parent 2f6cfbb6e6
commit 4a96357b2b
1 changed files with 125 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

182
crates/imap-proto/src/receiver.rs
View File

@@ -8,8 +8,6 @@ use super::{ResponseCode, ResponseType};
use compact_str::{CompactString, format_compact}; use compact_str::{CompactString, format_compact};
use std::fmt::Display; use std::fmt::Display;
const QUOTED_ARG_MAX_LEN: usize = 4096;
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub enum Error { pub enum Error {
NeedsMoreData, NeedsMoreData,
@@ -65,7 +63,7 @@ pub enum State {
} }
pub struct Receiver<T: CommandParser> { pub struct Receiver<T: CommandParser> {
buf: Vec<u8>, buf: ArgumentBuffer,
pub request: Request<T>, pub request: Request<T>,
pub state: State, pub state: State,
pub max_request_size: usize, pub max_request_size: usize,
@@ -73,6 +71,12 @@ pub struct Receiver<T: CommandParser> {
pub start_state: State, pub start_state: State,
} }
const ARG_MAX_LEN: usize = 4096;
struct ArgumentBuffer {
buf: Vec<u8>,
}
impl<T: CommandParser> Receiver<T> { impl<T: CommandParser> Receiver<T> {
pub fn new() -> Self { pub fn new() -> Self {
Receiver { Receiver {
@@ -104,7 +108,7 @@ impl<T: CommandParser> Receiver<T> {
}, },
message, message,
); );
self.buf = Vec::with_capacity(10); self.buf = ArgumentBuffer::default();
self.state = self.start_state; self.state = self.start_state;
self.current_request_size = 0; self.current_request_size = 0;
err err
@@ -119,8 +123,7 @@ impl<T: CommandParser> Receiver<T> {
self.max_request_size self.max_request_size
))); )));
} }
self.request.tokens.push(Token::Argument(self.buf.clone())); self.request.tokens.push(Token::Argument(self.buf.take()));
self.buf.clear();
} else if in_quote { } else if in_quote {
self.request.tokens.push(Token::Nil); self.request.tokens.push(Token::Nil);
} }
@@ -145,18 +148,18 @@ impl<T: CommandParser> Receiver<T> {
match self.state { match self.state {
State::Start => { State::Start => {
if !ch.is_ascii_whitespace() { if !ch.is_ascii_whitespace() {
self.buf.push(ch); // SAFETY: This called just once
self.buf.push_unchecked(ch);
self.state = State::Tag; self.state = State::Tag;
} }
} }
State::Tag => match ch { State::Tag => match ch {
b' ' => { b' ' => {
if !self.buf.is_empty() { if !self.buf.is_empty() {
self.request.tag = String::from_utf8(std::mem::replace( self.request.tag =
&mut self.buf, String::from_utf8(self.buf.take()).map_err(|_| {
Vec::with_capacity(10), self.error_reset("Tag is not a valid UTF-8 string.")
)) })?;
.map_err(|_| self.error_reset("Tag is not a valid UTF-8 string."))?;
self.state = State::Command { is_uid: false }; self.state = State::Command { is_uid: false };
} }
} }
@@ -164,35 +167,32 @@ impl<T: CommandParser> Receiver<T> {
b'\n' => { b'\n' => {
return Err(self.error_reset(format_compact!( return Err(self.error_reset(format_compact!(
"Missing command after tag {:?}, found CRLF instead.", "Missing command after tag {:?}, found CRLF instead.",
std::str::from_utf8(&self.buf).unwrap_or_default() self.buf.as_str()
))); )));
} }
_ => { _ => {
if self.buf.len() < 128 { self.buf.push_checked(ch, 128).map_err(|_| {
self.buf.push(ch); self.error_reset("Tag exceeds maximum length of 128 characters.")
} else { })?;
return Err(self.error_reset("Tag too long."));
}
} }
}, },
State::Command { is_uid } => { State::Command { is_uid } => {
if ch.is_ascii_alphanumeric() { if ch.is_ascii_alphanumeric() {
if self.buf.len() < 15 { self.buf
self.buf.push(ch.to_ascii_uppercase()); .push_checked(ch.to_ascii_uppercase(), 15)
} else { .map_err(|_| {
return Err(self.error_reset("Command too long")); self.error_reset("Command exceeds maximum length of 15 characters.")
} })?;
} else if ch.is_ascii_whitespace() { } else if ch.is_ascii_whitespace() {
if !self.buf.is_empty() { if !self.buf.is_empty() {
if !self.buf.eq_ignore_ascii_case(b"UID") { if !self.buf.as_ref().eq_ignore_ascii_case(b"UID") {
self.request.command = self.request.command = T::parse(self.buf.as_ref(), is_uid)
T::parse(&self.buf, is_uid).ok_or_else(|| { .ok_or_else(|| {
let command = let err = format_compact!(
String::from_utf8_lossy(&self.buf).into_owned();
self.error_reset(format_compact!(
"Unrecognized command '{}'.", "Unrecognized command '{}'.",
command String::from_utf8_lossy(self.buf.as_ref())
)) );
self.error_reset(err)
})?; })?;
self.buf.clear(); self.buf.clear();
if ch != b'\n' { if ch != b'\n' {
@@ -268,7 +268,9 @@ impl<T: CommandParser> Receiver<T> {
self.state = State::Argument { last_ch: ch }; self.state = State::Argument { last_ch: ch };
} }
_ => { _ => {
self.buf.push(ch); self.buf.push_checked(ch, ARG_MAX_LEN).map_err(|_| {
self.error_reset("Argument exceeds maximum length of 4096 bytes.")
})?;
self.state = State::Argument { last_ch: ch }; self.state = State::Argument { last_ch: ch };
} }
}, },
@@ -277,20 +279,18 @@ impl<T: CommandParser> Receiver<T> {
if !escaped { if !escaped {
self.push_argument(true)?; self.push_argument(true)?;
self.state = State::Argument { last_ch: b' ' }; self.state = State::Argument { last_ch: b' ' };
} else if self.buf.len() < QUOTED_ARG_MAX_LEN {
self.buf.push(ch);
self.state = State::ArgumentQuoted { escaped: false };
} else { } else {
return Err(self.error_reset("Quoted argument too long.")); self.buf
.push_checked(ch, ARG_MAX_LEN)
.map_err(|_| self.error_reset("Quoted argument too long."))?;
self.state = State::ArgumentQuoted { escaped: false };
} }
} }
b'\\' => { b'\\' => {
if escaped { if escaped {
if self.buf.len() < QUOTED_ARG_MAX_LEN { self.buf
self.buf.push(ch); .push_checked(ch, ARG_MAX_LEN)
} else { .map_err(|_| self.error_reset("Quoted argument too long."))?;
return Err(self.error_reset("Quoted argument too long."));
}
} }
self.state = State::ArgumentQuoted { escaped: !escaped }; self.state = State::ArgumentQuoted { escaped: !escaped };
} }
@@ -298,25 +298,21 @@ impl<T: CommandParser> Receiver<T> {
return Err(self.error_reset("Unterminated quoted argument.")); return Err(self.error_reset("Unterminated quoted argument."));
} }
_ => { _ => {
if self.buf.len() < QUOTED_ARG_MAX_LEN { if escaped {
if escaped { // SAFETY: We check the size below
self.buf.push(b'\\'); self.buf.push_unchecked(b'\\');
}
self.buf.push(ch);
self.state = State::ArgumentQuoted { escaped: false };
} else {
return Err(self.error_reset("Quoted argument too long."));
} }
self.buf
.push_checked(ch, ARG_MAX_LEN)
.map_err(|_| self.error_reset("Quoted argument too long."))?;
self.state = State::ArgumentQuoted { escaped: false };
} }
}, },
State::Literal { non_sync } => { State::Literal { non_sync } => {
match ch { match ch {
b'}' => { b'}' => {
if !self.buf.is_empty() { if !self.buf.is_empty() {
let size = std::str::from_utf8(&self.buf) let size = self.buf.as_str().parse::<u32>().map_err(|_| {
.unwrap()
.parse::<u32>()
.map_err(|_| {
self.error_reset("Literal size is not a valid number.") self.error_reset("Literal size is not a valid number.")
})?; })?;
if self.current_request_size + size as usize > self.max_request_size if self.current_request_size + size as usize > self.max_request_size
@@ -327,7 +323,8 @@ impl<T: CommandParser> Receiver<T> {
))); )));
} }
self.state = State::LiteralSeek { size, non_sync }; self.state = State::LiteralSeek { size, non_sync };
self.buf = Vec::with_capacity(size as usize); self.buf.resize_buffer(size as usize);
self.buf.clear();
} else { } else {
return Err(self.error_reset("Invalid empty literal.")); return Err(self.error_reset("Invalid empty literal."));
} }
@@ -341,7 +338,9 @@ impl<T: CommandParser> Receiver<T> {
} }
_ if ch.is_ascii_digit() => { _ if ch.is_ascii_digit() => {
if !non_sync { if !non_sync {
self.buf.push(ch); self.buf.push_checked(ch, 15).map_err(|_| {
self.error_reset("Literal size exceeds maximum of 15 digits.")
})?;
} else { } else {
// Digit found after non-sync '+' flag // Digit found after non-sync '+' flag
return Err(self.error_reset("Invalid literal.")); return Err(self.error_reset("Invalid literal."));
@@ -373,7 +372,9 @@ impl<T: CommandParser> Receiver<T> {
} }
} }
State::LiteralData { remaining } => { State::LiteralData { remaining } => {
self.buf.push(ch); // SAFETY: We checked the size before entering this state
self.buf.push_unchecked(ch);
if remaining > 1 { if remaining > 1 {
self.state = State::LiteralData { self.state = State::LiteralData {
remaining: remaining - 1, remaining: remaining - 1,
@@ -390,6 +391,61 @@ impl<T: CommandParser> Receiver<T> {
} }
} }
impl ArgumentBuffer {
pub fn new() -> Self {
ArgumentBuffer {
buf: Vec::with_capacity(10),
}
}
pub fn resize_buffer(&mut self, size: usize) {
if self.buf.capacity() < size {
self.buf.reserve(size - self.buf.capacity());
}
}
#[inline(always)]
pub fn push_checked(&mut self, byte: u8, limit: usize) -> Result<(), ()> {
if self.buf.len() < limit {
self.buf.push(byte);
Ok(())
} else {
Err(())
}
}
#[inline(always)]
pub fn push_unchecked(&mut self, byte: u8) {
self.buf.push(byte);
}
pub fn take(&mut self) -> Vec<u8> {
let buf = self.buf.clone();
self.buf.clear();
buf
}
#[inline(always)]
pub fn len(&self) -> usize {
self.buf.len()
}
#[inline(always)]
pub fn is_empty(&self) -> bool {
self.buf.is_empty()
}
#[inline(always)]
pub fn clear(&mut self) {
self.buf.clear();
}
#[inline(always)]
pub fn as_str(&self) -> &str {
std::str::from_utf8(&self.buf).unwrap_or_default()
}
}
impl Token { impl Token {
pub fn unwrap_string(self) -> crate::parser::Result<String> { pub fn unwrap_string(self) -> crate::parser::Result<String> {
match self { match self {
@@ -450,6 +506,18 @@ impl Token {
} }
} }
impl AsRef<[u8]> for ArgumentBuffer {
fn as_ref(&self) -> &[u8] {
&self.buf
}
}
impl Default for ArgumentBuffer {
fn default() -> Self {
Self::new()
}
}
impl Display for Token { impl Display for Token {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
f.write_str(&String::from_utf8_lossy(self.as_bytes())) f.write_str(&String::from_utf8_lossy(self.as_bytes()))
@@ -487,7 +555,7 @@ impl Error {
impl<T: CommandParser> Default for Receiver<T> { impl<T: CommandParser> Default for Receiver<T> {
fn default() -> Self { fn default() -> Self {
Self { Self {
buf: Vec::with_capacity(10), buf: Default::default(),
request: Default::default(), request: Default::default(),
state: State::Start, state: State::Start,
start_state: State::Start, start_state: State::Start,