starweaver_model/transport/
audit.rs1use std::{
2 collections::{BTreeMap, BTreeSet},
3 sync::{Arc, Mutex},
4};
5
6use serde::{Deserialize, Serialize};
7use serde_json::{Map, Value};
8
9use super::{HttpMethod, HttpRequest};
10
11pub type DynProviderRequestAuditRecorder = Arc<dyn ProviderRequestAuditRecorder>;
13
14#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
16#[serde(rename_all = "snake_case")]
17pub enum ProviderRequestAuditPayloadPolicy {
18 #[default]
20 Omit,
21 Redacted,
23 Full,
25}
26
27#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
29pub struct ProviderRequestAuditPolicy {
30 #[serde(default)]
32 pub headers: ProviderRequestAuditPayloadPolicy,
33 #[serde(default)]
35 pub body: ProviderRequestAuditPayloadPolicy,
36 #[serde(default = "default_sensitive_header_keys")]
38 pub sensitive_header_keys: BTreeSet<String>,
39 #[serde(default = "default_sensitive_body_keys")]
41 pub sensitive_body_keys: BTreeSet<String>,
42 #[serde(default = "default_redaction_value")]
44 pub redaction_value: Value,
45}
46
47impl Default for ProviderRequestAuditPolicy {
48 fn default() -> Self {
49 Self {
50 headers: ProviderRequestAuditPayloadPolicy::Omit,
51 body: ProviderRequestAuditPayloadPolicy::Omit,
52 sensitive_header_keys: default_sensitive_header_keys(),
53 sensitive_body_keys: default_sensitive_body_keys(),
54 redaction_value: default_redaction_value(),
55 }
56 }
57}
58
59impl ProviderRequestAuditPolicy {
60 #[must_use]
62 pub fn metadata_only() -> Self {
63 Self::default()
64 }
65
66 #[must_use]
68 pub fn redacted_payloads() -> Self {
69 Self {
70 headers: ProviderRequestAuditPayloadPolicy::Redacted,
71 body: ProviderRequestAuditPayloadPolicy::Redacted,
72 ..Self::default()
73 }
74 }
75
76 #[must_use]
78 pub fn full_payloads() -> Self {
79 Self {
80 headers: ProviderRequestAuditPayloadPolicy::Full,
81 body: ProviderRequestAuditPayloadPolicy::Full,
82 ..Self::default()
83 }
84 }
85
86 #[must_use]
88 pub fn with_sensitive_header_key(mut self, key: impl AsRef<str>) -> Self {
89 self.sensitive_header_keys
90 .insert(normalize_key(key.as_ref()));
91 self
92 }
93
94 #[must_use]
96 pub fn with_sensitive_body_key(mut self, key: impl AsRef<str>) -> Self {
97 self.sensitive_body_keys.insert(normalize_key(key.as_ref()));
98 self
99 }
100}
101
102#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
104pub struct ProviderRequestAuditSnapshot {
105 pub provider_name: String,
107 pub model_name: String,
109 pub stream: bool,
111 pub method: HttpMethod,
113 pub url: String,
115 #[serde(default, skip_serializing_if = "Option::is_none")]
117 pub headers: Option<BTreeMap<String, String>>,
118 #[serde(default, skip_serializing_if = "Option::is_none")]
120 pub body: Option<Value>,
121 #[serde(default, skip_serializing_if = "Option::is_none")]
123 pub timeout_ms: Option<u64>,
124 #[serde(default, skip_serializing_if = "Map::is_empty")]
126 pub metadata: Map<String, Value>,
127}
128
129impl ProviderRequestAuditSnapshot {
130 #[must_use]
132 pub fn from_request(
133 provider_name: impl Into<String>,
134 model_name: impl Into<String>,
135 stream: bool,
136 request: &HttpRequest,
137 policy: &ProviderRequestAuditPolicy,
138 ) -> Self {
139 Self {
140 provider_name: provider_name.into(),
141 model_name: model_name.into(),
142 stream,
143 method: request.method,
144 url: request.url.clone(),
145 headers: capture_headers(&request.headers, policy),
146 body: capture_body(&request.body, policy),
147 timeout_ms: request
148 .timeout
149 .and_then(|duration| u64::try_from(duration.as_millis()).ok()),
150 metadata: request.metadata.clone(),
151 }
152 }
153}
154
155pub trait ProviderRequestAuditRecorder: Send + Sync {
157 fn record_provider_request(&self, snapshot: ProviderRequestAuditSnapshot);
159}
160
161#[derive(Default)]
163pub struct InMemoryProviderRequestAuditRecorder {
164 snapshots: Mutex<Vec<ProviderRequestAuditSnapshot>>,
165}
166
167impl InMemoryProviderRequestAuditRecorder {
168 #[must_use]
170 pub fn new() -> Self {
171 Self::default()
172 }
173
174 #[must_use]
176 pub fn snapshots(&self) -> Vec<ProviderRequestAuditSnapshot> {
177 self.snapshots
178 .lock()
179 .map_or_else(|_| Vec::new(), |v| v.clone())
180 }
181
182 #[must_use]
184 pub fn take_snapshots(&self) -> Vec<ProviderRequestAuditSnapshot> {
185 self.snapshots
186 .lock()
187 .map_or_else(|_| Vec::new(), |mut v| std::mem::take(&mut *v))
188 }
189}
190
191impl ProviderRequestAuditRecorder for InMemoryProviderRequestAuditRecorder {
192 fn record_provider_request(&self, snapshot: ProviderRequestAuditSnapshot) {
193 if let Ok(mut snapshots) = self.snapshots.lock() {
194 snapshots.push(snapshot);
195 }
196 }
197}
198
199#[derive(Clone)]
200pub struct ProviderRequestAuditCapture {
201 recorder: DynProviderRequestAuditRecorder,
202 policy: ProviderRequestAuditPolicy,
203}
204
205impl ProviderRequestAuditCapture {
206 pub(crate) fn new(
207 recorder: DynProviderRequestAuditRecorder,
208 policy: ProviderRequestAuditPolicy,
209 ) -> Self {
210 Self { recorder, policy }
211 }
212
213 pub(crate) fn record(
214 &self,
215 provider_name: &str,
216 model_name: &str,
217 stream: bool,
218 request: &HttpRequest,
219 ) {
220 self.recorder
221 .record_provider_request(ProviderRequestAuditSnapshot::from_request(
222 provider_name,
223 model_name,
224 stream,
225 request,
226 &self.policy,
227 ));
228 }
229}
230
231fn capture_headers(
232 headers: &BTreeMap<String, String>,
233 policy: &ProviderRequestAuditPolicy,
234) -> Option<BTreeMap<String, String>> {
235 match policy.headers {
236 ProviderRequestAuditPayloadPolicy::Omit => None,
237 ProviderRequestAuditPayloadPolicy::Full => Some(headers.clone()),
238 ProviderRequestAuditPayloadPolicy::Redacted => Some(
239 headers
240 .iter()
241 .map(|(key, value)| {
242 if policy.sensitive_header_keys.contains(&normalize_key(key)) {
243 (key.clone(), policy.redaction_value.to_string())
244 } else {
245 (key.clone(), value.clone())
246 }
247 })
248 .collect(),
249 ),
250 }
251}
252
253fn capture_body(body: &Value, policy: &ProviderRequestAuditPolicy) -> Option<Value> {
254 match policy.body {
255 ProviderRequestAuditPayloadPolicy::Omit => None,
256 ProviderRequestAuditPayloadPolicy::Full => Some(body.clone()),
257 ProviderRequestAuditPayloadPolicy::Redacted => {
258 let mut body = body.clone();
259 redact_json_value(&mut body, policy);
260 Some(body)
261 }
262 }
263}
264
265fn redact_json_value(value: &mut Value, policy: &ProviderRequestAuditPolicy) {
266 match value {
267 Value::Object(object) => {
268 for (key, value) in object {
269 if policy.sensitive_body_keys.contains(&normalize_key(key)) {
270 *value = policy.redaction_value.clone();
271 } else {
272 redact_json_value(value, policy);
273 }
274 }
275 }
276 Value::Array(values) => {
277 for value in values {
278 redact_json_value(value, policy);
279 }
280 }
281 Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {}
282 }
283}
284
285fn normalize_key(key: &str) -> String {
286 key.to_ascii_lowercase()
287}
288
289fn default_sensitive_header_keys() -> BTreeSet<String> {
290 [
291 "authorization",
292 "proxy-authorization",
293 "x-api-key",
294 "api-key",
295 "cookie",
296 "set-cookie",
297 ]
298 .into_iter()
299 .map(str::to_string)
300 .collect()
301}
302
303fn default_sensitive_body_keys() -> BTreeSet<String> {
304 [
305 "api_key",
306 "apikey",
307 "authorization",
308 "access_token",
309 "refresh_token",
310 "id_token",
311 "token",
312 "password",
313 "secret",
314 ]
315 .into_iter()
316 .map(str::to_string)
317 .collect()
318}
319
320fn default_redaction_value() -> Value {
321 serde_json::json!({ "redacted": true })
322}