use crate::error::ContextError;
use crate::secret::{SecretAccessError, SecretResolver, SecretResolverError};
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap};
use std::fmt::Debug;
use std::sync::Arc;
use tracing::{debug, warn};
pub struct Context<T = serde_json::Value>
where
T: Serialize + for<'de> Deserialize<'de> + Debug,
{
data: HashMap<String, T>,
secrets: Option<Arc<dyn SecretResolver>>,
}
impl<T> Debug for Context<T>
where
T: Serialize + for<'de> Deserialize<'de> + Debug,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Context")
.field("data", &self.data)
.field(
"secrets",
&self.secrets.as_ref().map(|_| "<redacted resolver>"),
)
.finish()
}
}
impl<T> Context<T>
where
T: Serialize + for<'de> Deserialize<'de> + Debug,
{
pub fn new() -> Self {
debug!("Creating new empty context");
Self {
data: HashMap::new(),
secrets: None,
}
}
pub fn clone_data(&self) -> Self
where
T: Clone,
{
debug!("Cloning context data");
Self {
data: self.data.clone(),
secrets: self.secrets.clone(),
}
}
pub fn insert(&mut self, key: impl Into<String>, value: T) -> Result<(), ContextError> {
let key = key.into();
if self.data.contains_key(&key) {
warn!("Attempted to insert duplicate key: {}", key);
return Err(ContextError::KeyExists(key));
}
debug!("Inserting value for key: {}", key);
self.data.insert(key, value);
Ok(())
}
pub fn update(&mut self, key: impl Into<String>, value: T) -> Result<(), ContextError> {
let key = key.into();
if !self.data.contains_key(&key) {
warn!("Attempted to update non-existent key: {}", key);
return Err(ContextError::KeyNotFound(key));
}
debug!("Updating value for key: {}", key);
self.data.insert(key, value);
Ok(())
}
pub fn get(&self, key: &str) -> Option<&T> {
debug!("Getting value for key: {}", key);
self.data.get(key)
}
pub fn remove(&mut self, key: &str) -> Option<T> {
debug!("Removing value for key: {}", key);
self.data.remove(key)
}
pub fn data(&self) -> &HashMap<String, T> {
&self.data
}
pub fn into_data(self) -> HashMap<String, T> {
self.data
}
pub fn from_data(data: HashMap<String, T>) -> Self {
Self {
data,
secrets: None,
}
}
pub fn to_json(&self) -> Result<String, ContextError> {
debug!("Serializing context to JSON");
let json = serde_json::to_string(&self.data)?;
debug!("Context serialized successfully");
Ok(json)
}
pub fn from_json(json: String) -> Result<Self, ContextError> {
debug!("Deserializing context from JSON");
let data = serde_json::from_str(&json)?;
debug!("Context deserialized successfully");
Ok(Self {
data,
secrets: None,
})
}
pub fn with_secret_resolver(mut self, resolver: Arc<dyn SecretResolver>) -> Self {
self.secrets = Some(resolver);
self
}
pub fn set_secret_resolver(&mut self, resolver: Arc<dyn SecretResolver>) {
self.secrets = Some(resolver);
}
pub fn has_secret_resolver(&self) -> bool {
self.secrets.is_some()
}
pub async fn secret(&self, name: &str) -> Result<BTreeMap<String, String>, SecretAccessError> {
let resolver = self
.secrets
.as_ref()
.ok_or(SecretAccessError::NotConfigured)?;
let effective = self.resolve_secret_alias(name);
resolver.resolve(&effective).await.map_err(|e| match e {
SecretResolverError::NotFound(n) => SecretAccessError::NotFound(n),
SecretResolverError::NotGranted(n) => SecretAccessError::NotGranted(n),
SecretResolverError::Backend(m) => SecretAccessError::Backend(m),
})
}
fn resolve_secret_alias(&self, name: &str) -> String {
if let Some(v) = self.data.get(crate::secret::SECRET_REFS_KEY) {
if let Ok(serde_json::Value::Object(map)) = serde_json::to_value(v) {
if let Some(serde_json::Value::String(target)) = map.get(name) {
return target.clone();
}
}
}
name.to_string()
}
pub async fn secret_field(&self, name: &str, field: &str) -> Result<String, SecretAccessError> {
let fields = self.secret(name).await?;
fields
.get(field)
.cloned()
.ok_or_else(|| SecretAccessError::FieldNotFound {
secret: name.to_string(),
field: field.to_string(),
})
}
}
impl Context<serde_json::Value> {
pub fn get_as<V>(&self, key: &str) -> Result<Option<V>, crate::error::TaskError>
where
V: serde::de::DeserializeOwned,
{
match self.data.get(key) {
None => Ok(None),
Some(value) => serde_json::from_value(value.clone())
.map(Some)
.map_err(|e| crate::error::TaskError::ValidationFailed {
message: format!(
"context key '{}' could not be read as {}: {}",
key,
std::any::type_name::<V>(),
e
),
}),
}
}
pub fn get_required<V>(&self, key: &str) -> Result<V, crate::error::TaskError>
where
V: serde::de::DeserializeOwned,
{
match self.get_as(key)? {
Some(value) => Ok(value),
None => Err(crate::error::TaskError::ValidationFailed {
message: format!(
"required context key '{}' is missing (expected {})",
key,
std::any::type_name::<V>()
),
}),
}
}
pub fn insert_as<V>(
&mut self,
key: impl Into<String>,
value: V,
) -> Result<(), crate::error::TaskError>
where
V: serde::Serialize,
{
let key = key.into();
let json =
serde_json::to_value(value).map_err(|e| crate::error::TaskError::ValidationFailed {
message: format!("context key '{}' could not be serialized: {}", key, e),
})?;
self.data.insert(key, json);
Ok(())
}
}
impl<T> Default for Context<T>
where
T: Serialize + for<'de> Deserialize<'de> + Debug,
{
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn setup_test_context() -> Context<i32> {
Context::new()
}
#[test]
fn test_context_operations() {
let mut context = setup_test_context();
assert!(context.data.is_empty());
context.insert("test", 42).unwrap();
assert_eq!(context.get("test"), Some(&42));
assert!(matches!(
context.insert("test", 43),
Err(ContextError::KeyExists(_))
));
context.update("test", 43).unwrap();
assert_eq!(context.get("test"), Some(&43));
assert!(matches!(
context.update("nonexistent", 42),
Err(ContextError::KeyNotFound(_))
));
}
#[test]
fn test_context_serialization() {
let mut context = setup_test_context();
context.insert("test", 42).unwrap();
let json = context.to_json().unwrap();
let deserialized = Context::<i32>::from_json(json).unwrap();
assert_eq!(deserialized.get("test"), Some(&42));
}
#[test]
fn test_context_clone_data() {
let mut context = Context::<i32>::new();
context.insert("a", 1).unwrap();
context.insert("b", 2).unwrap();
let cloned = context.clone_data();
assert_eq!(cloned.get("a"), Some(&1));
assert_eq!(cloned.get("b"), Some(&2));
}
#[test]
fn test_context_from_data() {
let mut data = HashMap::new();
data.insert("key".to_string(), 42);
let context = Context::from_data(data);
assert_eq!(context.get("key"), Some(&42));
}
#[test]
fn test_context_into_data() {
let mut context = Context::<i32>::new();
context.insert("key", 42).unwrap();
let data = context.into_data();
assert_eq!(data.get("key"), Some(&42));
}
#[test]
fn test_typed_accessors_roundtrip() {
let mut ctx = Context::new();
ctx.insert_as("count", 7u32).unwrap();
ctx.insert_as("name", "ada").unwrap();
assert_eq!(ctx.get_as::<u32>("count").unwrap(), Some(7));
assert_eq!(ctx.get_as::<String>("missing").unwrap(), None);
let name: String = ctx.get_required("name").unwrap();
assert_eq!(name, "ada");
ctx.insert_as("count", 100u32).unwrap();
assert_eq!(ctx.get_as::<u32>("count").unwrap(), Some(100));
}
use async_trait::async_trait;
struct StubResolver {
secrets: HashMap<String, BTreeMap<String, String>>,
}
#[async_trait]
impl SecretResolver for StubResolver {
async fn resolve(
&self,
name: &str,
) -> Result<BTreeMap<String, String>, SecretResolverError> {
self.secrets
.get(name)
.cloned()
.ok_or_else(|| SecretResolverError::NotFound(name.to_string()))
}
}
fn stub_resolver() -> Arc<dyn SecretResolver> {
let mut db = BTreeMap::new();
db.insert("host".to_string(), "db.internal".to_string());
db.insert("password".to_string(), "s3cr3t-p@ss".to_string());
let mut secrets = HashMap::new();
secrets.insert("db_prod".to_string(), db);
Arc::new(StubResolver { secrets })
}
#[test]
fn test_secret_resolver_field_is_not_serialized() {
let mut ctx = Context::<serde_json::Value>::new();
ctx.insert("visible", serde_json::json!("in-context"))
.unwrap();
ctx.set_secret_resolver(stub_resolver());
assert!(ctx.has_secret_resolver());
let json = ctx.to_json().unwrap();
assert!(json.contains("visible"));
assert!(
!json.contains("s3cr3t-p@ss"),
"secret leaked into serialized Context: {json}"
);
assert!(!json.contains("secrets"));
assert!(!json.contains("resolver"));
let restored = Context::<serde_json::Value>::from_json(json).unwrap();
assert!(!restored.has_secret_resolver());
}
#[test]
fn test_debug_redacts_resolver_and_never_prints_plaintext() {
let mut ctx = Context::<serde_json::Value>::new();
ctx.set_secret_resolver(stub_resolver());
let dbg = format!("{:?}", ctx);
assert!(dbg.contains("<redacted resolver>"), "debug: {dbg}");
assert!(
!dbg.contains("s3cr3t-p@ss"),
"secret leaked into Debug: {dbg}"
);
}
#[tokio::test]
async fn test_secret_accessor_not_configured_errors_clearly() {
let ctx = Context::<serde_json::Value>::new();
let err = ctx.secret("db_prod").await.unwrap_err();
assert!(matches!(err, SecretAccessError::NotConfigured));
}
#[tokio::test]
async fn test_secret_accessor_happy_path() {
let ctx = Context::<serde_json::Value>::new().with_secret_resolver(stub_resolver());
let fields = ctx.secret("db_prod").await.unwrap();
assert_eq!(fields.get("password").unwrap(), "s3cr3t-p@ss");
assert_eq!(
ctx.secret_field("db_prod", "host").await.unwrap(),
"db.internal"
);
}
#[tokio::test]
async fn test_secret_accessor_resolves_through_alias_map() {
let mut ctx = Context::<serde_json::Value>::new().with_secret_resolver(stub_resolver());
ctx.insert(
crate::secret::SECRET_REFS_KEY,
serde_json::json!({"dst": "db_prod"}),
)
.unwrap();
let fields = ctx.secret("dst").await.unwrap();
assert_eq!(fields.get("password").unwrap(), "s3cr3t-p@ss");
assert_eq!(
ctx.secret_field("dst", "host").await.unwrap(),
"db.internal"
);
assert_eq!(
ctx.secret("db_prod")
.await
.unwrap()
.get("password")
.unwrap(),
"s3cr3t-p@ss"
);
assert!(matches!(
ctx.secret("unmapped").await.unwrap_err(),
SecretAccessError::NotFound(_)
));
let json = ctx.to_json().unwrap();
assert!(!json.contains("s3cr3t-p@ss"), "secret leaked: {json}");
}
#[tokio::test]
async fn test_secret_accessor_missing_name_and_field() {
let ctx = Context::<serde_json::Value>::new().with_secret_resolver(stub_resolver());
assert!(matches!(
ctx.secret("absent").await.unwrap_err(),
SecretAccessError::NotFound(_)
));
assert!(matches!(
ctx.secret_field("db_prod", "absent_field")
.await
.unwrap_err(),
SecretAccessError::FieldNotFound { .. }
));
}
#[test]
fn test_typed_accessor_errors_are_actionable() {
let mut ctx = Context::new();
ctx.insert("count", serde_json::json!("not-a-number"))
.unwrap();
let err = ctx.get_as::<u32>("count").unwrap_err();
let msg = err.to_string();
assert!(msg.contains("count"), "msg should name the key: {msg}");
let err = ctx.get_required::<u32>("absent").unwrap_err();
let msg = err.to_string();
assert!(msg.contains("absent"), "msg should name the key: {msg}");
}
}