Skip to main content

assay_core/mcp/
proxy.rs

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/// Validated proxy configuration.
20///
21/// Use `ProxyConfig::try_from_raw()` to create from CLI/config input.
22#[derive(Clone, Debug)]
23pub struct ProxyConfig {
24    pub dry_run: bool,
25    pub verbose: bool,
26    /// NDJSON log for mandate lifecycle events (audit trail)
27    pub audit_log_path: Option<std::path::PathBuf>,
28    pub server_id: String,
29    /// NDJSON log for tool decision events (high volume)
30    pub decision_log_path: Option<std::path::PathBuf>,
31    /// CloudEvents source URI (validated, required when logging enabled)
32    pub event_source: Option<String>,
33}
34
35/// Raw config as provided by CLI/config files before validation.
36#[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    /// Create validated config from raw input.
48    ///
49    /// Fails if:
50    /// - Logging is enabled but event_source is missing
51    /// - event_source is not a valid absolute URI (scheme://...)
52    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
81/// Validate event_source URI (must be absolute with scheme://).
82fn 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    // Require scheme://...
92    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    // Validate scheme charset (RFC 3986: ALPHA *( ALPHA / DIGIT / "+" / "-" / "." ))
100    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    /// Cache of tool identities discovered during tools/list
118    identity_cache: Arc<Mutex<HashMap<String, super::identity::ToolIdentity>>>,
119    /// Cache of bounded tool-definition bindings discovered during tools/list
120    tool_definition_cache: Arc<Mutex<HashMap<String, ToolDefinitionBinding>>>,
121}
122
123impl Drop for McpProxy {
124    fn drop(&mut self) {
125        // Best-effort cleanup
126        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()) // protocol blijft op stdout
142            .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        // Initialize decision emitter (I1: always emit decision)
166        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        // Thread A: server -> client passthrough
178        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        // Thread B: client -> server passthrough with Policy Check
191        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        // Wacht tot client->server eindigt (stdin closed)
208        t_client_to_server
209            .join()
210            .map_err(|_| io::Error::other("client->server thread panicked"))??;
211
212        // Server->client thread kan nog even lopen; join best-effort
213        let _ = t_server_to_client.join();
214
215        // Wacht op child exit
216        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        // We require scheme:// not just scheme:
256        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}