use crate::Boundary;
use crate::syntax::{is_field_name_byte, is_field_value_byte};
use memchr::memmem::Finder;
use std::fmt;
#[derive(Clone, Debug)]
pub struct MultipartWriter {
boundary: Boundary,
delimiter: Vec<u8>,
finder: Finder<'static>,
state: State,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum State {
Start,
Content { matched: Option<usize> },
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum MultipartWriteError {
InvalidFieldName,
InvalidFieldValue,
Delimiter,
NoPart,
}
impl MultipartWriter {
pub fn new(boundary: Boundary) -> Self {
let delimiter = [b"\r\n--", boundary.as_str().as_bytes()].concat();
let finder = Finder::new(&delimiter[2..]).into_owned();
Self {
boundary,
delimiter,
finder,
state: State::Start,
}
}
pub fn boundary(&self) -> &Boundary {
&self.boundary
}
pub fn start_part<N, V>(
&mut self,
fields: impl IntoIterator<Item = (N, V)>,
output: &mut Vec<u8>,
) -> Result<(), MultipartWriteError>
where
N: AsRef<str>,
V: AsRef<[u8]>,
{
let length = output.len();
let result = self.write_head(fields, output);
if result.is_err() {
output.truncate(length);
}
result
}
fn write_head<N, V>(
&mut self,
fields: impl IntoIterator<Item = (N, V)>,
output: &mut Vec<u8>,
) -> Result<(), MultipartWriteError>
where
N: AsRef<str>,
V: AsRef<[u8]>,
{
let delimiter = if self.state == State::Start {
&self.delimiter[2..]
} else {
&self.delimiter[..]
};
output.extend_from_slice(delimiter);
output.extend_from_slice(b"\r\n");
for (name, value) in fields {
let (name, value) = (name.as_ref().as_bytes(), value.as_ref());
if name.is_empty() || !name.iter().all(|&byte| is_field_name_byte(byte)) {
return Err(MultipartWriteError::InvalidFieldName);
}
if !value.iter().all(|&byte| is_field_value_byte(byte)) {
return Err(MultipartWriteError::InvalidFieldValue);
}
let line = output.len();
output.extend_from_slice(name);
output.extend_from_slice(b": ");
output.extend_from_slice(value);
if output[line..].starts_with(&self.delimiter[2..]) {
return Err(MultipartWriteError::Delimiter);
}
output.extend_from_slice(b"\r\n");
}
output.extend_from_slice(b"\r\n");
self.state = State::Content { matched: Some(0) };
Ok(())
}
pub fn write_content(
&mut self,
content: &[u8],
output: &mut Vec<u8>,
) -> Result<(), MultipartWriteError> {
let State::Content { matched } = self.state else {
return Err(MultipartWriteError::NoPart);
};
let matched = self.scan(matched, content)?;
output.extend_from_slice(content);
self.state = State::Content { matched };
Ok(())
}
pub fn finish(self, output: &mut Vec<u8>) {
let delimiter = if self.state == State::Start {
&self.delimiter[2..]
} else {
&self.delimiter[..]
};
output.extend_from_slice(delimiter);
output.extend_from_slice(b"--\r\n");
}
fn scan(
&self,
matched: Option<usize>,
content: &[u8],
) -> Result<Option<usize>, MultipartWriteError> {
let dash_boundary = &self.delimiter[2..];
let continues = |matched: usize, after: &[u8]| {
let remaining = &dash_boundary[matched..];
let length = remaining.len().min(after.len());
(after[..length] == remaining[..length]).then_some(matched + length)
};
if let Some(matched) = matched.and_then(|matched| continues(matched, content)) {
return if matched == dash_boundary.len() {
Err(MultipartWriteError::Delimiter)
} else {
Ok(Some(matched))
};
}
if self
.finder
.find_iter(content)
.any(|index| index > 0 && matches!(content[index - 1], b'\r' | b'\n'))
{
return Err(MultipartWriteError::Delimiter);
}
let start = content.len().saturating_sub(dash_boundary.len());
Ok(memchr::memchr2_iter(b'\r', b'\n', &content[start..])
.find_map(|index| continues(0, &content[start + index + 1..])))
}
}
impl fmt::Display for MultipartWriteError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::InvalidFieldName => "invalid multipart part header field name",
Self::InvalidFieldValue => "invalid multipart part header field value",
Self::Delimiter => "multipart part contains its boundary delimiter",
Self::NoPart => "multipart content written before any part",
})
}
}
impl std::error::Error for MultipartWriteError {}