a3s_code_core/security/
mod.rs

1//! Security Module
2//!
3//! Provides a trait-based security interface for A3S Code sessions.
4//! External consumers implement `SecurityProvider` to plug in their own
5//! security logic (sanitization, taint tracking, injection detection, etc.).
6
7pub mod config;
8pub mod default;
9mod event_sanitizer;
10mod value;
11
12pub use config::{RedactionStrategy, SecurityConfig, SensitivityLevel};
13pub use default::{DefaultSecurityConfig, DefaultSecurityProvider, SensitivePattern};
14pub(crate) use event_sanitizer::AgentEventStreamSanitizer;
15pub use value::{
16    sanitize_tainted_json, sanitize_tainted_text, sanitize_text, SanitizationState, SecurityLabel,
17    TaintLabel, TaintedValue, TrustLevel,
18};
19
20use crate::hooks::HookEngine;
21
22/// Sanitize string fields carried by a structured tool error while preserving
23/// its machine-readable discriminant and numeric retry metadata.
24pub(crate) fn sanitize_tool_error_kind(
25    provider: &dyn SecurityProvider,
26    kind: &crate::tools::ToolErrorKind,
27) -> crate::tools::ToolErrorKind {
28    use crate::tools::ToolErrorKind;
29
30    let text = |value: &str| sanitize_text(provider, value);
31    match kind {
32        ToolErrorKind::VersionConflict {
33            path,
34            expected,
35            actual,
36        } => ToolErrorKind::VersionConflict {
37            path: text(path),
38            expected: text(expected),
39            actual: actual.as_deref().map(text),
40        },
41        ToolErrorKind::RemoteGitConflict { code, message } => ToolErrorKind::RemoteGitConflict {
42            code: text(code),
43            message: text(message),
44        },
45        ToolErrorKind::NotFound { path } => ToolErrorKind::NotFound { path: text(path) },
46        ToolErrorKind::InvalidArgument { message } => ToolErrorKind::InvalidArgument {
47            message: text(message),
48        },
49        ToolErrorKind::HookDenied {
50            reason,
51            retryable,
52            retry_after_ms,
53        } => ToolErrorKind::HookDenied {
54            reason: text(reason),
55            retryable: *retryable,
56            retry_after_ms: *retry_after_ms,
57        },
58        ToolErrorKind::Unsupported { message } => ToolErrorKind::Unsupported {
59            message: text(message),
60        },
61        ToolErrorKind::Timeout { op, duration_ms } => ToolErrorKind::Timeout {
62            op: text(op),
63            duration_ms: *duration_ms,
64        },
65        ToolErrorKind::Transport { op } => ToolErrorKind::Transport { op: text(op) },
66        ToolErrorKind::Cancelled { op } => ToolErrorKind::Cancelled { op: text(op) },
67        ToolErrorKind::PartialFailure { failed, total } => ToolErrorKind::PartialFailure {
68            failed: *failed,
69            total: *total,
70        },
71        ToolErrorKind::RateLimited { retry_after_ms } => ToolErrorKind::RateLimited {
72            retry_after_ms: *retry_after_ms,
73        },
74    }
75}
76
77/// Sanitize every data-bearing field of an agent event while preserving the
78/// identifiers and discriminants used to correlate the event stream.
79///
80/// Runtime boundaries call this immediately before events are observed,
81/// persisted, or exposed to an SDK. Implementations of [`SecurityProvider`]
82/// should therefore make [`SecurityProvider::sanitize_output`] idempotent.
83pub fn sanitize_agent_event(
84    provider: &dyn SecurityProvider,
85    event: &crate::agent::AgentEvent,
86) -> crate::agent::AgentEvent {
87    use crate::agent::AgentEvent;
88
89    fn text(provider: &dyn SecurityProvider, value: &str) -> String {
90        sanitize_text(provider, value)
91    }
92
93    fn optional_text(provider: &dyn SecurityProvider, value: &Option<String>) -> Option<String> {
94        value.as_deref().map(|value| text(provider, value))
95    }
96
97    fn strings(provider: &dyn SecurityProvider, values: &[String]) -> Vec<String> {
98        values.iter().map(|value| text(provider, value)).collect()
99    }
100
101    fn json(provider: &dyn SecurityProvider, value: &serde_json::Value) -> serde_json::Value {
102        sanitize_tainted_json(provider, TaintedValue::untrusted(value.clone()))
103            .into_parts()
104            .0
105    }
106
107    fn optional_json(
108        provider: &dyn SecurityProvider,
109        value: &Option<serde_json::Value>,
110    ) -> Option<serde_json::Value> {
111        value.as_ref().map(|value| json(provider, value))
112    }
113
114    fn task(
115        provider: &dyn SecurityProvider,
116        task: &crate::planning::Task,
117    ) -> crate::planning::Task {
118        let mut task = task.clone();
119        task.content = text(provider, &task.content);
120        task.success_criteria = optional_text(provider, &task.success_criteria);
121        task
122    }
123
124    fn plan(
125        provider: &dyn SecurityProvider,
126        plan: &crate::planning::ExecutionPlan,
127    ) -> crate::planning::ExecutionPlan {
128        let mut plan = plan.clone();
129        plan.goal = text(provider, &plan.goal);
130        plan.steps = plan.steps.iter().map(|item| task(provider, item)).collect();
131        plan
132    }
133
134    fn goal(
135        provider: &dyn SecurityProvider,
136        goal: &crate::planning::AgentGoal,
137    ) -> crate::planning::AgentGoal {
138        let mut goal = goal.clone();
139        goal.description = text(provider, &goal.description);
140        goal.success_criteria = strings(provider, &goal.success_criteria);
141        goal
142    }
143
144    fn verification_summary(
145        provider: &dyn SecurityProvider,
146        summary: &crate::verification::VerificationSummary,
147    ) -> crate::verification::VerificationSummary {
148        let mut summary = summary.clone();
149        summary.pending_subjects = strings(provider, &summary.pending_subjects);
150        summary.failed_subjects = strings(provider, &summary.failed_subjects);
151        summary
152    }
153
154    fn response_meta(
155        provider: &dyn SecurityProvider,
156        meta: &Option<crate::llm::LlmResponseMeta>,
157    ) -> Option<crate::llm::LlmResponseMeta> {
158        meta.as_ref().map(|meta| {
159            let mut meta = meta.clone();
160            meta.request_url = optional_text(provider, &meta.request_url);
161            meta
162        })
163    }
164
165    match event {
166        AgentEvent::Start { prompt } => AgentEvent::Start {
167            prompt: text(provider, prompt),
168        },
169        AgentEvent::AgentModeChanged {
170            mode,
171            agent,
172            description,
173        } => AgentEvent::AgentModeChanged {
174            mode: mode.clone(),
175            agent: agent.clone(),
176            description: text(provider, description),
177        },
178        AgentEvent::TextDelta { text: value } => AgentEvent::TextDelta {
179            text: text(provider, value),
180        },
181        AgentEvent::ReasoningDelta { text: value } => AgentEvent::ReasoningDelta {
182            text: text(provider, value),
183        },
184        AgentEvent::ToolInputDelta { id, delta } => AgentEvent::ToolInputDelta {
185            id: id.clone(),
186            delta: text(provider, delta),
187        },
188        AgentEvent::ToolExecutionStart { id, name, args } => AgentEvent::ToolExecutionStart {
189            id: id.clone(),
190            name: name.clone(),
191            args: json(provider, args),
192        },
193        AgentEvent::ToolEnd {
194            id,
195            name,
196            args,
197            output,
198            exit_code,
199            metadata,
200            error_kind,
201        } => AgentEvent::ToolEnd {
202            id: id.clone(),
203            name: name.clone(),
204            args: optional_json(provider, args),
205            output: text(provider, output),
206            exit_code: *exit_code,
207            metadata: optional_json(provider, metadata),
208            error_kind: error_kind
209                .as_ref()
210                .map(|kind| sanitize_tool_error_kind(provider, kind)),
211        },
212        AgentEvent::ToolOutputDelta { id, name, delta } => AgentEvent::ToolOutputDelta {
213            id: id.clone(),
214            name: name.clone(),
215            delta: text(provider, delta),
216        },
217        AgentEvent::End {
218            text: value,
219            usage,
220            verification_summary: summary,
221            meta,
222        } => AgentEvent::End {
223            text: text(provider, value),
224            usage: usage.clone(),
225            verification_summary: Box::new(verification_summary(provider, summary)),
226            meta: response_meta(provider, meta),
227        },
228        AgentEvent::Error { message } => AgentEvent::Error {
229            message: text(provider, message),
230        },
231        AgentEvent::ConfirmationRequired {
232            tool_id,
233            tool_name,
234            args,
235            timeout_ms,
236        } => AgentEvent::ConfirmationRequired {
237            tool_id: tool_id.clone(),
238            tool_name: tool_name.clone(),
239            args: json(provider, args),
240            timeout_ms: *timeout_ms,
241        },
242        AgentEvent::ConfirmationReceived {
243            tool_id,
244            approved,
245            reason,
246        } => AgentEvent::ConfirmationReceived {
247            tool_id: tool_id.clone(),
248            approved: *approved,
249            reason: optional_text(provider, reason),
250        },
251        AgentEvent::ConfirmationTimeout {
252            tool_id,
253            action_taken,
254        } => AgentEvent::ConfirmationTimeout {
255            tool_id: tool_id.clone(),
256            action_taken: action_taken.clone(),
257        },
258        AgentEvent::ExternalTaskPending {
259            task_id,
260            session_id,
261            lane,
262            command_type,
263            payload,
264            timeout_ms,
265        } => AgentEvent::ExternalTaskPending {
266            task_id: task_id.clone(),
267            session_id: session_id.clone(),
268            lane: *lane,
269            command_type: command_type.clone(),
270            payload: json(provider, payload),
271            timeout_ms: *timeout_ms,
272        },
273        AgentEvent::PermissionDenied {
274            tool_id,
275            tool_name,
276            args,
277            reason,
278        } => AgentEvent::PermissionDenied {
279            tool_id: tool_id.clone(),
280            tool_name: tool_name.clone(),
281            args: json(provider, args),
282            reason: text(provider, reason),
283        },
284        AgentEvent::CommandDeadLettered {
285            command_id,
286            command_type,
287            lane,
288            error,
289            attempts,
290        } => AgentEvent::CommandDeadLettered {
291            command_id: command_id.clone(),
292            command_type: command_type.clone(),
293            lane: lane.clone(),
294            error: text(provider, error),
295            attempts: *attempts,
296        },
297        AgentEvent::QueueAlert {
298            level,
299            alert_type,
300            message,
301        } => AgentEvent::QueueAlert {
302            level: level.clone(),
303            alert_type: alert_type.clone(),
304            message: text(provider, message),
305        },
306        AgentEvent::TaskUpdated { session_id, tasks } => AgentEvent::TaskUpdated {
307            session_id: session_id.clone(),
308            tasks: tasks.iter().map(|item| task(provider, item)).collect(),
309        },
310        AgentEvent::MemoryStored {
311            memory_id,
312            memory_type,
313            importance,
314            tags,
315        } => AgentEvent::MemoryStored {
316            memory_id: memory_id.clone(),
317            memory_type: memory_type.clone(),
318            importance: *importance,
319            tags: strings(provider, tags),
320        },
321        AgentEvent::MemoryRecalled {
322            memory_id,
323            content,
324            relevance,
325        } => AgentEvent::MemoryRecalled {
326            memory_id: memory_id.clone(),
327            content: text(provider, content),
328            relevance: *relevance,
329        },
330        AgentEvent::MemoriesSearched {
331            query,
332            tags,
333            result_count,
334        } => AgentEvent::MemoriesSearched {
335            query: optional_text(provider, query),
336            tags: strings(provider, tags),
337            result_count: *result_count,
338        },
339        AgentEvent::SubagentStart {
340            task_id,
341            session_id,
342            parent_session_id,
343            agent,
344            description,
345            started_ms,
346        } => AgentEvent::SubagentStart {
347            task_id: task_id.clone(),
348            session_id: session_id.clone(),
349            parent_session_id: parent_session_id.clone(),
350            agent: agent.clone(),
351            description: text(provider, description),
352            started_ms: *started_ms,
353        },
354        AgentEvent::SubagentProgress {
355            task_id,
356            session_id,
357            status,
358            metadata,
359        } => AgentEvent::SubagentProgress {
360            task_id: task_id.clone(),
361            session_id: session_id.clone(),
362            status: text(provider, status),
363            metadata: json(provider, metadata),
364        },
365        AgentEvent::SubagentEnd {
366            task_id,
367            session_id,
368            agent,
369            output,
370            success,
371            finished_ms,
372        } => AgentEvent::SubagentEnd {
373            task_id: task_id.clone(),
374            session_id: session_id.clone(),
375            agent: agent.clone(),
376            output: text(provider, output),
377            success: *success,
378            finished_ms: *finished_ms,
379        },
380        AgentEvent::PlanningStart { prompt } => AgentEvent::PlanningStart {
381            prompt: text(provider, prompt),
382        },
383        AgentEvent::PlanningEnd {
384            plan: value,
385            estimated_steps,
386        } => AgentEvent::PlanningEnd {
387            plan: plan(provider, value),
388            estimated_steps: *estimated_steps,
389        },
390        AgentEvent::StepStart {
391            step_id,
392            description,
393            step_number,
394            total_steps,
395        } => AgentEvent::StepStart {
396            step_id: step_id.clone(),
397            description: text(provider, description),
398            step_number: *step_number,
399            total_steps: *total_steps,
400        },
401        AgentEvent::GoalExtracted { goal: value } => AgentEvent::GoalExtracted {
402            goal: goal(provider, value),
403        },
404        AgentEvent::GoalProgress {
405            goal,
406            progress,
407            completed_steps,
408            total_steps,
409        } => AgentEvent::GoalProgress {
410            goal: text(provider, goal),
411            progress: *progress,
412            completed_steps: *completed_steps,
413            total_steps: *total_steps,
414        },
415        AgentEvent::GoalAchieved {
416            goal,
417            total_steps,
418            duration_ms,
419        } => AgentEvent::GoalAchieved {
420            goal: text(provider, goal),
421            total_steps: *total_steps,
422            duration_ms: *duration_ms,
423        },
424        AgentEvent::PersistenceFailed {
425            session_id,
426            operation,
427            error,
428        } => AgentEvent::PersistenceFailed {
429            session_id: session_id.clone(),
430            operation: operation.clone(),
431            error: text(provider, error),
432        },
433        AgentEvent::BudgetThresholdHit {
434            resource,
435            kind,
436            consumed,
437            limit,
438            message,
439        } => AgentEvent::BudgetThresholdHit {
440            resource: resource.clone(),
441            kind: kind.clone(),
442            consumed: *consumed,
443            limit: *limit,
444            message: optional_text(provider, message),
445        },
446        AgentEvent::PassivationRequested {
447            reason,
448            deadline_ms,
449        } => AgentEvent::PassivationRequested {
450            reason: text(provider, reason),
451            deadline_ms: *deadline_ms,
452        },
453        _ => event.clone(),
454    }
455}
456
457/// Trait for pluggable security providers.
458///
459/// Implement this trait to provide custom security logic for sessions.
460/// The default `NoOpSecurityProvider` passes everything through unchanged.
461pub trait SecurityProvider: Send + Sync {
462    /// Classify and register sensitive data found in input text
463    fn taint_input(&self, _text: &str) {}
464
465    /// Sanitize output text by redacting sensitive data.
466    /// Returns the sanitized text.
467    fn sanitize_output(&self, text: &str) -> String {
468        text.to_string()
469    }
470
471    /// Securely wipe all session security state
472    fn wipe(&self) {}
473
474    /// Register security hooks with the given engine
475    fn register_hooks(&self, _hook_engine: &HookEngine) {}
476
477    /// Unregister all hooks from the engine
478    fn teardown(&self, _hook_engine: &HookEngine) {}
479}
480
481/// No-op security provider (default when security is disabled)
482pub struct NoOpSecurityProvider;
483
484impl SecurityProvider for NoOpSecurityProvider {}
485
486#[cfg(test)]
487mod tests {
488    use super::*;
489
490    #[test]
491    fn test_noop_provider_passthrough() {
492        let provider = NoOpSecurityProvider;
493        provider.taint_input("SSN: 123-45-6789");
494        let output = provider.sanitize_output("SSN: 123-45-6789");
495        assert_eq!(output, "SSN: 123-45-6789");
496    }
497
498    #[test]
499    fn test_noop_provider_wipe() {
500        let provider = NoOpSecurityProvider;
501        provider.wipe(); // Should not panic
502    }
503
504    #[test]
505    fn test_noop_provider_hooks() {
506        let engine = HookEngine::new();
507        let provider = NoOpSecurityProvider;
508        provider.register_hooks(&engine);
509        provider.teardown(&engine);
510        assert_eq!(engine.hook_count(), 0);
511    }
512
513    #[test]
514    fn agent_event_sanitization_redacts_streams_arguments_and_outputs() {
515        let provider = DefaultSecurityProvider::new();
516        let secret = "user@example.com";
517        let events = [
518            crate::agent::AgentEvent::TextDelta {
519                text: secret.to_string(),
520            },
521            crate::agent::AgentEvent::ReasoningDelta {
522                text: secret.to_string(),
523            },
524            crate::agent::AgentEvent::ToolInputDelta {
525                id: Some("tool-1".to_string()),
526                delta: format!(r#"{{"token":"{secret}"}}"#),
527            },
528            crate::agent::AgentEvent::ToolExecutionStart {
529                id: "tool-1".to_string(),
530                name: "bash".to_string(),
531                args: serde_json::json!({"command": format!("echo {secret}")}),
532            },
533            crate::agent::AgentEvent::ToolOutputDelta {
534                id: "tool-1".to_string(),
535                name: "bash".to_string(),
536                delta: secret.to_string(),
537            },
538            crate::agent::AgentEvent::ToolEnd {
539                id: "tool-1".to_string(),
540                name: "bash".to_string(),
541                args: Some(serde_json::json!({"command": format!("echo {secret}")})),
542                output: secret.to_string(),
543                exit_code: 0,
544                metadata: Some(serde_json::json!({"contact": secret})),
545                error_kind: None,
546            },
547        ];
548
549        for event in &events {
550            let sanitized = sanitize_agent_event(&provider, event);
551            let json = serde_json::to_string(&sanitized).unwrap();
552            assert!(!json.contains(secret), "unsanitized event: {json}");
553            assert!(json.contains("REDACTED:EMAIL"));
554        }
555
556        let sanitized = sanitize_agent_event(&provider, &events[3]);
557        assert!(matches!(
558            sanitized,
559            crate::agent::AgentEvent::ToolExecutionStart { id, name, .. }
560                if id == "tool-1" && name == "bash"
561        ));
562    }
563
564    #[test]
565    fn agent_event_sanitization_redacts_every_tool_error_kind_string() {
566        use crate::agent::AgentEvent;
567        use crate::tools::ToolErrorKind;
568
569        let provider = DefaultSecurityProvider::new();
570        let secret = "user@example.com";
571        let redacted = "[REDACTED:EMAIL]";
572        let cases = [
573            (
574                ToolErrorKind::VersionConflict {
575                    path: format!("path/{secret}"),
576                    expected: format!("expected: {secret}"),
577                    actual: Some(format!("actual: {secret}")),
578                },
579                ToolErrorKind::VersionConflict {
580                    path: format!("path/{redacted}"),
581                    expected: format!("expected: {redacted}"),
582                    actual: Some(format!("actual: {redacted}")),
583                },
584            ),
585            (
586                ToolErrorKind::RemoteGitConflict {
587                    code: format!("code: {secret}"),
588                    message: format!("message: {secret}"),
589                },
590                ToolErrorKind::RemoteGitConflict {
591                    code: format!("code: {redacted}"),
592                    message: format!("message: {redacted}"),
593                },
594            ),
595            (
596                ToolErrorKind::NotFound {
597                    path: format!("path/{secret}"),
598                },
599                ToolErrorKind::NotFound {
600                    path: format!("path/{redacted}"),
601                },
602            ),
603            (
604                ToolErrorKind::InvalidArgument {
605                    message: format!("message: {secret}"),
606                },
607                ToolErrorKind::InvalidArgument {
608                    message: format!("message: {redacted}"),
609                },
610            ),
611            (
612                ToolErrorKind::HookDenied {
613                    reason: format!("reason: {secret}"),
614                    retryable: true,
615                    retry_after_ms: Some(250),
616                },
617                ToolErrorKind::HookDenied {
618                    reason: format!("reason: {redacted}"),
619                    retryable: true,
620                    retry_after_ms: Some(250),
621                },
622            ),
623            (
624                ToolErrorKind::Unsupported {
625                    message: format!("message: {secret}"),
626                },
627                ToolErrorKind::Unsupported {
628                    message: format!("message: {redacted}"),
629                },
630            ),
631            (
632                ToolErrorKind::Timeout {
633                    op: format!("operation: {secret}"),
634                    duration_ms: 42,
635                },
636                ToolErrorKind::Timeout {
637                    op: format!("operation: {redacted}"),
638                    duration_ms: 42,
639                },
640            ),
641            (
642                ToolErrorKind::Transport {
643                    op: format!("operation: {secret}"),
644                },
645                ToolErrorKind::Transport {
646                    op: format!("operation: {redacted}"),
647                },
648            ),
649        ];
650
651        for (error_kind, expected) in cases {
652            let event = AgentEvent::ToolEnd {
653                id: "tool-1".to_string(),
654                name: "test".to_string(),
655                args: None,
656                output: String::new(),
657                exit_code: 1,
658                metadata: None,
659                error_kind: Some(error_kind),
660            };
661
662            let sanitized = sanitize_agent_event(&provider, &event);
663            let AgentEvent::ToolEnd {
664                id,
665                name,
666                exit_code,
667                error_kind,
668                ..
669            } = sanitized
670            else {
671                panic!("sanitization changed the event variant");
672            };
673
674            assert_eq!(id, "tool-1");
675            assert_eq!(name, "test");
676            assert_eq!(exit_code, 1);
677            assert_eq!(error_kind, Some(expected));
678        }
679    }
680}
a3s_code_core/security/mod.rs

a3s_code_core/security/
mod.rs
