1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
4pub struct Probe {
5 pub standard: String,
6 pub path: String,
7 pub accept: String,
8 pub content_type_prefix: String,
9}
10
11#[derive(Debug, Clone, PartialEq, Eq)]
12pub struct ProbeRequest {
13 pub probe: Probe,
14 pub url: String,
15}
16
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct ProbeMatch {
19 pub probe: Probe,
20 pub url: String,
21 pub body: String,
22}
23
24pub struct ProbeRegistry {
25 probes: Vec<Probe>,
26}
27
28impl ProbeRegistry {
29 #[must_use]
30 pub fn new() -> Self {
31 Self { probes: Vec::new() }
32 }
33
34 pub fn register(&mut self, probe: Probe) {
35 self.probes.push(probe);
36 }
37
38 #[must_use]
39 pub fn probes_for_origin(&self, origin: &str) -> Vec<ProbeRequest> {
40 self.probes
41 .iter()
42 .map(|p| {
43 let url = format!("{}{}", origin.trim_end_matches('/'), p.path);
44 ProbeRequest {
45 probe: p.clone(),
46 url,
47 }
48 })
49 .collect()
50 }
51
52 #[must_use]
53 pub fn check_response(
54 &self,
55 standard: &str,
56 status: u16,
57 content_type: &str,
58 body: String,
59 url: String,
60 ) -> Option<ProbeMatch> {
61 if status != 200 {
62 return None;
63 }
64 let probe = self.probes.iter().find(|p| p.standard == standard)?;
65 if is_html_content_type(content_type) {
66 return None;
67 }
68 let ct_lower = content_type.to_ascii_lowercase();
69 let prefix_lower = probe.content_type_prefix.to_ascii_lowercase();
70 if !ct_lower.starts_with(&prefix_lower) {
71 return None;
72 }
73 Some(ProbeMatch {
74 probe: probe.clone(),
75 url,
76 body,
77 })
78 }
79}
80
81impl Default for ProbeRegistry {
82 fn default() -> Self {
83 Self::new()
84 }
85}
86
87fn is_html_content_type(ct: &str) -> bool {
88 let lower = ct.to_ascii_lowercase();
89 let base = lower.split(';').next().unwrap_or("").trim();
90 base == "text/html" || base == "application/xhtml+xml"
91}
92
93#[cfg(test)]
94mod tests {
95 use super::*;
96
97 fn registry() -> ProbeRegistry {
98 let mut r = ProbeRegistry::new();
99 r.register(Probe {
100 standard: "ard".into(),
101 path: "/.well-known/ard.json".into(),
102 accept: "application/json".into(),
103 content_type_prefix: "application/json".into(),
104 });
105 r.register(Probe {
106 standard: "ai-catalog".into(),
107 path: "/.well-known/ai-catalog.json".into(),
108 accept: "application/json".into(),
109 content_type_prefix: "application/json".into(),
110 });
111 r
112 }
113
114 #[test]
115 fn probes_for_origin_builds_urls() {
116 let r = registry();
117 let probes = r.probes_for_origin("https://example.com");
118 assert_eq!(probes.len(), 2);
119 assert_eq!(probes[0].url, "https://example.com/.well-known/ard.json");
120 assert_eq!(probes[0].probe.accept, "application/json");
121 assert_eq!(probes[0].probe.standard, "ard");
122 }
123
124 #[test]
125 fn trailing_slash_origin() {
126 let r = registry();
127 let probes = r.probes_for_origin("https://example.com/");
128 assert_eq!(probes[0].url, "https://example.com/.well-known/ard.json");
129 }
130
131 #[test]
132 fn probes_expose_accept_header() {
133 let r = registry();
134 let probes = r.probes_for_origin("https://example.com");
135 for req in &probes {
136 assert!(!req.probe.accept.is_empty());
137 }
138 }
139
140 #[test]
141 fn exact_200_json_matches() {
142 let r = registry();
143 let m = r.check_response("ard", 200, "application/json", "{}".into(), "https://example.com/.well-known/ard.json".into());
144 assert!(m.is_some());
145 let m = m.unwrap();
146 assert_eq!(m.probe.standard, "ard");
147 assert_eq!(m.body, "{}");
148 }
149
150 #[test]
151 fn json_charset_matches() {
152 let r = registry();
153 let m = r.check_response("ard", 200, "application/json; charset=utf-8", "{}".into(), "u".into());
154 assert!(m.is_some());
155 }
156
157 #[test]
158 fn redirect_301_does_not_match() {
159 let r = registry();
160 let m = r.check_response("ard", 301, "application/json", "{}".into(), "u".into());
161 assert!(m.is_none());
162 }
163
164 #[test]
165 fn status_204_does_not_match() {
166 let r = registry();
167 let m = r.check_response("ard", 204, "application/json", "".into(), "u".into());
168 assert!(m.is_none());
169 }
170
171 #[test]
172 fn status_206_does_not_match() {
173 let r = registry();
174 let m = r.check_response("ard", 206, "application/json", "{}".into(), "u".into());
175 assert!(m.is_none());
176 }
177
178 #[test]
179 fn html_content_type_does_not_match() {
180 let r = registry();
181 let m = r.check_response("ard", 200, "text/html", "<html>".into(), "u".into());
182 assert!(m.is_none());
183 }
184
185 #[test]
186 fn xhtml_does_not_match() {
187 let r = registry();
188 let m = r.check_response("ard", 200, "application/xhtml+xml", "<html>".into(), "u".into());
189 assert!(m.is_none());
190 }
191
192 #[test]
193 fn wrong_content_type_does_not_match() {
194 let r = registry();
195 let m = r.check_response("ard", 200, "text/plain", "hello".into(), "u".into());
196 assert!(m.is_none());
197 }
198
199 #[test]
200 fn unknown_standard_returns_none() {
201 let r = registry();
202 let m = r.check_response("unknown", 200, "application/json", "{}".into(), "u".into());
203 assert!(m.is_none());
204 }
205}