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

a3s_code_core/security/
mod.rs
