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            .filter(|s| !s.is_empty())
33            .map(TraceId::new)
34            .or_else(|| {
35                Self::extract_header(headers, headers::TRACEPARENT)
36                    .and_then(|value| TraceId::from_traceparent(&value))
37            })
38            .unwrap_or_else(TraceId::generate)
39    }
40
41    pub fn extract_context_id(headers: &HeaderMap) -> Option<ContextId> {
42        Self::extract_header(headers, headers::CONTEXT_ID)
43            .filter(|s| !s.is_empty())
44            .and_then(|s| {
45                ContextId::try_new(s)
46                    .map_err(|e| {
47                        tracing::warn!(error = %e, "Invalid context_id header value, ignoring");
48                        e
49                    })
50                    .ok()
51            })
52    }
53
54    pub fn extract_gateway_conversation_id(headers: &HeaderMap) -> Option<GatewayConversationId> {
55        Self::extract_header(headers, headers::GATEWAY_CONVERSATION_ID)
56            .filter(|s| !s.is_empty())
57            .and_then(|s| {
58                GatewayConversationId::try_new(s)
59                    .map_err(|e| {
60                        tracing::warn!(error = %e, "Invalid gateway_conversation_id header value, ignoring");
61                        e
62                    })
63                    .ok()
64            })
65    }
66
67    pub fn extract_provider_request_id(headers: &HeaderMap) -> Option<ProviderRequestId> {
68        Self::extract_header(headers, headers::PROVIDER_REQUEST_ID)
69            .filter(|s| !s.is_empty())
70            .and_then(|s| {
71                ProviderRequestId::try_new(s)
72                    .map_err(|e| {
73                        tracing::warn!(error = %e, "Invalid provider_request_id header value, ignoring");
74                        e
75                    })
76                    .ok()
77            })
78    }
79
80    pub fn extract_task_id(headers: &HeaderMap) -> Option<TaskId> {
81        Self::extract_header(headers, headers::TASK_ID).map(TaskId::new)
82    }
83
84    pub fn extract_agent_name(headers: &HeaderMap) -> AgentName {
85        Self::extract_header(headers, headers::AGENT_NAME)
86            .and_then(|s| {
87                AgentName::try_new(s)
88                    .map_err(|e| {
89                        tracing::warn!(error = %e, "Invalid agent_name header value, using system");
90                        e
91                    })
92                    .ok()
93            })
94            .unwrap_or_else(AgentName::system)
95    }
96
97    fn extract_header(headers: &HeaderMap, name: &str) -> Option<String> {
98        headers
99            .get(name)
100            .and_then(|v| {
101                v.to_str()
102                    .map_err(|e| {
103                        tracing::debug!(error = %e, header = %name, "Header contains non-ASCII characters");
104                        e
105                    })
106                    .ok()
107            })
108            .map(str::to_owned)
109    }
110}
111
112#[derive(Debug, Clone, Copy)]
113pub struct HeaderInjector;
114
115impl HeaderInjector {
116    pub fn inject_session_id(
117        headers: &mut HeaderMap,
118        session_id: &SessionId,
119    ) -> Result<(), HeaderInjectionError> {
120        Self::inject_header(headers, headers::SESSION_ID, session_id.as_str())
121    }
122
123    pub fn inject_user_id(
124        headers: &mut HeaderMap,
125        user_id: &UserId,
126    ) -> Result<(), HeaderInjectionError> {
127        Self::inject_header(headers, headers::USER_ID, user_id.as_str())
128    }
129
130    pub fn inject_trace_id(
131        headers: &mut HeaderMap,
132        trace_id: &TraceId,
133    ) -> Result<(), HeaderInjectionError> {
134        Self::inject_header(headers, headers::TRACE_ID, trace_id.as_str())
135    }
136
137    pub fn inject_context_id(
138        headers: &mut HeaderMap,
139        context_id: &ContextId,
140    ) -> Result<(), HeaderInjectionError> {
141        Self::inject_header(headers, headers::CONTEXT_ID, context_id.as_str())
142    }
143
144    pub fn inject_gateway_conversation_id(
145        headers: &mut HeaderMap,
146        id: &GatewayConversationId,
147    ) -> Result<(), HeaderInjectionError> {
148        Self::inject_header(headers, headers::GATEWAY_CONVERSATION_ID, id.as_str())
149    }
150
151    pub fn inject_provider_request_id(
152        headers: &mut HeaderMap,
153        id: &ProviderRequestId,
154    ) -> Result<(), HeaderInjectionError> {
155        Self::inject_header(headers, headers::PROVIDER_REQUEST_ID, id.as_str())
156    }
157
158    pub fn inject_task_id(
159        headers: &mut HeaderMap,
160        task_id: &TaskId,
161    ) -> Result<(), HeaderInjectionError> {
162        Self::inject_header(headers, headers::TASK_ID, task_id.as_str())
163    }
164
165    pub fn inject_agent_name(
166        headers: &mut HeaderMap,
167        agent_name: &AgentName,
168    ) -> Result<(), HeaderInjectionError> {
169        Self::inject_header(headers, headers::AGENT_NAME, agent_name.as_str())
170    }
171
172    pub fn inject_from_request_context(
173        headers: &mut HeaderMap,
174        ctx: &RequestContext,
175    ) -> Result<(), HeaderInjectionError> {
176        Self::inject_session_id(headers, &ctx.request.session_id)?;
177        Self::inject_user_id(headers, &ctx.auth.actor.user_id)?;
178        Self::inject_trace_id(headers, &ctx.execution.trace_id)?;
179        Self::inject_context_id(headers, &ctx.execution.context_id)?;
180        Self::inject_agent_name(headers, &ctx.execution.agent_name)?;
181        Ok(())
182    }
183
184    fn inject_header(
185        headers: &mut HeaderMap,
186        name: &'static str,
187        value: &str,
188    ) -> Result<(), HeaderInjectionError> {
189        HeaderValue::from_str(value).map_or(Err(HeaderInjectionError), |header_value| {
190            headers.insert(name, header_value);
191            Ok(())
192        })
193    }
194}
systemprompt_security/extraction/header.rs

systemprompt_security/extraction/
header.rs
