at_jet/

dual_format.rs

//! Dual-format support for AT-Jet
//!
//! Provides request/response types that support both Protobuf and JSON formats.
//! - Protobuf: Production format (efficient, compact, with evolution guarantees)
//! - JSON: Debug/testing format (human-readable, requires debug header)
//!
//! ## Why require a debug header for JSON?
//!
//! Protobuf provides schema evolution guarantees that JSON lacks:
//! - Field numbers enable backward/forward compatibility
//! - Unknown fields are preserved
//! - Optional fields have well-defined defaults
//!
//! If clients accidentally use JSON in production, they lose these guarantees
//! and may break when the schema evolves. The debug header requirement ensures
//! JSON is only used intentionally for debugging/testing.
//!
//! ## Format selection:
//!
//! - Request: Based on `Content-Type` header (JSON requires debug header)
//! - Response: Based on `Accept` header (JSON requires debug header)
//! - Debug header: `X-Debug-Format: <configured-secret>` (empty string = no secret required)
//!
//! ## Example:
//!
//! ```bash
//! # JSON request (requires debug header)
//! curl -X POST http://localhost:8080/api/quests/my \
//!   -H "Content-Type: application/json" \
//!   -H "Accept: application/json" \
//!   -H "X-Debug-Format: debug-secret-123" \
//!   -d '{"ucid": "user123"}'
//!
//! # Protobuf request (no debug header needed)
//! curl -X POST http://localhost:8080/api/quests/my \
//!   -H "Content-Type: application/x-protobuf" \
//!   --data-binary @request.pb
//! ```

use {crate::{content_types::{APPLICATION_JSON,
                             APPLICATION_PROTOBUF},
             error::JetError},
     axum::{async_trait,
            body::Bytes,
            extract::{FromRequest,
                      FromRequestParts,
                      Request},
            http::{StatusCode,
                   header::{ACCEPT,
                            CONTENT_TYPE},
                   request::Parts},
            response::{IntoResponse,
                       Response}},
     prost::Message,
     serde::{Serialize,
             de::DeserializeOwned},
     std::sync::OnceLock};

// ============================================================================
// Debug Header Configuration
// ============================================================================

/// Header name for enabling JSON debug format
pub const DEBUG_FORMAT_HEADER: &str = "x-debug-format";

/// Global list of authorized debug keys
static DEBUG_KEYS: OnceLock<Vec<String>> = OnceLock::new();

/// Configure the authorized debug keys for JSON format.
///
/// Call this once at application startup before handling any requests.
///
/// - Non-empty list: JSON requires `X-Debug-Format: <key>` header with one of the authorized keys
/// - Empty list: JSON format is completely disabled
///
/// If not configured, defaults to empty (JSON disabled).
///
/// # Example
///
/// ```rust,ignore
/// // In main.rs - enable JSON for specific keys
/// at_jet::dual_format::configure_debug_keys(vec![
///     "alice-dev-key".to_string(),
///     "bob-dev-key".to_string(),
///     "qa-team-key".to_string(),
/// ]);
///
/// // Or disable JSON completely
/// at_jet::dual_format::configure_debug_keys(vec![]);
/// ```
pub fn configure_debug_keys(keys: Vec<String>) {
  DEBUG_KEYS.set(keys).ok(); // Ignore if already set
}

/// Check if JSON format is allowed based on the debug header value
fn is_json_allowed(debug_header_value: Option<&str>) -> bool {
  match (DEBUG_KEYS.get(), debug_header_value) {
    | (Some(keys), Some(provided_key)) if !keys.is_empty() => {
      // Check if provided key is in the authorized list
      keys.iter().any(|k| k == provided_key)
    }
    | _ => {
      // No keys configured, empty list, or no header provided
      false
    }
  }
}

// ============================================================================
// Response Format
// ============================================================================

/// Response format preference
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ResponseFormat {
  /// Protobuf format (default, production)
  #[default]
  Protobuf,
  /// JSON format (debug/testing, requires valid debug header)
  Json,
}

impl ResponseFormat {
  /// Parse from Accept header, considering debug header authorization
  ///
  /// JSON format is only returned if:
  /// 1. Accept header contains "application/json"
  /// 2. Debug header is valid (checked via `is_json_allowed`)
  ///
  /// Otherwise, Protobuf is returned (safe default for production).
  pub fn from_headers(accept: Option<&str>, debug_header: Option<&str>) -> Self {
    let wants_json = accept.map(|s| s.contains("application/json")).unwrap_or(false);

    if wants_json && is_json_allowed(debug_header) {
      ResponseFormat::Json
    } else {
      ResponseFormat::Protobuf
    }
  }

