rig_http/http_client/

multipart.rs

//! Transport-neutral multipart fields and byte encoding.
//!
//! ```
//! use rig_http::http_client::multipart::MultipartForm;
//!
//! let (boundary, body) = MultipartForm::new().text("model", "example").encode();
//! assert!(!boundary.is_empty());
//! assert!(!body.is_empty());
//! ```

use bytes::Bytes;
use mime::Mime;
use std::borrow::Cow;

/// A generic multipart form part that can represent text or binary data
#[derive(Clone, Debug)]
pub struct Part {
    name: String,
    content: PartContent,
    filename: Option<String>,
    content_type: Option<Mime>,
}

/// The payload of a multipart [`Part`].
#[derive(Clone, Debug)]
pub enum PartContent {
    /// A text field.
    Text(String),
    /// Binary data (e.g. a file upload).
    Binary(Bytes),
}

impl Part {
    /// Create a text part
    pub fn text(name: impl Into<String>, value: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            content: PartContent::Text(value.into()),
            filename: None,
            content_type: None,
        }
    }

    /// Create a binary part (e.g., file upload)
    pub fn bytes(name: impl Into<String>, data: impl Into<Bytes>) -> Self {
        Self {
            name: name.into(),
            content: PartContent::Binary(data.into()),
            filename: None,
            content_type: None,
        }
    }

    /// Set the filename for this part
    pub fn filename(mut self, filename: impl Into<String>) -> Self {
        self.filename = Some(filename.into());
        self
    }

    /// Set the content type for this part
    pub fn content_type(mut self, content_type: Mime) -> Self {
        self.content_type = Some(content_type);
        self
    }

    /// Get the part name
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Get the filename if set
    pub fn get_filename(&self) -> Option<&str> {
        self.filename.as_deref()
    }

    /// Get the content type if set
    pub fn get_content_type(&self) -> Option<&Mime> {
        self.content_type.as_ref()
    }

    /// The part's payload.
    pub fn content(&self) -> &PartContent {
        &self.content
    }

    /// Split the part into its owned pieces: `(name, content, filename, content_type)`.
    /// Transports use this to render the part in their native multipart type.
    pub fn into_pieces(self) -> (String, PartContent, Option<String>, Option<Mime>) {
        (self.name, self.content, self.filename, self.content_type)
    }
}

/// Generic multipart form data container
#[derive(Clone, Debug, Default)]
pub struct MultipartForm {
    parts: Vec<Part>,
    boundary: Option<String>,
}

impl MultipartForm {
    /// Create a new empty multipart form
    pub fn new() -> Self {
        Self::default()
    }

    /// Add a part to the form
    pub fn part(mut self, part: Part) -> Self {
        self.parts.push(part);
        self
    }

    /// Add a text field
    pub fn text(self, name: impl Into<String>, value: impl Into<String>) -> Self {
        self.part(Part::text(name, value))
    }

    /// Add a file/binary field
    pub fn file(
        self,
        name: impl Into<String>,
        filename: impl Into<String>,
        content_type: Mime,
        data: impl Into<Bytes>,
    ) -> Self {
        self.part(
            Part::bytes(name, data)
                .filename(filename)
                .content_type(content_type),
        )
    }

    /// Set a custom boundary (optional, one will be generated if not set)
    pub fn boundary(mut self, boundary: impl Into<String>) -> Self {
        self.boundary = Some(boundary.into());
        self
    }

    /// Get the parts
    pub fn parts(&self) -> &[Part] {
        &self.parts
    }

    /// Consume the form, yielding its parts in order. Transports use this to
    /// render the form in their native multipart type.
    pub fn into_parts(self) -> Vec<Part> {
        self.parts
    }

    /// Generate a boundary string
    fn generate_boundary() -> String {
        use std::time::{SystemTime, UNIX_EPOCH};
        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_nanos();
        format!("----boundary{timestamp}")
    }

    /// Get or generate boundary
    fn get_boundary(&self) -> Cow<'_, str> {
        match &self.boundary {
            Some(b) => Cow::Borrowed(b),
            None => Cow::Owned(Self::generate_boundary()),
        }
    }

    /// Encodes parts in insertion order, returning the boundary and body bytes.
    /// Uses the configured boundary or generates one for this call. Names,
    /// filenames, and boundaries are inserted verbatim; callers must ensure they
    /// are safe for multipart headers and that the boundary does not occur in data.
    pub fn encode(&self) -> (String, Bytes) {
        let boundary = self.get_boundary();
        let mut body = Vec::new();

        for part in &self.parts {
            body.extend_from_slice(b"--");
            body.extend_from_slice(boundary.as_bytes());
            body.extend_from_slice(b"\r\n");

            body.extend_from_slice(b"Content-Disposition: form-data; name=\"");
            body.extend_from_slice(part.name.as_bytes());
            body.extend_from_slice(b"\"");

            if let Some(filename) = &part.filename {
                body.extend_from_slice(b"; filename=\"");
                body.extend_from_slice(filename.as_bytes());
                body.extend_from_slice(b"\"");
            }
            body.extend_from_slice(b"\r\n");

            if let Some(content_type) = &part.content_type {
                body.extend_from_slice(b"Content-Type: ");
                body.extend_from_slice(content_type.as_ref().as_bytes());
                body.extend_from_slice(b"\r\n");
            }

            body.extend_from_slice(b"\r\n");

            match &part.content {
                PartContent::Text(text) => body.extend_from_slice(text.as_bytes()),
                PartContent::Binary(bytes) => body.extend_from_slice(bytes),
            }

            body.extend_from_slice(b"\r\n");
        }

        body.extend_from_slice(b"--");
        body.extend_from_slice(boundary.as_bytes());
        body.extend_from_slice(b"--\r\n");

        (boundary.into_owned(), Bytes::from(body))
    }
}

#[cfg(test)]
mod tests;