systemprompt_security/extraction/
header.rs

1//! Header-based credential extraction and injection.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use axum::http::{HeaderMap, HeaderValue};
7use std::error::Error;
8use std::fmt;
9use systemprompt_identifiers::{
10    AgentName, ContextId, GatewayConversationId, ProviderRequestId, SessionId, TaskId, TraceId,
11    UserId, headers,
12};
13use systemprompt_models::execution::context::RequestContext;
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub struct HeaderInjectionError;
17
18impl fmt::Display for HeaderInjectionError {
19    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
20        write!(f, "Header value contains invalid characters")
21    }
22}
23
24impl Error for HeaderInjectionError {}
25
26#[derive(Debug, Clone, Copy)]
27pub struct HeaderExtractor;
28
29impl HeaderExtractor {
30    pub fn extract_trace_id(headers: &HeaderMap) -> TraceId {
31        Self::extract_header(headers, headers::TRACE_ID)
32            .map_or_else(TraceId::generate, TraceId::new)
33    }
34
35    pub fn extract_context_id(headers: &HeaderMap) -> Option<ContextId> {
36        Self::extract_header(headers, headers::CONTEXT_ID)
37            .filter(|s| !s.is_empty())
38            .and_then(|s| {
39                ContextId::try_new(s)
40                    .map_err(|e| {
41                        tracing::warn!(error = %e, "Invalid context_id header value, ignoring");
42                        e
43                    })
44                    .ok()
45            })
46    }
47
48    pub fn extract_gateway_conversation_id(headers: &HeaderMap) -> Option<GatewayConversationId> {
49        Self::extract_header(headers, headers::GATEWAY_CONVERSATION_ID)
50            .filter(|s| !s.is_empty())
51            .and_then(|s| {
52                GatewayConversationId::try_new(s)
53                    .map_err(|e| {
54                        tracing::warn!(error = %e, "Invalid gateway_conversation_id header value, ignoring");
55                        e
56                    })
57                    .ok()
58            })
59    }
60
61    pub fn extract_provider_request_id(headers: &HeaderMap) -> Option<ProviderRequestId> {
62        Self::extract_header(headers, headers::PROVIDER_REQUEST_ID)
63            .filter(|s| !s.is_empty())
64            .and_then(|s| {
65                ProviderRequestId::try_new(s)
66                    .map_err(|e| {
67                        tracing::warn!(error = %e, "Invalid provider_request_id header value, ignoring");
68                        e
69                    })
70                    .ok()
71            })
72    }
73
74    pub fn extract_task_id(headers: &HeaderMap) -> Option<TaskId> {
75        Self::extract_header(headers, headers::TASK_ID).map(TaskId::new)
76    }
77
78    pub fn extract_agent_name(headers: &HeaderMap) -> AgentName {
79        Self::extract_header(headers, headers::AGENT_NAME)
80            .map_or_else(AgentName::system, AgentName::new)
81    }
82
83    fn extract_header(headers: &HeaderMap, name: &str) -> Option<String> {
84        headers
85            .get(name)
86            .and_then(|v| {
87                v.to_str()
88                    .map_err(|e| {
89                        tracing::debug!(error = %e, header = %name, "Header contains non-ASCII characters");
90                        e
91                    })
92                    .ok()
93            })
94            .map(str::to_owned)
95    }
96}
97
98#[derive(Debug, Clone, Copy)]
99pub struct HeaderInjector;
100
101impl HeaderInjector {
102    pub fn inject_session_id(
103        headers: &mut HeaderMap,
104        session_id: &SessionId,
105    ) -> Result<(), HeaderInjectionError> {
106        Self::inject_header(headers, headers::SESSION_ID, session_id.as_str())
107    }
108
109    pub fn inject_user_id(
110        headers: &mut HeaderMap,
111        user_id: &UserId,
112    ) -> Result<(), HeaderInjectionError> {
113        Self::inject_header(headers, headers::USER_ID, user_id.as_str())
114    }
115
116    pub fn inject_trace_id(
117        headers: &mut HeaderMap,
118        trace_id: &TraceId,
119    ) -> Result<(), HeaderInjectionError> {
120        Self::inject_header(headers, headers::TRACE_ID, trace_id.as_str())
121    }
122
123    pub fn inject_context_id(
124        headers: &mut HeaderMap,
125        context_id: &ContextId,
126    ) -> Result<(), HeaderInjectionError> {
127        Self::inject_header(headers, headers::CONTEXT_ID, context_id.as_str())
128    }
129
130    pub fn inject_gateway_conversation_id(
131        headers: &mut HeaderMap,
132        id: &GatewayConversationId,
133    ) -> Result<(), HeaderInjectionError> {
134        Self::inject_header(headers, headers::GATEWAY_CONVERSATION_ID, id.as_str())
135    }
136
137    pub fn inject_provider_request_id(
138        headers: &mut HeaderMap,
139        id: &ProviderRequestId,
140    ) -> Result<(), HeaderInjectionError> {
141        Self::inject_header(headers, headers::PROVIDER_REQUEST_ID, id.as_str())
142    }
143
144    pub fn inject_task_id(
145        headers: &mut HeaderMap,
146        task_id: &TaskId,
147    ) -> Result<(), HeaderInjectionError> {
148        Self::inject_header(headers, headers::TASK_ID, task_id.as_str())
149    }
150
151    pub fn inject_agent_name(
152        headers: &mut HeaderMap,
153        agent_name: &str,
154    ) -> Result<(), HeaderInjectionError> {
155        Self::inject_header(headers, headers::AGENT_NAME, agent_name)
156    }
157
158    pub fn inject_from_request_context(
159        headers: &mut HeaderMap,
160        ctx: &RequestContext,
161    ) -> Result<(), HeaderInjectionError> {
162        Self::inject_session_id(headers, &ctx.request.session_id)?;
163        Self::inject_user_id(headers, &ctx.auth.actor.user_id)?;
164        Self::inject_trace_id(headers, &ctx.execution.trace_id)?;
165        Self::inject_context_id(headers, &ctx.execution.context_id)?;
166        Self::inject_agent_name(headers, ctx.execution.agent_name.as_str())?;
167        Ok(())
168    }
169
170    fn inject_header(
171        headers: &mut HeaderMap,
172        name: &'static str,
173        value: &str,
174    ) -> Result<(), HeaderInjectionError> {
175        HeaderValue::from_str(value).map_or(Err(HeaderInjectionError), |header_value| {
176            headers.insert(name, header_value);
177            Ok(())
178        })
179    }
180}
systemprompt_security/extraction/header.rs

systemprompt_security/extraction/
header.rs
