#![forbid(unsafe_code)]
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use bytes::{BufMut, Bytes, BytesMut};
static BOUNDARY_COUNTER: AtomicU64 = AtomicU64::new(0);
fn generate_boundary() -> String {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.subsec_nanos();
let counter = BOUNDARY_COUNTER.fetch_add(1, Ordering::Relaxed);
format!("----OxiHTTPBoundary{nanos:08x}{counter:04x}")
}
#[derive(Debug, Clone)]
pub struct Part {
headers: Vec<(String, String)>,
body: Bytes,
}
impl Part {
pub fn text(name: &str, value: impl Into<String>) -> Self {
Self {
headers: vec![(
"Content-Disposition".into(),
format!("form-data; name=\"{name}\""),
)],
body: Bytes::from(value.into()),
}
}
pub fn file(name: &str, filename: &str, content_type: &str, body: impl Into<Bytes>) -> Self {
Self {
headers: vec![
(
"Content-Disposition".into(),
format!("form-data; name=\"{name}\"; filename=\"{filename}\""),
),
("Content-Type".into(), content_type.to_owned()),
],
body: body.into(),
}
}
pub fn custom(headers: Vec<(String, String)>, body: impl Into<Bytes>) -> Self {
Self {
headers,
body: body.into(),
}
}
}
#[derive(Debug, Clone)]
pub struct MultipartBuilder {
boundary: String,
parts: Vec<Part>,
}
impl MultipartBuilder {
pub fn new() -> Self {
Self {
boundary: generate_boundary(),
parts: Vec::new(),
}
}
pub fn boundary(&self) -> &str {
&self.boundary
}
pub fn content_type(&self) -> String {
format!("multipart/form-data; boundary={}", self.boundary)
}
pub fn add_text(mut self, name: &str, value: impl Into<String>) -> Self {
self.parts.push(Part::text(name, value));
self
}
pub fn add_file(
mut self,
name: &str,
filename: &str,
content_type: &str,
body: impl Into<Bytes>,
) -> Self {
self.parts
.push(Part::file(name, filename, content_type, body));
self
}
pub fn add_part(mut self, part: Part) -> Self {
self.parts.push(part);
self
}
pub fn build(self) -> Bytes {
let boundary = self.find_unique_boundary();
let dash_boundary = format!("--{boundary}");
let final_boundary = format!("--{boundary}--\r\n");
let mut buf = BytesMut::new();
for part in &self.parts {
buf.put_slice(dash_boundary.as_bytes());
buf.put_slice(b"\r\n");
for (k, v) in &part.headers {
buf.put_slice(k.as_bytes());
buf.put_slice(b": ");
buf.put_slice(v.as_bytes());
buf.put_slice(b"\r\n");
}
buf.put_slice(b"\r\n");
buf.put_slice(&part.body);
buf.put_slice(b"\r\n");
}
buf.put_slice(final_boundary.as_bytes());
buf.freeze()
}
fn find_unique_boundary(&self) -> String {
let mut boundary = self.boundary.clone();
let mut suffix = 0u32;
loop {
let has_collision = self.parts.iter().any(|p| {
p.body
.windows(boundary.len())
.any(|w| w == boundary.as_bytes())
});
if !has_collision {
return boundary;
}
suffix += 1;
boundary = format!("{}{suffix:04x}", self.boundary);
}
}
}
impl Default for MultipartBuilder {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_empty_builder() {
let builder = MultipartBuilder::new();
let ct = builder.content_type();
assert!(ct.starts_with("multipart/form-data; boundary="));
let bytes = builder.build();
let s = String::from_utf8(bytes.to_vec()).unwrap();
assert!(s.contains("----OxiHTTPBoundary"));
assert!(s.ends_with("--\r\n"));
}
#[test]
fn test_text_field() {
let bytes = MultipartBuilder::new()
.add_text("field1", "hello world")
.build();
let s = String::from_utf8(bytes.to_vec()).unwrap();
assert!(s.contains("name=\"field1\""));
assert!(s.contains("hello world"));
}
#[test]
fn test_file_part() {
let bytes = MultipartBuilder::new()
.add_file("upload", "test.txt", "text/plain", "file contents")
.build();
let s = String::from_utf8(bytes.to_vec()).unwrap();
assert!(s.contains("filename=\"test.txt\""));
assert!(s.contains("Content-Type: text/plain"));
assert!(s.contains("file contents"));
}
#[test]
fn test_mixed_parts() {
let bytes = MultipartBuilder::new()
.add_text("name", "Alice")
.add_file("avatar", "pic.png", "image/png", b"\x89PNG\r\n".as_ref())
.build();
let bytes_vec = bytes.to_vec();
let header_section = &bytes_vec[..];
assert!(
bytes_vec
.windows(b"name=\"name\"".len())
.any(|w| w == b"name=\"name\""),
"missing name field"
);
assert!(
bytes_vec.windows(b"Alice".len()).any(|w| w == b"Alice"),
"missing Alice"
);
assert!(
header_section
.windows(b"filename=\"pic.png\"".len())
.any(|w| w == b"filename=\"pic.png\""),
"missing filename"
);
}
#[test]
fn test_boundary_collision_resolved() {
let mut builder = MultipartBuilder::new();
let boundary_clone = builder.boundary().to_owned();
builder = builder.add_text("field", boundary_clone.as_str());
let bytes = builder.build();
let s = String::from_utf8(bytes.to_vec()).unwrap();
assert!(s.ends_with("--\r\n"));
}
#[test]
fn test_content_type_header() {
let b = MultipartBuilder::new();
let ct = b.content_type();
let bnd = b.boundary().to_owned();
assert_eq!(ct, format!("multipart/form-data; boundary={bnd}"));
}
#[test]
fn test_crlf_format() {
let bytes = MultipartBuilder::new().add_text("x", "y").build();
let s = String::from_utf8(bytes.to_vec()).unwrap();
assert!(s.contains("\r\n\r\n"));
assert!(s.contains("y\r\n"));
}
#[test]
fn test_boundary_collision_unique() {
let mut b = MultipartBuilder::new();
let bnd = b.boundary().to_owned();
let body_with_boundary = format!("some text {bnd} more text");
b = b.add_text("collision_field", &body_with_boundary);
let bytes = b.build();
let s = String::from_utf8(bytes.to_vec()).unwrap();
assert!(s.ends_with("--\r\n"));
}
#[test]
fn test_custom_part() {
let bytes = MultipartBuilder::new()
.add_part(Part::custom(
vec![
(
"Content-Disposition".into(),
"form-data; name=\"raw\"".into(),
),
("X-Custom".into(), "header-value".into()),
],
Bytes::from("raw body"),
))
.build();
let s = String::from_utf8(bytes.to_vec()).unwrap();
assert!(s.contains("X-Custom: header-value"));
assert!(s.contains("raw body"));
}
}