  /// Check if this format is JSON
  pub fn is_json(&self) -> bool {
    matches!(self, ResponseFormat::Json)
  }
}

// ============================================================================
// Format Preference Extractor
// ============================================================================

/// Extractor for response format preference from Accept header
///
/// JSON format requires a valid debug header (`X-Debug-Format: <secret>`).
/// If debug header is missing or invalid, Protobuf is returned even if
/// Accept header requests JSON.
///
/// # Example
///
/// ```rust,ignore
/// async fn handler(format: AcceptFormat) -> ApiResponse<MyMessage> {
///     ApiResponse::ok(format.0, message)
/// }
/// ```
pub struct AcceptFormat(pub ResponseFormat);

#[async_trait]
impl<S> FromRequestParts<S> for AcceptFormat
where
  S: Send + Sync,
{
  type Rejection = std::convert::Infallible;

  async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
    let accept = parts.headers.get(ACCEPT).and_then(|v| v.to_str().ok());
    let debug_header = parts.headers.get(DEBUG_FORMAT_HEADER).and_then(|v| v.to_str().ok());
    Ok(AcceptFormat(ResponseFormat::from_headers(accept, debug_header)))
  }
}

// ============================================================================
// Dual-Format Request Extractor
// ============================================================================

/// Maximum request body size (10MB default)
const MAX_BODY_SIZE: usize = 10 * 1024 * 1024;

/// Dual-format request extractor
///
/// Extracts request body from either Protobuf or JSON format based on Content-Type.
/// Also captures the response format preference from Accept header.
///
/// # Example
///
/// ```rust,ignore
/// async fn create_user(
///     ApiRequest { body, format }: ApiRequest<CreateUserRequest>
/// ) -> ApiResponse<User> {
///     let user = process(body);
///     ApiResponse::ok(format, user)
/// }
/// ```
pub struct ApiRequest<T> {
  /// The decoded request body
  pub body:   T,
  /// The preferred response format (from Accept header)
  pub format: ResponseFormat,
}

impl<T> ApiRequest<T> {
  /// Create an OK response with the captured format preference
  pub fn ok<R>(self, response: R) -> ApiResponse<R>
  where
    R: Message + Serialize, {
    ApiResponse::ok(self.format, response)
  }

  /// Create a response with custom status and the captured format preference
  pub fn respond<R>(self, status: StatusCode, response: R) -> ApiResponse<R>
  where
    R: Message + Serialize, {
    ApiResponse::new(self.format, status, response)
  }

  /// Create a Created (201) response with the captured format preference
  pub fn created<R>(self, response: R) -> ApiResponse<R>
  where
    R: Message + Serialize, {
    ApiResponse::created(self.format, response)
  }
}

#[async_trait]
impl<S, T> FromRequest<S> for ApiRequest<T>
where
  S: Send + Sync,
  T: Message + Default + DeserializeOwned,
{
  type Rejection = JetError;

  async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
    // Extract headers before consuming the request
    let accept = req
      .headers()
      .get(ACCEPT)
      .and_then(|v| v.to_str().ok())
      .map(|s| s.to_string());

    let debug_header = req
      .headers()
      .get(DEBUG_FORMAT_HEADER)
      .and_then(|v| v.to_str().ok())
      .map(|s| s.to_string());

    let content_type = req
      .headers()
      .get(CONTENT_TYPE)
      .and_then(|v| v.to_str().ok())
      .unwrap_or("")
      .to_string();

    // Determine response format (considers debug header)
    let format = ResponseFormat::from_headers(accept.as_deref(), debug_header.as_deref());

    // Check if JSON input is allowed
    let json_allowed = is_json_allowed(debug_header.as_deref());
    let wants_json_input = content_type.contains("application/json");

    // Extract body bytes (consumes req)
    let bytes = Bytes::from_request(req, state)
      .await
      .map_err(|e| JetError::BadRequest(format!("Failed to read body: {}", e)))?;

    // Check size
    if bytes.len() > MAX_BODY_SIZE {
      return Err(JetError::BodyTooLarge {
        size: bytes.len(),
        max:  MAX_BODY_SIZE,
      });
    }

    // Decode based on Content-Type
    let body = if wants_json_input {
      if !json_allowed {
        return Err(JetError::InvalidContentType {
          expected: APPLICATION_PROTOBUF.to_string(),
          actual:   "application/json (requires valid X-Debug-Format header)".to_string(),
        });
      }
      serde_json::from_slice(&bytes).map_err(|e| JetError::BadRequest(format!("Invalid JSON: {}", e)))?
    } else if content_type.contains(APPLICATION_PROTOBUF) || bytes.is_empty() {
      // Default to protobuf, also handle empty body (for GET-like requests)
      if bytes.is_empty() {
        T::default()
      } else {
        T::decode(bytes)?
      }
    } else {
      // Unknown content type - try protobuf (production default)
      T::decode(bytes)?
    };

    Ok(ApiRequest { body, format })
  }
}

