use crate::syntax::{is_field_name_byte, is_field_value_byte};
use crate::{
Boundary, MultipartError, MultipartErrorKind, MultipartEvent, MultipartLimits, Parsed, PartHead,
};
use memchr::memmem::Finder;
#[derive(Clone, Debug)]
pub struct MultipartParser {
delimiter: Vec<u8>,
finder: Finder<'static>,
limits: MultipartLimits,
state: State,
offset: u64,
preamble_size: usize,
part_count: usize,
failed: Option<MultipartError>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum State {
Start,
Preamble,
Delimiter {
padding: usize,
},
Closing {
padding: usize,
},
Headers {
padding: usize,
searched: usize,
},
BodyStart,
Content,
End,
}
enum Found {
Delimiter(usize),
Bare(usize),
Absent,
}
impl MultipartParser {
pub fn new(boundary: &Boundary) -> Self {
Self::with_limits(boundary, MultipartLimits::default())
}
pub fn with_limits(boundary: &Boundary, limits: MultipartLimits) -> Self {
let delimiter = [b"\r\n--", boundary.as_str().as_bytes()].concat();
let finder = Finder::new(&delimiter[2..]).into_owned();
Self {
delimiter,
finder,
limits,
state: State::Start,
offset: 0,
preamble_size: 0,
part_count: 0,
failed: None,
}
}
pub fn limits(&self) -> MultipartLimits {
self.limits
}
pub fn minimum_buffer_size(&self) -> usize {
self.limits
.header_size_limit
.saturating_add(2)
.max(self.delimiter.len())
}
pub fn is_finished(&self) -> bool {
self.state == State::End
}
pub fn parse<'a>(&mut self, input: &'a [u8]) -> Result<Parsed<'a>, MultipartError> {
if let Some(error) = self.failed {
return Err(error);
}
let mut consumed = 0;
match self.step(input, &mut consumed) {
Ok(event) => {
self.offset += consumed as u64;
Ok(Parsed { consumed, event })
}
Err(kind) => {
let error = MultipartError::new(kind, self.offset + consumed as u64);
self.failed = Some(error);
Err(error)
}
}
}
pub fn finish(&self) -> Result<(), MultipartError> {
if let Some(error) = self.failed {
return Err(error);
}
match self.state {
State::End | State::Closing { .. } => Ok(()),
_ => Err(MultipartError::new(
MultipartErrorKind::UnexpectedEnd,
self.offset,
)),
}
}
fn step<'a>(
&mut self,
input: &'a [u8],
consumed: &mut usize,
) -> Result<MultipartEvent<'a>, MultipartErrorKind> {
loop {
let rest = &input[*consumed..];
match self.state {
State::Start => {
let dash_boundary = &self.delimiter[2..];
if rest.starts_with(dash_boundary) {
*consumed += dash_boundary.len();
self.state = State::Delimiter { padding: 0 };
} else if rest.len() < dash_boundary.len() && dash_boundary.starts_with(rest) {
return Ok(MultipartEvent::Incomplete);
} else {
self.state = State::Preamble;
}
}
State::Preamble => {
let (skipped, found) = match self.find(rest) {
Found::Delimiter(index) => (index, true),
Found::Bare(index) => {
*consumed += index;
return Err(MultipartErrorKind::InvalidDelimiter);
}
Found::Absent => (self.safe_length(rest), false),
};
self.preamble_size = self.preamble_size.saturating_add(skipped);
if self.preamble_size > self.limits.preamble_size_limit {
return Err(MultipartErrorKind::PreambleSizeLimit);
}
*consumed += skipped;
if !found {
return Ok(MultipartEvent::Incomplete);
}
*consumed += self.delimiter.len();
self.state = State::Delimiter { padding: 0 };
}
State::Delimiter { padding } => {
if padding == 0 && rest.starts_with(b"--") {
*consumed += 2;
self.state = State::Closing { padding: 0 };
continue;
}
let padding = self.pad(rest, padding, consumed)?;
self.state = State::Delimiter { padding };
let rest = &input[*consumed..];
if rest.starts_with(b"\r\n") {
self.part_count += 1;
if self.part_count > self.limits.part_count_limit {
return Err(MultipartErrorKind::PartCountLimit);
}
*consumed += 2;
self.state = State::Headers {
padding,
searched: 0,
};
} else if rest.is_empty() || rest == b"\r" || (padding == 0 && rest == b"-") {
return Ok(MultipartEvent::Incomplete);
} else {
return Err(MultipartErrorKind::InvalidDelimiter);
}
}
State::Closing { padding } => {
let padding = self.pad(rest, padding, consumed)?;
self.state = State::Closing { padding };
let rest = &input[*consumed..];
if rest.starts_with(b"\r\n") {
*consumed += 2;
self.state = State::End;
return Ok(MultipartEvent::End);
} else if rest.is_empty() || rest == b"\r" {
return Ok(MultipartEvent::Incomplete);
} else {
return Err(MultipartErrorKind::InvalidDelimiter);
}
}
State::Headers { padding, searched } => {
let limit = self.limits.header_size_limit - padding;
let length = if rest.starts_with(b"\r\n") {
Some(0)
} else {
let start = searched.saturating_sub(3).min(rest.len());
memchr::memmem::find(&rest[start..], b"\r\n\r\n")
.map(|index| start + index + 2)
};
let Some(length) = length else {
if rest.len() >= limit.saturating_add(2) {
return Err(MultipartErrorKind::HeaderSizeLimit);
}
self.state = State::Headers {
padding,
searched: rest.len(),
};
return Ok(MultipartEvent::Incomplete);
};
if length > limit {
return Err(MultipartErrorKind::HeaderSizeLimit);
}
let section = &rest[..length];
self.check_section(section)?;
*consumed += length;
self.state = State::BodyStart;
return Ok(MultipartEvent::Part(PartHead::new(section)));
}
State::BodyStart => {
if rest.starts_with(&self.delimiter) {
*consumed += self.delimiter.len();
self.state = State::Delimiter { padding: 0 };
return Ok(MultipartEvent::PartEnd);
}
if rest.len() < self.delimiter.len() && self.delimiter.starts_with(rest) {
return Ok(MultipartEvent::Incomplete);
}
*consumed += 2;
self.state = State::Content;
}
State::Content => {
return match self.find(rest) {
Found::Delimiter(0) => {
*consumed += self.delimiter.len();
self.state = State::Delimiter { padding: 0 };
Ok(MultipartEvent::PartEnd)
}
Found::Delimiter(index) => {
*consumed += index;
Ok(MultipartEvent::Content(&rest[..index]))
}
Found::Bare(index) => {
*consumed += index;
Err(MultipartErrorKind::InvalidDelimiter)
}
Found::Absent => match self.safe_length(rest) {
0 => Ok(MultipartEvent::Incomplete),
length => {
*consumed += length;
Ok(MultipartEvent::Content(&rest[..length]))
}
},
};
}
State::End => {
*consumed += rest.len();
return Ok(MultipartEvent::End);
}
}
}
}
fn pad(
&self,
rest: &[u8],
padding: usize,
consumed: &mut usize,
) -> Result<usize, MultipartErrorKind> {
let spaces = rest
.iter()
.take_while(|&&byte| byte == b' ' || byte == b'\t')
.count();
let padding = padding + spaces;
if padding > self.limits.header_size_limit {
return Err(MultipartErrorKind::HeaderSizeLimit);
}
*consumed += spaces;
Ok(padding)
}
fn find(&self, input: &[u8]) -> Found {
let mut from = 0;
while let Some(index) = self.finder.find(&input[from..]).map(|index| from + index) {
if index >= 2 && &input[index - 2..index] == b"\r\n" {
return Found::Delimiter(index - 2);
}
if index >= 1 && matches!(input[index - 1], b'\r' | b'\n') {
return Found::Bare(index - 1);
}
from = index + 1;
}
Found::Absent
}
fn safe_length(&self, input: &[u8]) -> usize {
let dash_boundary = &self.delimiter[2..];
let start = input.len().saturating_sub(self.delimiter.len() - 1);
memchr::memchr2_iter(b'\r', b'\n', &input[start..])
.map(|index| start + index)
.find(|&index| {
self.delimiter.starts_with(&input[index..])
|| dash_boundary.starts_with(&input[index + 1..])
})
.unwrap_or(input.len())
}
fn check_section(&self, section: &[u8]) -> Result<(), MultipartErrorKind> {
let (mut count, mut rest) = (0, section);
while let Some(end) = memchr::memmem::find(rest, b"\r\n") {
let line = &rest[..end];
rest = &rest[end + 2..];
count += 1;
if count > self.limits.header_count_limit {
return Err(MultipartErrorKind::HeaderCountLimit);
}
if line
.first()
.is_some_and(|&byte| byte == b' ' || byte == b'\t')
{
return Err(MultipartErrorKind::FoldedField);
}
let colon = memchr::memchr(b':', line).ok_or(MultipartErrorKind::InvalidField)?;
let (name, value) = (&line[..colon], &line[colon + 1..]);
if name.is_empty()
|| !name.iter().all(|&byte| is_field_name_byte(byte))
|| !value.iter().all(|&byte| is_field_value_byte(byte))
{
return Err(MultipartErrorKind::InvalidField);
}
}
Ok(())
}
}