rialo_types/

headers.rs

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use std::fmt;

use borsh::{BorshDeserialize, BorshSerialize};
use serde::{Deserialize, Serialize};
use thiserror::Error;

use crate::RexValue;

type HeadersInnerType = std::collections::BTreeMap<String, RexValue>;

/// A structure representing the headers of an HTTP request.
#[derive(
    Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default, BorshSerialize, BorshDeserialize,
)]
pub struct Headers(HeadersInnerType);

#[derive(Error, Debug)]
pub enum HeadersError {
    #[error("Failed to decode {encoding} for parameter '{param}': {reason}")]
    DecodeError {
        param: String,
        encoding: String,
        reason: String,
    },

    #[error("JSON parsing error: {reason}")]
    JsonError { reason: String },
}

impl Headers {
    /// Creates a new Headers instance from the provided HeadersType map.
    ///
    /// # Arguments
    /// * `headers` - A HeadersInnerType containing header key-value pairs
    ///
    /// # Returns
    /// A new Headers instance wrapping the provided map
    pub fn new(headers: HeadersInnerType) -> Self {
        Headers(headers)
    }

    /// Inserts a plain text header value into the headers map.
    ///
    /// This method adds a header with a plain (unencrypted) value to the headers collection.
    ///
    /// # Arguments
    /// * `key` - The header name
    /// * `value` - The plain text header value
    pub fn insert_plain(&mut self, key: String, value: String) {
        self.0.insert(key, RexValue::Plain(value.into_bytes()));
    }

    /// Inserts an encrypted header value into the headers map.
    ///
    /// This method adds a header with an encrypted value to the headers collection.
    /// The encrypted value should be base64-encoded ciphertext that can be decrypted
    /// within the TEE using the stored secret key.
    ///
    /// # Arguments
    /// * `key` - The header name
    /// * `encrypted_value` - The base64-encoded encrypted header value
    pub fn insert_encrypted(&mut self, key: String, encrypted_value: String) {
        self.0
            .insert(key, RexValue::Encrypted(encrypted_value.into_bytes()));
    }

    /// Inserts encrypted raw bytes as a header value.
    ///
    /// Use this when you have raw ciphertext bytes (e.g., from HPKE encryption)
    /// rather than a base64-encoded string. The bytes will be internally base64-encoded
    /// to match the expected format for TEE decryption.
    ///
    /// # Arguments
    /// * `key` - The header name
    /// * `ciphertext` - The raw encrypted bytes
    #[cfg(feature = "non-pdk")]
    pub fn insert_encrypted_bytes(&mut self, key: String, ciphertext: Vec<u8>) {
        use fastcrypto::encoding::{Base64, Encoding};
        // Base64-encode the raw ciphertext so it can be decoded by the TEE's decode_headers
        let encoded = Base64::encode(&ciphertext);
        self.0
            .insert(key, RexValue::Encrypted(encoded.into_bytes()));
    }

    /// Provide a public accessor to move out the inner map.
    /// This enables other crates to consume the headers safely.
    ///
    /// # Returns
    /// The inner HeadersInnerType map, consuming the Headers instance
    pub fn into_inner(self) -> HeadersInnerType {
        self.0
    }

    /// Parses base64-encoded headers string into a Headers instance.
    ///
    /// This method decodes a base64-encoded string containing JSON headers data
    /// and deserializes it into a Headers structure.
    ///
    /// # Arguments
    /// * `headers_b64` - Base64-encoded string containing JSON headers
    ///
    /// # Returns
    /// * `Ok(Headers)` - Successfully parsed headers
    /// * `Err(HeadersError)` - If base64 decoding, UTF-8 conversion, or JSON parsing fails
    #[cfg(feature = "non-pdk")]
    pub fn parse_b64(headers_b64: &str) -> Result<Self, HeadersError> {
        use fastcrypto::{encoding, encoding::Encoding};
        let headers_bytes =
            encoding::Base64::decode(headers_b64).map_err(|err| HeadersError::DecodeError {
                param: "headers".to_string(),
                encoding: "base64".to_string(),
                reason: err.to_string(),
            })?;

        let headers_str =
            std::str::from_utf8(&headers_bytes).map_err(|err| HeadersError::DecodeError {
                param: "headers".to_string(),
                encoding: "utf-8".to_string(),
                reason: err.to_string(),
            })?;

        let parsed_headers: Headers =
            serde_json::from_str(headers_str).map_err(|error| HeadersError::JsonError {
                reason: format!("Invalid headers JSON: {error}"),
            })?;

        Ok(parsed_headers)
    }
}

impl std::str::FromStr for Headers {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let headers: HeadersInnerType = serde_json::from_str(s).map_err(|e| e.to_string())?;
        Ok(Headers(headers))
    }
}

impl fmt::Display for Headers {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", serde_json::to_string(&self.0).unwrap())
    }
}

#[cfg(test)]
mod tests {
    use std::str::FromStr;

    use super::*;

    #[test]
    fn test_headers_from_str_json() {
        let json = r#"{"A": {"Plain":[66]}}"#;
        let headers = Headers::from_str(json).unwrap();
        let got = headers.0.into_iter().collect::<Vec<(_, RexValue)>>();
        let expected = vec![("A".to_string(), RexValue::plain_string("B"))];
        assert_eq!(got, expected);
    }

    #[test]
    fn test_headers_from_str_with_string_values() {
        let json = r#"{"Mynonce": {"Plain":"abc123"}}"#;
        let headers = Headers::from_str(json).unwrap();
        let got = headers.0.into_iter().collect::<Vec<(_, RexValue)>>();
        let expected = vec![("Mynonce".to_string(), RexValue::plain_string("abc123"))];
        assert_eq!(got, expected);
    }
}