eric / rustwarden

//! Custom extractors for form data with duplicate fields
//!
//! This module provides extractors that handle HTML form data with duplicate field names,
//! which is common when multiple checkboxes share the same name (e.g., roles checkboxes).

use axum::{
    extract::{FromRequest, Request},
    http::StatusCode,
    response::{IntoResponse, Response},
};
use serde::Deserialize;
use std::collections::HashMap;

/// Extract form data that handles duplicate field names
///
/// This extractor parses URL-encoded form data and handles duplicate field names
/// by collecting all values into a vector. This is needed for HTML forms where
/// multiple checkboxes can have the same name (e.g., `roles=user&roles=admin`).
///
/// # Usage
///
/// ```rust,ignore
/// pub async fn create_user(
///     FormWithDuplicates(form_data): FormWithDuplicates<CreateUserRequest>,
/// ) -> Result<impl IntoResponse, AppError> {
///     // form_data.roles will be a Vec<String> with all selected roles
/// }
/// ```
#[derive(Debug)]
pub struct FormWithDuplicates<T>(pub T);

impl<T, S> FromRequest<S> for FormWithDuplicates<T>
where
    T: for<'de> Deserialize<'de>,
    S: Send + Sync,
{
    type Rejection = FormRejection;

    async fn from_request(req: Request, _state: &S) -> Result<Self, Self::Rejection> {
        // Check content type
        let content_type = req
            .headers()
            .get("content-type")
            .and_then(|h| h.to_str().ok())
            .unwrap_or("");

        if !content_type.starts_with("application/x-www-form-urlencoded") {
            return Err(FormRejection::InvalidContentType);
        }

        // Extract body
        let body_bytes = axum::body::to_bytes(req.into_body(), usize::MAX)
            .await
            .map_err(|_| FormRejection::BodyReadError)?;

        let body_str =
            std::str::from_utf8(&body_bytes).map_err(|_| FormRejection::InvalidEncoding)?;

        // Parse URL-encoded data and handle duplicates
        let parsed_map = parse_urlencoded_with_duplicates(body_str)?;

        // Convert to JSON value for deserialization
        let json_value = serde_json::to_value(&parsed_map)
            .map_err(|e| FormRejection::DeserializeError(e.to_string()))?;

        // Deserialize into target type
        let deserialized = T::deserialize(&json_value)
            .map_err(|e| FormRejection::DeserializeError(e.to_string()))?;

        Ok(FormWithDuplicates(deserialized))
    }
}

/// Parse URL-encoded string and handle duplicate keys by collecting them into vectors
fn parse_urlencoded_with_duplicates(
    input: &str,
) -> Result<HashMap<String, serde_json::Value>, FormRejection> {
    let mut map: HashMap<String, Vec<String>> = HashMap::new();

    for pair in input.split('&') {
        if pair.is_empty() {
            continue;
        }

        let (key, value) = if let Some(eq_pos) = pair.find('=') {
            // Simple URL decoding - replace %XX with actual characters
            let key =
                percent_decode(&pair[..eq_pos]).map_err(|_| FormRejection::InvalidEncoding)?;
            let value =
                percent_decode(&pair[eq_pos + 1..]).map_err(|_| FormRejection::InvalidEncoding)?;
            (key, value)
        } else {
            let key = percent_decode(pair).map_err(|_| FormRejection::InvalidEncoding)?;
            (key, String::new())
        };

        map.entry(key).or_default().push(value);
    }

    // Convert to HashMap<String, Value> where single values are strings and multiple values are arrays
    let mut result = HashMap::new();
    for (key, values) in map {
        if values.len() == 1 {
            result.insert(key, serde_json::Value::String(values[0].clone()));
        } else {
            let array: Vec<serde_json::Value> =
                values.into_iter().map(serde_json::Value::String).collect();
            result.insert(key, serde_json::Value::Array(array));
        }
    }

    Ok(result)
}

/// Simple percent decoding for URL-encoded strings
/// Handles UTF-8 characters correctly by collecting bytes first
fn percent_decode(input: &str) -> Result<String, ()> {
    let mut result = Vec::new();
    let bytes = input.as_bytes();
    let mut i = 0;

    while i < bytes.len() {
        if bytes[i] == b'%' {
            if i + 2 >= bytes.len() {
                return Err(());
            }
            let hex = &input[i + 1..i + 3];
            if let Ok(byte) = u8::from_str_radix(hex, 16) {
                result.push(byte);
                i += 3;
            } else {
                return Err(());
            }
        } else if bytes[i] == b'+' {
            result.push(b' ');
            i += 1;
        } else {
            result.push(bytes[i]);
            i += 1;
        }
    }

    String::from_utf8(result).map_err(|_| ())
}

/// Rejection type for form extraction errors
#[derive(Debug)]
pub enum FormRejection {
    /// Content-Type header is not `application/x-www-form-urlencoded`.
    InvalidContentType,
    /// Failed to read the request body bytes.
    BodyReadError,
    /// Body bytes are not valid UTF-8 or percent-encoding is malformed.
    InvalidEncoding,
    /// Deserialization of the form data into the target type failed.
    DeserializeError(String),
}

impl IntoResponse for FormRejection {
    fn into_response(self) -> Response {
        let (status, message) = match self {
            FormRejection::InvalidContentType => (
                StatusCode::BAD_REQUEST,
                "Invalid content type. Expected application/x-www-form-urlencoded".to_string(),
            ),
            FormRejection::BodyReadError => (
                StatusCode::BAD_REQUEST,
                "Failed to read request body".to_string(),
            ),
            FormRejection::InvalidEncoding => {
                (StatusCode::BAD_REQUEST, "Invalid URL encoding".to_string())
            }
            FormRejection::DeserializeError(msg) => (
                StatusCode::BAD_REQUEST,
                format!("Failed to deserialize form: {}", msg),
            ),
        };

        (status, message).into_response()
    }
}