// ============================================================================
// Dual-Format Response
// ============================================================================

/// Dual-format response type
///
/// Returns either Protobuf or JSON based on the format preference.
///
/// # Example
///
/// ```rust,ignore
/// async fn get_user(
///     AcceptFormat(format): AcceptFormat,
///     Path(id): Path<i32>,
/// ) -> ApiResponse<User> {
///     let user = fetch_user(id);
///     ApiResponse::ok(format, user)
/// }
/// ```
pub struct ApiResponse<T>
where
  T: Message + Serialize, {
  format:  ResponseFormat,
  status:  StatusCode,
  message: T,
}

impl<T> ApiResponse<T>
where
  T: Message + Serialize,
{
  /// Create a new response with specified format and status
  pub fn new(format: ResponseFormat, status: StatusCode, message: T) -> Self {
    Self {
      format,
      status,
      message,
    }
  }

  /// Create a 200 OK response
  pub fn ok(format: ResponseFormat, message: T) -> Self {
    Self::new(format, StatusCode::OK, message)
  }

  /// Create a 201 Created response
  pub fn created(format: ResponseFormat, message: T) -> Self {
    Self::new(format, StatusCode::CREATED, message)
  }

  /// Create a 202 Accepted response
  pub fn accepted(format: ResponseFormat, message: T) -> Self {
    Self::new(format, StatusCode::ACCEPTED, message)
  }
}

impl<T> IntoResponse for ApiResponse<T>
where
  T: Message + Serialize,
{
  fn into_response(self) -> Response {
    match self.format {
      | ResponseFormat::Json => {
        match serde_json::to_vec(&self.message) {
          | Ok(bytes) => (self.status, [(CONTENT_TYPE, APPLICATION_JSON)], bytes).into_response(),
          | Err(e) => {
            // Fallback to error response if JSON serialization fails
            (
              StatusCode::INTERNAL_SERVER_ERROR,
              [(CONTENT_TYPE, APPLICATION_JSON)],
              format!("{{\"error\": \"JSON serialization failed: {}\"}}", e),
            )
              .into_response()
          }
        }
      }
      | ResponseFormat::Protobuf => {
        let bytes = self.message.encode_to_vec();
        (self.status, [(CONTENT_TYPE, APPLICATION_PROTOBUF)], bytes).into_response()
      }
    }
  }
}

// ============================================================================
// Result Type
// ============================================================================

/// Result type for dual-format responses
pub type ApiResult<T> = Result<ApiResponse<T>, JetError>;

#[cfg(test)]
mod tests {
  use super::*;

  #[test]
  fn test_response_format_without_debug_keys() {
    // Without configuring debug keys, JSON should be disabled
    // (default state - OnceLock not set returns None)

    // All requests should get Protobuf when debug keys not configured
    assert_eq!(ResponseFormat::from_headers(None, None), ResponseFormat::Protobuf);
    assert_eq!(
      ResponseFormat::from_headers(Some("application/x-protobuf"), None),
      ResponseFormat::Protobuf
    );
    // Even if Accept is JSON, without debug header it should be Protobuf
    assert_eq!(
      ResponseFormat::from_headers(Some("application/json"), None),
      ResponseFormat::Protobuf
    );
    assert_eq!(
      ResponseFormat::from_headers(Some("application/json"), Some("any-key")),
      ResponseFormat::Protobuf
    );
  }

  #[test]
  fn test_is_json_allowed_not_configured() {
    // When not configured, JSON should be disabled
    assert!(!is_json_allowed(None));
    assert!(!is_json_allowed(Some("any-value")));
  }

  #[test]
  fn test_response_format_is_json() {
    assert!(!ResponseFormat::Protobuf.is_json());
    assert!(ResponseFormat::Json.is_json());
  }
}