1use std::collections::HashMap;
6use std::fs;
7use std::path::PathBuf;
8use std::time::Duration;
9
10use serde::{Deserialize, Serialize};
11use serde_json::{Map, Value};
12
13use crate::errors::{AcpError, AcpResult};
14use crate::http_security::{HttpSecurityPolicy, build_http_client, validate_http_url};
15use crate::identity::{parse_agent_id, verify_identity_document};
16use crate::well_known::{
17 parse_well_known_document, resolve_identity_document_reference, well_known_url_from_base,
18};
19
20#[derive(Debug, Clone, Serialize, Deserialize)]
21struct CachedDocument {
22 identity_document: Map<String, Value>,
23 fetched_at: String,
24}
25
26#[derive(Debug, Clone)]
27pub struct DiscoveryClient {
28 cache_path: Option<PathBuf>,
29 default_scheme: String,
30 relay_hints: Vec<String>,
31 enterprise_directory_hints: Vec<String>,
32 timeout_seconds: u64,
33 policy: HttpSecurityPolicy,
34 cache: HashMap<String, CachedDocument>,
35 registry: HashMap<String, Map<String, Value>>,
36 client: reqwest::blocking::Client,
37}
38
39impl DiscoveryClient {
40 #[allow(clippy::too_many_arguments)]
41 pub fn new(
42 cache_path: Option<PathBuf>,
43 default_scheme: Option<String>,
44 relay_hints: Option<Vec<String>>,
45 enterprise_directory_hints: Option<Vec<String>>,
46 timeout_seconds: u64,
47 allow_insecure_http: bool,
48 allow_insecure_tls: bool,
49 ca_file: Option<String>,
50 mtls_enabled: bool,
51 cert_file: Option<String>,
52 key_file: Option<String>,
53 ) -> AcpResult<Self> {
54 let policy = HttpSecurityPolicy {
55 allow_insecure_http,
56 allow_insecure_tls,
57 mtls_enabled,
58 ca_file,
59 cert_file,
60 key_file,
61 };
62 let client = build_http_client(timeout_seconds.max(1), &policy)?;
63 let mut instance = Self {
64 cache_path,
65 default_scheme: default_scheme.unwrap_or_else(|| "https".to_string()),
66 relay_hints: relay_hints.unwrap_or_default(),
67 enterprise_directory_hints: enterprise_directory_hints.unwrap_or_default(),
68 timeout_seconds: timeout_seconds.max(1),
69 policy,
70 cache: HashMap::new(),
71 registry: HashMap::new(),
72 client,
73 };
74 instance.load_cache();
75 Ok(instance)
76 }
77
78 pub fn seed(&mut self, identity_document: Map<String, Value>) -> AcpResult<()> {
79 if let Some(agent_id) = identity_document.get("agent_id").and_then(Value::as_str) {
80 self.cache.insert(
81 agent_id.to_string(),
82 CachedDocument {
83 identity_document,
84 fetched_at: chrono::Utc::now().to_rfc3339(),
85 },
86 );
87 self.persist_cache()?;
88 }
89 Ok(())
90 }
91
92 pub fn register_identity_document(
93 &mut self,
94 identity_document: Map<String, Value>,
95 ) -> AcpResult<()> {
96 let agent_id = identity_document
97 .get("agent_id")
98 .and_then(Value::as_str)
99 .map(str::to_string)
100 .ok_or_else(|| {
101 AcpError::Validation("Identity document missing agent_id".to_string())
102 })?;
103 self.registry
104 .insert(agent_id.clone(), identity_document.clone());
105 self.cache.insert(
106 agent_id,
107 CachedDocument {
108 identity_document,
109 fetched_at: chrono::Utc::now().to_rfc3339(),
110 },
111 );
112 self.persist_cache()
113 }
114
115 pub fn resolve(&mut self, agent_id: &str) -> AcpResult<Map<String, Value>> {
116 if let Some(registry_doc) = self.registry.get(agent_id) {
117 return Ok(registry_doc.clone());
118 }
119 if let Some(cached) = self.try_cache(agent_id)? {
120 return Ok(cached);
121 }
122 if let Some(well_known) = self.try_well_known(agent_id)? {
123 self.cache_identity(agent_id, well_known.clone())?;
124 return Ok(well_known);
125 }
126 if let Some(relay_doc) = self.try_hint_lookups(&self.relay_hints, agent_id)? {
127 self.cache_identity(agent_id, relay_doc.clone())?;
128 return Ok(relay_doc);
129 }
130 if let Some(enterprise_doc) =
131 self.try_hint_lookups(&self.enterprise_directory_hints, agent_id)?
132 {
133 self.cache_identity(agent_id, enterprise_doc.clone())?;
134 return Ok(enterprise_doc);
135 }
136 Err(AcpError::Discovery(format!(
137 "Unable to resolve identity document for {agent_id}"
138 )))
139 }
140
141 pub fn resolve_well_known(
142 &mut self,
143 base_url: &str,
144 expected_agent_id: Option<&str>,
145 ) -> AcpResult<Map<String, Value>> {
146 let well_known_url = well_known_url_from_base(base_url)?;
147 let resolved = self
148 .resolve_well_known_url(&well_known_url, expected_agent_id)?
149 .ok_or_else(|| {
150 AcpError::Discovery(format!(
151 "Unable to resolve well-known metadata from {well_known_url}"
152 ))
153 })?;
154 let identity_document = resolved
155 .get("identity_document")
156 .and_then(Value::as_object)
157 .cloned()
158 .ok_or_else(|| {
159 AcpError::Discovery(
160 "Well-known discovery returned invalid identity document".to_string(),
161 )
162 })?;
163 let agent_id = identity_document
164 .get("agent_id")
165 .and_then(Value::as_str)
166 .map(str::to_string)
167 .ok_or_else(|| {
168 AcpError::Discovery(
169 "Well-known discovery returned identity document without agent_id".to_string(),
170 )
171 })?;
172 self.cache_identity(&agent_id, identity_document)?;
173 let mut response = resolved;
174 response.insert("well_known_url".to_string(), Value::String(well_known_url));
175 Ok(response)
176 }
177
178 fn try_cache(&mut self, agent_id: &str) -> AcpResult<Option<Map<String, Value>>> {
179 let Some(cached) = self.cache.get(agent_id).cloned() else {
180 return Ok(None);
181 };
182 if cache_valid(&cached.identity_document) {
183 return Ok(Some(cached.identity_document));
184 }
185 self.cache.remove(agent_id);
186 self.persist_cache()?;
187 Ok(None)
188 }
189
190 fn try_well_known(&mut self, agent_id: &str) -> AcpResult<Option<Map<String, Value>>> {
191 let parts = parse_agent_id(agent_id)?;
192 let Some(domain) = parts.domain else {
193 return Ok(None);
194 };
195 let well_known_url = format!("{}://{domain}/.well-known/acp", self.default_scheme);
196 let Some(resolved) = self.resolve_well_known_url(&well_known_url, Some(agent_id))? else {
197 return Ok(None);
198 };
199 let identity_document = resolved
200 .get("identity_document")
201 .and_then(Value::as_object)
202 .cloned();
203 Ok(identity_document)
204 }
205
206 fn try_hint_lookups(
207 &self,
208 hints: &[String],
209 agent_id: &str,
210 ) -> AcpResult<Option<Map<String, Value>>> {
211 for hint in hints {
212 let url = format!("{}/discover", hint.trim_end_matches('/'));
213 let body = self.fetch_json(
214 &url,
215 Some(&[("agent_id", agent_id)]),
216 "Discovery hint lookup",
217 )?;
218 let Some(body) = body else {
219 continue;
220 };
221 if let Some(identity_document) = extract_identity_document(&body)
222 && validate_identity_document(&identity_document)
223 {
224 return Ok(Some(identity_document));
225 }
226 }
227 Ok(None)
228 }
229
230 fn resolve_well_known_url(
231 &self,
232 well_known_url: &str,
233 expected_agent_id: Option<&str>,
234 ) -> AcpResult<Option<Map<String, Value>>> {
235 let body = self.fetch_json(well_known_url, None, "Discovery .well-known lookup")?;
236 let Some(body) = body else {
237 return Ok(None);
238 };
239 let well_known = match parse_well_known_document(&Value::Object(body.clone())) {
240 Ok(value) => value,
241 Err(_) => return Ok(None),
242 };
243 if let Some(expected) = expected_agent_id
244 && well_known.get("agent_id").and_then(Value::as_str) != Some(expected)
245 {
246 return Ok(None);
247 }
248 let identity_reference = resolve_identity_document_reference(&well_known, well_known_url)?;
249 let identity_body = self.fetch_json(
250 &identity_reference,
251 None,
252 "Discovery identity document lookup",
253 )?;
254 let Some(identity_body) = identity_body else {
255 return Ok(None);
256 };
257 let Some(identity_document) = extract_identity_document(&identity_body) else {
258 return Ok(None);
259 };
260 if !validate_identity_document(&identity_document) {
261 return Ok(None);
262 }
263 if let Some(expected) = expected_agent_id
264 && identity_document.get("agent_id").and_then(Value::as_str) != Some(expected)
265 {
266 return Ok(None);
267 }
268 let mut response = Map::new();
269 response.insert("well_known".to_string(), Value::Object(well_known));
270 response.insert(
271 "identity_document".to_string(),
272 Value::Object(identity_document),
273 );
274 Ok(Some(response))
275 }
276
277 fn fetch_json(
278 &self,
279 url: &str,
280 query: Option<&[(&str, &str)]>,
281 context: &str,
282 ) -> AcpResult<Option<Map<String, Value>>> {
283 validate_http_url(
284 url,
285 self.policy.allow_insecure_http,
286 self.policy.mtls_enabled,
287 context,
288 )?;
289 let mut request = self
290 .client
291 .get(url)
292 .timeout(Duration::from_secs(self.timeout_seconds));
293 if let Some(query) = query {
294 request = request.query(query);
295 }
296 let response = match request.send() {
297 Ok(response) => response,
298 Err(_) => return Ok(None),
299 };
300 if response.status().as_u16() != 200 {
301 return Ok(None);
302 }
303 let value: Value = match response.json() {
304 Ok(value) => value,
305 Err(_) => return Ok(None),
306 };
307 Ok(value.as_object().cloned())
308 }
309
310 fn cache_identity(
311 &mut self,
312 agent_id: &str,
313 identity_document: Map<String, Value>,
314 ) -> AcpResult<()> {
315 self.cache.insert(
316 agent_id.to_string(),
317 CachedDocument {
318 identity_document,
319 fetched_at: chrono::Utc::now().to_rfc3339(),
320 },
321 );
322 self.persist_cache()
323 }
324
325 fn load_cache(&mut self) {
326 let Some(path) = &self.cache_path else {
327 return;
328 };
329 let Ok(raw) = fs::read_to_string(path) else {
330 return;
331 };
332 if let Ok(value) = serde_json::from_str::<HashMap<String, CachedDocument>>(&raw) {
333 self.cache = value;
334 }
335 }
336
337 fn persist_cache(&self) -> AcpResult<()> {
338 let Some(path) = &self.cache_path else {
339 return Ok(());
340 };
341 if let Some(parent) = path.parent() {
342 fs::create_dir_all(parent)?;
343 }
344 let value = serde_json::to_string(&self.cache)?;
345 fs::write(path, value)?;
346 Ok(())
347 }
348}
349
350fn cache_valid(identity_document: &Map<String, Value>) -> bool {
351 let Some(valid_until) = identity_document.get("valid_until").and_then(Value::as_str) else {
352 return false;
353 };
354 chrono::DateTime::parse_from_rfc3339(valid_until)
355 .map(|ts| ts > chrono::Utc::now())
356 .unwrap_or(false)
357}
358
359fn validate_identity_document(identity_document: &Map<String, Value>) -> bool {
360 verify_identity_document(identity_document) && cache_valid(identity_document)
361}
362
363fn extract_identity_document(body: &Map<String, Value>) -> Option<Map<String, Value>> {
364 if let Some(identity) = body.get("identity_document").and_then(Value::as_object) {
365 return Some(identity.clone());
366 }
367 if body.get("agent_id").is_some() {
368 return Some(body.clone());
369 }
370 None
371}