use std::collections::HashMap;
use std::io::{self, BufRead, BufReader, Read, Write};
use crate::header::Header;
use super::MimeError;
#[allow(dead_code)]
pub struct Reader<R: Read> {
inner: BufReader<R>,
boundary: String,
delim: Vec<u8>,
done: bool,
}
impl<R: Read> Reader<R> {
pub fn new(inner: R, boundary: &str) -> Self {
let delim = format!("--{boundary}").into_bytes();
Self {
inner: BufReader::new(inner),
boundary: boundary.to_owned(),
delim,
done: false,
}
}
pub fn next_part(&mut self) -> Result<Option<Part<'_>>, MimeError> {
if self.done {
return Ok(None);
}
loop {
let mut line = String::new();
let n = self.inner
.read_line(&mut line)
.map_err(|e| MimeError::Other(e.to_string()))?;
if n == 0 {
self.done = true;
return Ok(None);
}
let trimmed = line.trim_end_matches(|c| c == '\r' || c == '\n');
if trimmed.as_bytes() == self.delim.as_slice() {
let header = read_part_headers(&mut self.inner)?;
return Ok(Some(Part {
header,
boundary: self.delim.clone(),
inner: &mut self.inner as *mut BufReader<R> as *mut (),
_marker: std::marker::PhantomData,
}));
}
let mut close = self.delim.clone();
close.extend_from_slice(b"--");
if trimmed.as_bytes() == close.as_slice() {
self.done = true;
return Ok(None);
}
}
}
pub fn read_form(&mut self, max_memory: i64) -> Result<Form, MimeError> {
let mut form = Form::default();
let mut mem_used: i64 = 0;
while let Some(mut part) = self.next_part()? {
let cd = part.header.get("Content-Disposition").unwrap_or("").to_owned();
let (_, params) = super::parse_media_type(&cd)
.unwrap_or_else(|_| ("".into(), HashMap::new()));
let name = params.get("name").cloned().unwrap_or_default();
let filename = params.get("filename").cloned();
let mut body = Vec::new();
part.read_to_end(&mut body)
.map_err(|e| MimeError::Other(e.to_string()))?;
if let Some(fname) = filename {
let size = body.len() as i64;
form.file
.entry(name)
.or_default()
.push(FileHeader {
filename: fname,
header: part.header.clone(),
size,
content: body,
});
} else {
mem_used += body.len() as i64;
if max_memory >= 0 && mem_used > max_memory {
return Err(MimeError::Other("multipart: message too large".into()));
}
if let Ok(s) = String::from_utf8(body) {
form.value.entry(name).or_default().push(s);
}
}
}
Ok(form)
}
}
fn read_part_headers<R: Read>(r: &mut BufReader<R>) -> Result<Header, MimeError> {
let mut h = Header::new();
loop {
let mut line = String::new();
let n = r
.read_line(&mut line)
.map_err(|e| MimeError::Other(e.to_string()))?;
if n == 0 {
break;
}
let trimmed = line.trim_end_matches(|c| c == '\r' || c == '\n');
if trimmed.is_empty() {
break;
}
if let Some(colon) = trimmed.find(':') {
let name = trimmed[..colon].trim();
let value = trimmed[colon + 1..].trim();
h.add(name, value);
}
}
Ok(h)
}
#[allow(dead_code)]
pub struct Part<'a> {
pub header: Header,
boundary: Vec<u8>,
inner: *mut (),
_marker: std::marker::PhantomData<&'a mut ()>,
}
impl<'a> Read for Part<'a> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
if buf.is_empty() {
return Ok(0);
}
unsafe {
let r = &mut *(self.inner as *mut BufReader<std::io::Empty>);
r.read(&mut buf[..1])
}
}
}
#[derive(Debug, Default)]
pub struct Form {
pub value: HashMap<String, Vec<String>>,
pub file: HashMap<String, Vec<FileHeader>>,
}
#[derive(Debug)]
pub struct FileHeader {
pub filename: String,
pub header: Header,
pub size: i64,
content: Vec<u8>,
}
impl FileHeader {
pub fn open(&self) -> impl Read + '_ {
std::io::Cursor::new(self.content.as_slice())
}
}
#[allow(dead_code)]
pub struct Writer<W: Write> {
inner: W,
boundary: String,
last: bool,
}
impl<W: Write> Writer<W> {
pub fn new(inner: W) -> Self {
let boundary = random_boundary();
Self { inner, boundary, last: false }
}
pub fn boundary(&self) -> &str {
&self.boundary
}
pub fn form_data_content_type(&self) -> String {
format!("multipart/form-data; boundary={}", self.boundary)
}
pub fn create_part(&mut self, header: Header) -> io::Result<PartWriter<'_, W>> {
write!(self.inner, "--{}\r\n", self.boundary)?;
header.write_to(&mut self.inner)?;
self.inner.write_all(b"\r\n")?;
Ok(PartWriter { inner: &mut self.inner })
}
pub fn create_form_field(&mut self, fieldname: &str) -> io::Result<PartWriter<'_, W>> {
let mut h = Header::new();
h.set(
"Content-Disposition",
format!("form-data; name=\"{fieldname}\""),
);
self.create_part(h)
}
pub fn create_form_file(
&mut self,
fieldname: &str,
filename: &str,
) -> io::Result<PartWriter<'_, W>> {
let mut h = Header::new();
h.set(
"Content-Disposition",
format!("form-data; name=\"{fieldname}\"; filename=\"{filename}\""),
);
h.set("Content-Type", "application/octet-stream");
self.create_part(h)
}
pub fn close(mut self) -> io::Result<W> {
write!(self.inner, "--{}--\r\n", self.boundary)?;
self.inner.flush()?;
Ok(self.inner)
}
}
pub struct PartWriter<'a, W: Write> {
inner: &'a mut W,
}
impl<'a, W: Write> Write for PartWriter<'a, W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.inner.write(buf)
}
fn flush(&mut self) -> io::Result<()> {
self.inner.flush()
}
}
fn random_boundary() -> String {
use std::time::{SystemTime, UNIX_EPOCH};
let t = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.subsec_nanos();
format!("{t:016x}x")
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
#[test]
fn writer_produces_valid_output() {
let mut buf = Vec::new();
{
let mut w = Writer::new(&mut buf);
let boundary = w.boundary().to_owned();
{
let mut part = w.create_form_field("greeting").unwrap();
part.write_all(b"Hello").unwrap();
}
w.close().unwrap();
let s = String::from_utf8(buf.clone()).unwrap();
assert!(s.contains(&format!("--{boundary}")));
assert!(s.contains("Content-Disposition: form-data; name=\"greeting\""));
assert!(s.contains("Hello"));
assert!(s.contains(&format!("--{boundary}--")));
}
}
}