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

a3s_code_core/security/
mod.rs
