use std::collections::HashMap;
use std::sync::Arc;
use adk_core::{AdkError, Result, SecretRequest, SecretService};
use async_trait::async_trait;
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct SecretGrant {
names: Vec<String>,
prefixes: Vec<String>,
}
impl SecretGrant {
pub fn none() -> Self {
Self::default()
}
#[must_use]
pub fn name(mut self, name: impl Into<String>) -> Self {
self.names.push(name.into());
self
}
#[must_use]
pub fn prefix(mut self, prefix: impl Into<String>) -> Self {
self.prefixes.push(prefix.into());
self
}
fn allows(&self, name: &str) -> bool {
self.names.iter().any(|allowed| allowed == name)
|| self.prefixes.iter().any(|prefix| name.starts_with(prefix.as_str()))
}
}
pub trait SecretAuditSink: Send + Sync {
fn record(&self, decision: SecretAccessDecision<'_>);
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SecretAccessDecision<'a> {
pub allowed: bool,
pub name: &'a str,
pub tool_name: Option<&'a str>,
pub user_id: Option<&'a str>,
pub invocation_id: Option<&'a str>,
pub reason: &'static str,
}
pub struct AuthorizingSecretService {
inner: Arc<dyn SecretService>,
grants: HashMap<String, SecretGrant>,
untooled: SecretGrant,
audit: Option<Arc<dyn SecretAuditSink>>,
}
impl AuthorizingSecretService {
pub fn new(inner: Arc<dyn SecretService>) -> Self {
Self { inner, grants: HashMap::new(), untooled: SecretGrant::none(), audit: None }
}
#[must_use]
pub fn grant(mut self, tool_name: impl Into<String>, grant: SecretGrant) -> Self {
self.grants.insert(tool_name.into(), grant);
self
}
#[must_use]
pub fn grant_untooled(mut self, grant: SecretGrant) -> Self {
self.untooled = grant;
self
}
#[must_use]
pub fn with_audit_sink(mut self, sink: Arc<dyn SecretAuditSink>) -> Self {
self.audit = Some(sink);
self
}
fn decide(&self, request: &SecretRequest) -> (bool, &'static str) {
match &request.tool_name {
Some(tool_name) => match self.grants.get(tool_name) {
Some(grant) if grant.allows(&request.name) => (true, "granted to tool"),
Some(_) => (false, "secret not in the tool's grant"),
None => (false, "no grant for tool"),
},
None => {
if self.untooled.allows(&request.name) {
(true, "granted without tool identity")
} else {
(false, "no grant for a request without tool identity")
}
}
}
}
fn record(&self, request: &SecretRequest, allowed: bool, reason: &'static str) {
let decision = SecretAccessDecision {
allowed,
name: &request.name,
tool_name: request.tool_name.as_deref(),
user_id: request.user_id.as_deref(),
invocation_id: request.invocation_id.as_deref(),
reason,
};
if allowed {
tracing::info!(
secret.name = %decision.name,
tool.name = decision.tool_name.unwrap_or("<none>"),
user.id = decision.user_id.unwrap_or("<unknown>"),
invocation.id = decision.invocation_id.unwrap_or("<unknown>"),
decision.reason = reason,
"secret access allowed"
);
} else {
tracing::warn!(
secret.name = %decision.name,
tool.name = decision.tool_name.unwrap_or("<none>"),
user.id = decision.user_id.unwrap_or("<unknown>"),
invocation.id = decision.invocation_id.unwrap_or("<unknown>"),
decision.reason = reason,
"secret access denied"
);
}
if let Some(sink) = &self.audit {
sink.record(decision);
}
}
}
impl std::fmt::Debug for AuthorizingSecretService {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AuthorizingSecretService")
.field("granted_tools", &self.grants.keys().collect::<Vec<_>>())
.field("audited", &self.audit.is_some())
.finish()
}
}
#[async_trait]
impl SecretService for AuthorizingSecretService {
async fn get_secret(&self, name: &str) -> Result<String> {
let request = SecretRequest::new(name);
self.record(&request, false, "no identity supplied");
Err(AdkError::unauthorized(
adk_core::ErrorComponent::Tool,
"secret.no_identity",
format!("secret '{name}' was requested without identity, so it cannot be authorized"),
))
}
async fn get_secret_for(&self, request: &SecretRequest) -> Result<String> {
let (allowed, reason) = self.decide(request);
self.record(request, allowed, reason);
if !allowed {
return Err(AdkError::unauthorized(
adk_core::ErrorComponent::Tool,
"secret.access_denied",
format!(
"tool {} is not permitted to read secret '{}': {reason}",
request.tool_name.as_deref().unwrap_or("<none>"),
request.name
),
));
}
self.inner.get_secret_for(request).await
}
}