1mod client;
2mod decisions;
3mod server;
4mod tools;
5
6use self::client::run_client_to_server;
7use self::server::run_server_to_client;
8use super::decision::{DecisionEmitter, FileDecisionEmitter, NullDecisionEmitter};
9use super::policy::McpPolicy;
10use super::tool_definition::ToolDefinitionBinding;
11use std::{
12 collections::HashMap,
13 io,
14 process::{Child, Command, Stdio},
15 sync::{Arc, Mutex},
16 thread,
17};
18
19#[derive(Clone, Debug)]
23pub struct ProxyConfig {
24 pub dry_run: bool,
25 pub verbose: bool,
26 pub audit_log_path: Option<std::path::PathBuf>,
28 pub server_id: String,
29 pub decision_log_path: Option<std::path::PathBuf>,
31 pub event_source: Option<String>,
33}
34
35#[derive(Clone, Debug, Default)]
37pub struct ProxyConfigRaw {
38 pub dry_run: bool,
39 pub verbose: bool,
40 pub audit_log_path: Option<std::path::PathBuf>,
41 pub server_id: String,
42 pub decision_log_path: Option<std::path::PathBuf>,
43 pub event_source: Option<String>,
44}
45
46impl ProxyConfig {
47 pub fn try_from_raw(raw: ProxyConfigRaw) -> anyhow::Result<Self> {
53 let logging_enabled = raw.audit_log_path.is_some() || raw.decision_log_path.is_some();
54
55 let event_source = raw
56 .event_source
57 .map(|s| s.trim().to_string())
58 .filter(|s| !s.is_empty());
59
60 if logging_enabled && event_source.is_none() {
61 anyhow::bail!(
62 "event_source is required when logging is enabled (e.g. --event-source assay://org/app)"
63 );
64 }
65
66 if let Some(ref src) = event_source {
67 validate_event_source(src)?;
68 }
69
70 Ok(ProxyConfig {
71 dry_run: raw.dry_run,
72 verbose: raw.verbose,
73 audit_log_path: raw.audit_log_path,
74 server_id: raw.server_id,
75 decision_log_path: raw.decision_log_path,
76 event_source,
77 })
78 }
79}
80
81fn validate_event_source(s: &str) -> anyhow::Result<()> {
83 let s = s.trim();
84 if s.is_empty() {
85 anyhow::bail!("event_source must be absolute URI with scheme (e.g. assay://org/app)");
86 }
87 if s.chars().any(|c| c.is_whitespace()) {
88 anyhow::bail!("event_source must not contain whitespace");
89 }
90
91 let Some(pos) = s.find("://") else {
93 anyhow::bail!("event_source must be absolute URI with scheme (e.g. assay://org/app)");
94 };
95 if pos == 0 {
96 anyhow::bail!("event_source must have scheme before :// (e.g. assay://org/app)");
97 }
98
99 let scheme = &s[..pos];
101 let mut chars = scheme.chars();
102 match chars.next() {
103 Some(c) if c.is_ascii_alphabetic() => {}
104 _ => anyhow::bail!("event_source URI scheme must start with a letter"),
105 }
106 if !chars.all(|c| c.is_ascii_alphanumeric() || c == '+' || c == '-' || c == '.') {
107 anyhow::bail!("event_source URI scheme contains invalid characters");
108 }
109
110 Ok(())
111}
112
113pub struct McpProxy {
114 child: Child,
115 policy: McpPolicy,
116 config: ProxyConfig,
117 identity_cache: Arc<Mutex<HashMap<String, super::identity::ToolIdentity>>>,
119 tool_definition_cache: Arc<Mutex<HashMap<String, ToolDefinitionBinding>>>,
121}
122
123impl Drop for McpProxy {
124 fn drop(&mut self) {
125 let _ = self.child.kill();
127 }
128}
129
130impl McpProxy {
131 pub fn spawn(
132 command: &str,
133 args: &[String],
134 policy: McpPolicy,
135 config: ProxyConfig,
136 ) -> io::Result<Self> {
137 let child = Command::new(command)
138 .args(args)
139 .stdin(Stdio::piped())
140 .stdout(Stdio::piped())
141 .stderr(Stdio::inherit()) .spawn()?;
143
144 Ok(Self {
145 child,
146 policy,
147 config,
148 identity_cache: Arc::new(Mutex::new(HashMap::new())),
149 tool_definition_cache: Arc::new(Mutex::new(HashMap::new())),
150 })
151 }
152
153 pub fn run(mut self) -> io::Result<i32> {
154 let child_stdin = self.child.stdin.take().expect("child stdin");
155 let child_stdout = self.child.stdout.take().expect("child stdout");
156
157 let stdout = Arc::new(Mutex::new(io::stdout()));
158 let policy = self.policy.clone();
159 let config = self.config.clone();
160 let identity_cache_a = self.identity_cache.clone();
161 let identity_cache_b = self.identity_cache.clone();
162 let tool_definition_cache_a = self.tool_definition_cache.clone();
163 let tool_definition_cache_b = self.tool_definition_cache.clone();
164
165 let decision_emitter: Arc<dyn DecisionEmitter> =
167 if let Some(path) = &config.decision_log_path {
168 Arc::new(FileDecisionEmitter::new(path)?)
169 } else {
170 Arc::new(NullDecisionEmitter)
171 };
172 let event_source = config
173 .event_source
174 .clone()
175 .unwrap_or_else(|| format!("assay://{}", config.server_id));
176
177 let stdout_a = stdout.clone();
179 let server_id_a = config.server_id.clone();
180 let t_server_to_client = thread::spawn(move || {
181 run_server_to_client(
182 child_stdout,
183 stdout_a,
184 server_id_a,
185 identity_cache_a,
186 tool_definition_cache_a,
187 )
188 });
189
190 let stdout_b = stdout.clone();
192 let emitter_b = decision_emitter.clone();
193 let event_source_b = event_source.clone();
194 let t_client_to_server = thread::spawn(move || {
195 run_client_to_server(
196 child_stdin,
197 stdout_b,
198 policy,
199 config,
200 emitter_b,
201 event_source_b,
202 identity_cache_b,
203 tool_definition_cache_b,
204 )
205 });
206
207 t_client_to_server
209 .join()
210 .map_err(|_| io::Error::other("client->server thread panicked"))??;
211
212 let _ = t_server_to_client.join();
214
215 let status = self.child.wait()?;
217 Ok(status.code().unwrap_or(1))
218 }
219}
220
221#[cfg(test)]
222mod tests {
223 use super::*;
224
225 #[test]
226 fn event_source_accepts_assay_uri() {
227 validate_event_source("assay://myorg/myapp").unwrap();
228 }
229
230 #[test]
231 fn event_source_accepts_https_uri() {
232 validate_event_source("https://example.com/agent").unwrap();
233 }
234
235 #[test]
236 fn event_source_rejects_empty() {
237 assert!(validate_event_source("").is_err());
238 assert!(validate_event_source(" ").is_err());
239 }
240
241 #[test]
242 fn event_source_rejects_whitespace() {
243 assert!(validate_event_source("assay://myorg/my app").is_err());
244 assert!(validate_event_source("assay://myorg/\tmyapp").is_err());
245 }
246
247 #[test]
248 fn event_source_rejects_missing_scheme() {
249 assert!(validate_event_source("myorg/myapp").is_err());
250 assert!(validate_event_source("://myorg/myapp").is_err());
251 }
252
253 #[test]
254 fn event_source_rejects_did_and_urn() {
255 assert!(validate_event_source("did:example:123").is_err());
257 assert!(validate_event_source("urn:example:foo").is_err());
258 }
259
260 #[test]
261 fn event_source_rejects_scheme_starting_with_non_letter() {
262 assert!(validate_event_source("1assay://myorg/myapp").is_err());
263 assert!(validate_event_source("-assay://myorg/myapp").is_err());
264 }
265
266 #[test]
267 fn event_source_rejects_scheme_with_invalid_chars() {
268 assert!(validate_event_source("as_say://myorg/myapp").is_err());
269 assert!(validate_event_source("as@say://myorg/myapp").is_err());
270 }
271
272 #[test]
273 fn config_requires_event_source_when_logging_enabled() {
274 let raw = ProxyConfigRaw {
275 dry_run: false,
276 verbose: false,
277 audit_log_path: None,
278 decision_log_path: Some(std::path::PathBuf::from("decisions.ndjson")),
279 event_source: None,
280 server_id: "srv".to_string(),
281 };
282
283 let err = ProxyConfig::try_from_raw(raw).unwrap_err();
284 let msg = format!("{err:#}");
285 assert!(msg.contains("event_source is required"));
286 }
287
288 #[test]
289 fn config_allows_no_event_source_when_logging_disabled() {
290 let raw = ProxyConfigRaw {
291 dry_run: false,
292 verbose: false,
293 audit_log_path: None,
294 decision_log_path: None,
295 event_source: None,
296 server_id: "srv".to_string(),
297 };
298
299 ProxyConfig::try_from_raw(raw).unwrap();
300 }
301
302 #[test]
303 fn config_accepts_valid_event_source() {
304 let raw = ProxyConfigRaw {
305 dry_run: false,
306 verbose: false,
307 audit_log_path: None,
308 decision_log_path: Some(std::path::PathBuf::from("decisions.ndjson")),
309 event_source: Some("assay://myorg/myapp".to_string()),
310 server_id: "srv".to_string(),
311 };
312
313 let cfg = ProxyConfig::try_from_raw(raw).unwrap();
314 assert_eq!(cfg.event_source.as_deref(), Some("assay://myorg/myapp"));
315 }
316
317 #[test]
318 fn config_rejects_invalid_event_source_uri() {
319 let raw = ProxyConfigRaw {
320 dry_run: false,
321 verbose: false,
322 audit_log_path: None,
323 decision_log_path: Some(std::path::PathBuf::from("decisions.ndjson")),
324 event_source: Some("not a uri".to_string()),
325 server_id: "srv".to_string(),
326 };
327
328 assert!(ProxyConfig::try_from_raw(raw).is_err());
329 }
330}