use crate::crypto::TOKEN_SIGNATURE_ALG_ED25519;
use crate::models::{AgentIdentityDocument, Violation};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use std::collections::HashMap;
pub const DISCOVERY_ISSUER_PATH: &str = "/.well-known/delegated-issuer";
pub const DISCOVERY_JWKS_PATH: &str = "/.well-known/jwks.json";
pub const DISCOVERY_REGISTRY_PREFIX: &str = "/registry/";
pub const DISCOVERY_RESOLVE_PREFIX: &str = "/resolve/";
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct IssuerMetadata {
pub issuer: String,
pub jwks_uri: String,
pub registry_uri: String,
pub resolution_uri: String,
pub revocation_uri: String,
pub approval_uri: String,
pub supported_signature_algorithms: Vec<String>,
pub supported_protocols: Vec<String>,
pub supported_profiles: Vec<String>,
pub spec_version: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct JwksDocument {
pub keys: Vec<JwkRecord>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct JwkRecord {
pub kid: String,
pub kty: String,
pub crv: String,
pub x: String,
#[serde(rename = "use")]
pub use_field: String,
pub alg: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct DiscoveryHttpRequest {
pub method: String,
pub path: String,
#[serde(default)]
pub query: HashMap<String, String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct DiscoveryHttpResponse {
pub status_code: u16,
pub body: Value,
}
#[derive(Debug, Clone, Default)]
pub struct DiscoveryService {
issuer_metadata: HashMap<String, IssuerMetadata>,
registry: HashMap<String, AgentIdentityDocument>,
}
impl DiscoveryService {
pub fn new() -> Self {
Self {
issuer_metadata: HashMap::new(),
registry: HashMap::new(),
}
}
pub fn register_issuer_metadata(&mut self, metadata: IssuerMetadata) {
self.issuer_metadata
.insert(metadata.issuer.clone(), metadata);
}
pub fn register_identity_document(&mut self, document: AgentIdentityDocument) {
self.registry.insert(document.agent_id.clone(), document);
}
pub fn get_issuer_metadata(&self, issuer: &str) -> Option<&IssuerMetadata> {
self.issuer_metadata.get(issuer)
}
pub fn get_identity_document(&self, agent_id: &str) -> Option<&AgentIdentityDocument> {
self.registry.get(agent_id)
}
pub fn resolve_agent_endpoint(&self, agent_id: &str, protocol: &str) -> Option<String> {
self.registry.get(agent_id).and_then(|document| {
document
.endpoints
.iter()
.find(|endpoint| endpoint.protocol == protocol)
.map(|endpoint| endpoint.url.clone())
})
}
}
pub fn build_jwks_document(
identity_document: &AgentIdentityDocument,
) -> Result<JwksDocument, Violation> {
let mut keys = Vec::with_capacity(identity_document.public_keys.len());
for key in &identity_document.public_keys {
let crv = key.crv.clone().ok_or_else(|| {
Violation::new(
"build_jwks_document",
"identity public key must include crv for JWKS conversion",
)
})?;
let x = key.x.clone().ok_or_else(|| {
Violation::new(
"build_jwks_document",
"identity public key must include x for JWKS conversion",
)
})?;
if key.kty != "OKP" || crv != TOKEN_SIGNATURE_ALG_ED25519 {
return Err(Violation::new(
"build_jwks_document",
"only OKP/Ed25519 keys are supported in JWKS conversion",
));
}
keys.push(JwkRecord {
kid: key.kid.clone(),
kty: key.kty.clone(),
crv,
x,
use_field: "sig".to_string(),
alg: "EdDSA".to_string(),
});
}
Ok(JwksDocument { keys })
}
pub fn handle_discovery_http_request(
service: &DiscoveryService,
request: &DiscoveryHttpRequest,
) -> DiscoveryHttpResponse {
if request.method != "GET" {
return discovery_error(
405,
"discovery handlers only support GET requests".to_string(),
);
}
if request.path == DISCOVERY_ISSUER_PATH {
return serve_issuer_metadata(service, &request.query);
}
if request.path == DISCOVERY_JWKS_PATH {
return serve_jwks(service, &request.query);
}
if let Some(agent_id) = request.path.strip_prefix(DISCOVERY_REGISTRY_PREFIX) {
return serve_registry_lookup(service, agent_id);
}
if let Some(agent_id) = request.path.strip_prefix(DISCOVERY_RESOLVE_PREFIX) {
return serve_resolution_lookup(service, agent_id, &request.query);
}
discovery_error(404, format!("unknown discovery path: {}", request.path))
}
fn serve_issuer_metadata(
service: &DiscoveryService,
query: &HashMap<String, String>,
) -> DiscoveryHttpResponse {
let issuer = match query.get("issuer") {
Some(issuer) if !issuer.trim().is_empty() => issuer.clone(),
Some(_) => {
return discovery_error(400, "query parameter issuer must be non-empty".to_string());
}
None => {
if service.issuer_metadata.len() == 1 {
service
.issuer_metadata
.keys()
.next()
.cloned()
.expect("issuer metadata length checked to be one")
} else {
return discovery_error(
400,
"query parameter issuer is required when multiple issuers are registered"
.to_string(),
);
}
}
};
match service.get_issuer_metadata(&issuer) {
Some(metadata) => discovery_ok(
200,
serde_json::to_value(metadata).expect("serialize metadata"),
),
None => discovery_error(
404,
format!("issuer metadata not found for issuer: {issuer}"),
),
}
}
fn serve_jwks(
service: &DiscoveryService,
query: &HashMap<String, String>,
) -> DiscoveryHttpResponse {
let agent_id = match query.get("agent_id") {
Some(agent_id) if !agent_id.trim().is_empty() => agent_id,
Some(_) => {
return discovery_error(
400,
"query parameter agent_id must be non-empty".to_string(),
);
}
None => {
return discovery_error(400, "query parameter agent_id is required".to_string());
}
};
let identity_document = match service.get_identity_document(agent_id) {
Some(document) => document,
None => {
return discovery_error(
404,
format!("identity document not found for agent_id: {agent_id}"),
);
}
};
match build_jwks_document(identity_document) {
Ok(jwks) => discovery_ok(200, serde_json::to_value(jwks).expect("serialize jwks")),
Err(error) => discovery_error(422, error.reason),
}
}
fn serve_registry_lookup(service: &DiscoveryService, agent_id: &str) -> DiscoveryHttpResponse {
if agent_id.trim().is_empty() {
return discovery_error(400, "registry path must include agent_id".to_string());
}
match service.get_identity_document(agent_id) {
Some(document) => discovery_ok(
200,
serde_json::to_value(document).expect("serialize identity document"),
),
None => discovery_error(
404,
format!("identity document not found for agent_id: {agent_id}"),
),
}
}
fn serve_resolution_lookup(
service: &DiscoveryService,
agent_id: &str,
query: &HashMap<String, String>,
) -> DiscoveryHttpResponse {
if agent_id.trim().is_empty() {
return discovery_error(400, "resolution path must include agent_id".to_string());
}
let protocol = match query.get("protocol") {
Some(protocol) if !protocol.trim().is_empty() => protocol,
Some(_) => {
return discovery_error(
400,
"query parameter protocol must be non-empty".to_string(),
);
}
None => {
return discovery_error(400, "query parameter protocol is required".to_string());
}
};
match service.resolve_agent_endpoint(agent_id, protocol) {
Some(endpoint) => discovery_ok(
200,
json!({
"agent_id": agent_id,
"protocol": protocol,
"endpoint": endpoint
}),
),
None => discovery_error(
404,
format!("endpoint not found for agent_id {agent_id} and protocol {protocol}"),
),
}
}
fn discovery_ok(status_code: u16, body: Value) -> DiscoveryHttpResponse {
DiscoveryHttpResponse { status_code, body }
}
fn discovery_error(status_code: u16, message: String) -> DiscoveryHttpResponse {
DiscoveryHttpResponse {
status_code,
body: json!({ "error": message }),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::models::{AgentEndpoint, AgentIdentityDocument, PublicKeyRecord};
use chrono::{TimeZone, Utc};
fn identity_document() -> AgentIdentityDocument {
AgentIdentityDocument {
spec_version: "0.1".to_string(),
kind: "AgentIdentityDocument".to_string(),
agent_id: "agent:example:scheduler:v1".to_string(),
display_name: Some("Scheduler".to_string()),
owner_id: "org:example".to_string(),
issuer: "https://trust.example.ai".to_string(),
identity_type: "spiffe".to_string(),
subject: "spiffe://example.ai/agents/scheduler".to_string(),
public_keys: vec![PublicKeyRecord {
kid: "key-2026-01".to_string(),
kty: "OKP".to_string(),
crv: Some(TOKEN_SIGNATURE_ALG_ED25519.to_string()),
x: Some("abc".to_string()),
}],
supported_protocols: vec!["http".to_string(), "mcp".to_string()],
supported_auth_methods: vec!["delegation_token".to_string()],
capabilities: None,
endpoints: vec![AgentEndpoint {
protocol: "mcp".to_string(),
url: "https://agents.example.ai/scheduler/mcp".to_string(),
}],
attestation: None,
created_at: Utc
.with_ymd_and_hms(2026, 6, 1, 0, 0, 0)
.single()
.expect("valid timestamp"),
expires_at: Utc
.with_ymd_and_hms(2026, 6, 8, 0, 0, 0)
.single()
.expect("valid timestamp"),
signature: "sig".to_string(),
}
}
fn issuer_metadata() -> IssuerMetadata {
IssuerMetadata {
issuer: "https://trust.example.ai".to_string(),
jwks_uri: "https://trust.example.ai/.well-known/jwks.json".to_string(),
registry_uri: "https://trust.example.ai/registry".to_string(),
resolution_uri: "https://trust.example.ai/resolve".to_string(),
revocation_uri: "https://trust.example.ai/revoke".to_string(),
approval_uri: "https://trust.example.ai/approve".to_string(),
supported_signature_algorithms: vec!["Ed25519".to_string()],
supported_protocols: vec!["http".to_string(), "mcp".to_string()],
supported_profiles: vec!["developer".to_string(), "spiffe".to_string()],
spec_version: "0.1".to_string(),
}
}
fn seeded_discovery_service() -> DiscoveryService {
let mut service = DiscoveryService::new();
service.register_issuer_metadata(issuer_metadata());
service.register_identity_document(identity_document());
service
}
#[test]
fn builds_rfc_style_jwks_document() {
let jwks = build_jwks_document(&identity_document()).expect("JWKS conversion should work");
assert_eq!(jwks.keys.len(), 1);
assert_eq!(jwks.keys[0].use_field, "sig");
assert_eq!(jwks.keys[0].alg, "EdDSA");
}
#[test]
fn resolves_registered_agent_endpoint() {
let service = seeded_discovery_service();
assert_eq!(
service.resolve_agent_endpoint("agent:example:scheduler:v1", "mcp"),
Some("https://agents.example.ai/scheduler/mcp".to_string())
);
assert!(
service
.get_issuer_metadata("https://trust.example.ai")
.is_some()
);
}
#[test]
fn serves_canonical_discovery_paths() {
let service = seeded_discovery_service();
let issuer_response = handle_discovery_http_request(
&service,
&DiscoveryHttpRequest {
method: "GET".to_string(),
path: DISCOVERY_ISSUER_PATH.to_string(),
query: HashMap::new(),
},
);
assert_eq!(issuer_response.status_code, 200);
assert_eq!(
issuer_response.body["issuer"],
json!("https://trust.example.ai")
);
let mut jwks_query = HashMap::new();
jwks_query.insert(
"agent_id".to_string(),
"agent:example:scheduler:v1".to_string(),
);
let jwks_response = handle_discovery_http_request(
&service,
&DiscoveryHttpRequest {
method: "GET".to_string(),
path: DISCOVERY_JWKS_PATH.to_string(),
query: jwks_query,
},
);
assert_eq!(jwks_response.status_code, 200);
assert_eq!(jwks_response.body["keys"][0]["alg"], json!("EdDSA"));
let registry_response = handle_discovery_http_request(
&service,
&DiscoveryHttpRequest {
method: "GET".to_string(),
path: "/registry/agent:example:scheduler:v1".to_string(),
query: HashMap::new(),
},
);
assert_eq!(registry_response.status_code, 200);
assert_eq!(
registry_response.body["agent_id"],
json!("agent:example:scheduler:v1")
);
let mut resolve_query = HashMap::new();
resolve_query.insert("protocol".to_string(), "mcp".to_string());
let resolve_response = handle_discovery_http_request(
&service,
&DiscoveryHttpRequest {
method: "GET".to_string(),
path: "/resolve/agent:example:scheduler:v1".to_string(),
query: resolve_query,
},
);
assert_eq!(resolve_response.status_code, 200);
assert_eq!(
resolve_response.body["endpoint"],
json!("https://agents.example.ai/scheduler/mcp")
);
}
#[test]
fn rejects_invalid_discovery_requests() {
let service = seeded_discovery_service();
let wrong_method = handle_discovery_http_request(
&service,
&DiscoveryHttpRequest {
method: "POST".to_string(),
path: DISCOVERY_ISSUER_PATH.to_string(),
query: HashMap::new(),
},
);
assert_eq!(wrong_method.status_code, 405);
let missing_jwks_query = handle_discovery_http_request(
&service,
&DiscoveryHttpRequest {
method: "GET".to_string(),
path: DISCOVERY_JWKS_PATH.to_string(),
query: HashMap::new(),
},
);
assert_eq!(missing_jwks_query.status_code, 400);
let not_found = handle_discovery_http_request(
&service,
&DiscoveryHttpRequest {
method: "GET".to_string(),
path: "/resolve/agent:example:scheduler:v1".to_string(),
query: HashMap::new(),
},
);
assert_eq!(not_found.status_code, 400);
}
}