1use std::collections::{HashMap, HashSet};
2use std::fmt;
3use std::net::{IpAddr, SocketAddr};
4use std::sync::Arc;
5use std::time::{Duration, Instant};
6
7use axum::body::{Body, Bytes};
8use axum::extract::{ConnectInfo, State};
9use axum::http::{HeaderMap, HeaderName, HeaderValue, Method, Request, StatusCode, Uri};
10use axum::response::{IntoResponse, Response};
11use axum::routing::get;
12use axum::Router;
13use futures_util::StreamExt;
14use gaze::{
15 token_shape, BoundaryRefusal, CleanDocument, CommittedSessionSnapshot, DictionaryBundle,
16 EmittedTokenSpan, LocaleTag, Pipeline, PrefixCacheWriteMode, RawDocument, Scope, Session,
17 SessionTransaction, SessionTransactionError,
18};
19use reqwest::{Client, ClientBuilder};
20use serde::de::{MapAccess, SeqAccess, Visitor};
21use serde::{Deserialize, Deserializer, Serialize};
22use serde_json::{Map, Number, Value};
23use tokio::net::TcpListener;
24use tokio::sync::RwLock;
25use url::{Host, Url};
26
27use crate::adapter::{
28 CoveragePolicy, ProtocolContract, ProviderAdapter, SessionPolicy, SseEvent, SurfaceSyntax,
29};
30use crate::adapters::anthropic::{AnthropicAdapter, DEFAULT_ANTHROPIC_VERSION};
31use crate::codec::{
32 BodyCodec, CodecError, CodecErrorCode, CodecLimits, CodecPhase, InspectionParsedFactsV1,
33 ProvedRequestBody, RequestPseudonymizer, RequestTransformContext, ResponseResidualValidator,
34 ResponseTransformContext, WireFormat,
35};
36use crate::error::{
37 DirectProxyError, HeaderRejectionReason, ProxyError, ProxyErrorCode, ProxyErrorPhase,
38};
39use crate::inspection::{ProxyInspectionEndpointCodesV1, ProxyInspectionLogicalV1};
40use crate::principal::{
41 invoke_resolver, ListenerScope, PeerScope, PrincipalContext, PrincipalResolver,
42 ProcessLocalLoopbackResolver,
43};
44use crate::session_registry::{RegistryError, SessionLease, SessionRegistry};
45use crate::ProxyConfig;
46
47const SESSION_HEADER: &str = "x-gaze-session-id";
48
49pub const DIRECT_PROXY_PING_FRAME: &[u8] = b"event: ping\ndata: {\"type\":\"ping\"}\n\n";
51
52pub const DIRECT_PROXY_ERROR_FRAME: &[u8] = b"event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"api_error\",\"message\":\"proxy_validation_failed\"}}\n\n";
54
55const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
56const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120);
57const DEFAULT_TOTAL_TIMEOUT: Duration = Duration::from_secs(180);
58const DEFAULT_MAX_REQUEST_BYTES: usize = 2 * 1024 * 1024;
59const DEFAULT_MAX_RESPONSE_BYTES: usize = 32 * 1024 * 1024;
60const DEFAULT_MAX_ERROR_BODY_BYTES: usize = 64 * 1024;
61const DEFAULT_MAX_HEADERS: usize = 64;
62const DEFAULT_MAX_HEADER_NAME_BYTES: usize = 128;
63const DEFAULT_MAX_HEADER_VALUE_BYTES: usize = 8 * 1024;
64const DEFAULT_MAX_HEADER_BYTES: usize = 64 * 1024;
65
66#[derive(Clone, Copy, Debug, Eq, PartialEq)]
68pub(crate) struct DirectClientConfig {
69 connect_timeout: Duration,
70 request_timeout: Duration,
71 total_timeout: Duration,
72 max_request_bytes: usize,
73 max_response_bytes: usize,
74 max_error_body_bytes: usize,
75 max_headers: usize,
76 max_header_name_bytes: usize,
77 max_header_value_bytes: usize,
78 max_header_bytes: usize,
79}
80
81impl Default for DirectClientConfig {
82 fn default() -> Self {
83 Self {
84 connect_timeout: DEFAULT_CONNECT_TIMEOUT,
85 request_timeout: DEFAULT_REQUEST_TIMEOUT,
86 total_timeout: DEFAULT_TOTAL_TIMEOUT,
87 max_request_bytes: DEFAULT_MAX_REQUEST_BYTES,
88 max_response_bytes: DEFAULT_MAX_RESPONSE_BYTES,
89 max_error_body_bytes: DEFAULT_MAX_ERROR_BODY_BYTES,
90 max_headers: DEFAULT_MAX_HEADERS,
91 max_header_name_bytes: DEFAULT_MAX_HEADER_NAME_BYTES,
92 max_header_value_bytes: DEFAULT_MAX_HEADER_VALUE_BYTES,
93 max_header_bytes: DEFAULT_MAX_HEADER_BYTES,
94 }
95 }
96}
97
98#[allow(dead_code)]
100impl DirectClientConfig {
101 fn from_anthropic(adapter: Option<&AnthropicAdapter>) -> Self {
102 let Some(adapter) = adapter else {
103 return Self::default();
104 };
105 let (max_headers, max_header_name_bytes, max_header_value_bytes, max_header_bytes) =
106 adapter.header_limits();
107 Self {
108 connect_timeout: adapter.connect_timeout(),
109 request_timeout: adapter.request_timeout(),
110 total_timeout: adapter.total_timeout(),
111 max_headers,
112 max_header_name_bytes,
113 max_header_value_bytes,
114 max_header_bytes,
115 ..Self::default()
116 }
117 }
118
119 #[must_use]
120 pub(crate) const fn connect_timeout(self) -> Duration {
121 self.connect_timeout
122 }
123
124 #[must_use]
125 pub(crate) const fn request_timeout(self) -> Duration {
126 self.request_timeout
127 }
128
129 #[must_use]
130 pub(crate) const fn total_timeout(self) -> Duration {
131 self.total_timeout
132 }
133
134 #[must_use]
135 pub(crate) const fn max_request_bytes(self) -> usize {
136 self.max_request_bytes
137 }
138
139 #[must_use]
140 pub(crate) const fn max_response_bytes(self) -> usize {
141 self.max_response_bytes
142 }
143
144 #[must_use]
145 pub(crate) const fn max_error_body_bytes(self) -> usize {
146 self.max_error_body_bytes
147 }
148
149 pub(crate) fn try_with_timeouts(
151 mut self,
152 connect: Duration,
153 request: Duration,
154 total: Duration,
155 ) -> Result<Self, DirectProxyError> {
156 if connect.is_zero()
157 || request.is_zero()
158 || total.is_zero()
159 || connect > DEFAULT_CONNECT_TIMEOUT
160 || request > DEFAULT_REQUEST_TIMEOUT
161 || total > DEFAULT_TOTAL_TIMEOUT
162 {
163 return Err(
164 ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::UpstreamConfiguration)
165 );
166 }
167 self.connect_timeout = connect;
168 self.request_timeout = request;
169 self.total_timeout = total;
170 Ok(self)
171 }
172
173 pub(crate) fn try_with_max_response_bytes(
175 mut self,
176 value: usize,
177 ) -> Result<Self, DirectProxyError> {
178 if value == 0 || value > DEFAULT_MAX_RESPONSE_BYTES {
179 return Err(
180 ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::UpstreamConfiguration)
181 );
182 }
183 self.max_response_bytes = value;
184 Ok(self)
185 }
186
187 pub(crate) fn try_with_max_error_body_bytes(
189 mut self,
190 value: usize,
191 ) -> Result<Self, DirectProxyError> {
192 if value == 0 || value > DEFAULT_MAX_ERROR_BODY_BYTES {
193 return Err(
194 ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::UpstreamConfiguration)
195 );
196 }
197 self.max_error_body_bytes = value;
198 Ok(self)
199 }
200
201 pub(crate) fn try_with_max_request_bytes(
203 mut self,
204 value: usize,
205 ) -> Result<Self, DirectProxyError> {
206 if value == 0 || value > DEFAULT_MAX_REQUEST_BYTES {
207 return Err(
208 ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::UpstreamConfiguration)
209 );
210 }
211 self.max_request_bytes = value;
212 Ok(self)
213 }
214}
215
216enum ResolvedDirectAdapter {
218 Configured(Arc<AnthropicAdapter>),
219 Compatibility(Arc<dyn ProviderAdapter>),
220}
221
222impl ResolvedDirectAdapter {
223 fn resolve(config: &ProxyConfig) -> Result<Option<Self>, ProxyError> {
224 let mut direct = config
225 .adapters
226 .iter()
227 .filter(|adapter| matches!(adapter.contract().protocol(), ProtocolContract::Codec(_)));
228 let Some(selected) = direct.next() else {
229 return if config.direct_anthropic.is_some() {
230 Err(ProxyError::DirectAdapterMismatch)
231 } else {
232 Ok(None)
233 };
234 };
235 if direct.next().is_some() {
236 return Err(direct_readiness_error("direct_adapter_count_invalid"));
237 }
238 if let Some(configured) = &config.direct_anthropic {
239 let dynamic: Arc<dyn ProviderAdapter> = configured.clone();
240 if !Arc::ptr_eq(&dynamic, selected) {
241 return Err(ProxyError::DirectAdapterMismatch);
242 }
243 Ok(Some(Self::Configured(Arc::clone(configured))))
244 } else {
245 Ok(Some(Self::Compatibility(Arc::clone(selected))))
246 }
247 }
248
249 fn provider(&self) -> &dyn ProviderAdapter {
250 match self {
251 Self::Configured(adapter) => adapter.as_ref(),
252 Self::Compatibility(adapter) => adapter.as_ref(),
253 }
254 }
255
256 fn configured(&self) -> Option<&AnthropicAdapter> {
257 match self {
258 Self::Configured(adapter) => Some(adapter),
259 Self::Compatibility(_) => None,
260 }
261 }
262}
263
264struct DirectRuntime {
265 adapter: ResolvedDirectAdapter,
266 client: DirectClient,
267 endpoint: Url,
268 session_policy: SessionPolicy,
269 registry: Option<SessionRegistry>,
270 resolver: Arc<dyn PrincipalResolver>,
271 listener_scope: ListenerScope,
272 allowed_versions: Arc<[String]>,
273 allowed_betas: Arc<[String]>,
274 codec_limits: CodecLimits,
275 client_config: DirectClientConfig,
276 ping_interval: Duration,
277 inspection_endpoint_codes: ProxyInspectionEndpointCodesV1,
278}
279
280enum DirectIngressSession {
281 Ephemeral,
282 Continuity(SessionLease),
283}
284
285struct PreparedDirectIngress {
286 outbound_headers: HeaderMap,
287 session: DirectIngressSession,
288}
289
290impl DirectRuntime {
291 fn from_config(config: &ProxyConfig) -> Result<Option<Self>, ProxyError> {
292 let Some(resolved) = ResolvedDirectAdapter::resolve(config)? else {
293 return Ok(None);
294 };
295 let adapter = resolved.provider();
296 let contract = adapter.contract();
297 let session_policy = contract
298 .session_policy()
299 .ok_or_else(|| direct_readiness_error("direct_session_policy_missing"))?;
300 let codec_limits = contract
301 .codec_limits()
302 .ok_or_else(|| direct_readiness_error("direct_codec_limits_missing"))?;
303 validate_direct_origin(adapter.upstream_base())
304 .map_err(|_| direct_readiness_error("direct_upstream_invalid"))?;
305 let mut endpoint = adapter.upstream_base().clone();
306 endpoint.set_path("/v1/messages");
307 let inspection_endpoint_codes =
308 ProxyInspectionEndpointCodesV1::from_validated_origin(&endpoint);
309
310 let explicit = resolved.configured();
311 let client_config = DirectClientConfig::from_anthropic(explicit);
312 let client = DirectClient::new(client_config)
313 .map_err(|_| direct_readiness_error("direct_client_invalid"))?;
314 let listener_scope = if config.bind.ip().is_loopback() {
315 ListenerScope::Loopback
316 } else {
317 ListenerScope::NonLoopback
318 };
319 let resolver: Arc<dyn PrincipalResolver> =
320 if let Some(resolver) = explicit.and_then(AnthropicAdapter::principal_resolver) {
321 Arc::clone(resolver)
322 } else if listener_scope == ListenerScope::Loopback {
323 Arc::new(
324 ProcessLocalLoopbackResolver::new()
325 .map_err(|_| direct_readiness_error("direct_principal_unavailable"))?,
326 )
327 } else {
328 return Err(direct_readiness_error("direct_principal_required"));
329 };
330 let registry = match session_policy {
331 SessionPolicy::EphemeralSingleRequest => None,
332 SessionPolicy::OpaqueHeaderContinuity(registry_config) => Some(
333 SessionRegistry::new(registry_config)
334 .map_err(|_| direct_readiness_error("direct_registry_invalid"))?,
335 ),
336 };
337 let allowed_versions = explicit.map_or_else(
338 || Arc::from([DEFAULT_ANTHROPIC_VERSION.to_string()]),
339 |value| Arc::from(value.allowed_versions()),
340 );
341 let allowed_betas = explicit.map_or_else(
342 || Arc::<[String]>::from([]),
343 |value| Arc::from(value.allowed_betas()),
344 );
345 let ping_interval =
346 explicit.map_or(Duration::from_secs(10), AnthropicAdapter::ping_interval);
347 Ok(Some(Self {
348 adapter: resolved,
349 client,
350 endpoint,
351 session_policy,
352 registry,
353 resolver,
354 listener_scope,
355 allowed_versions,
356 allowed_betas,
357 codec_limits,
358 client_config,
359 ping_interval,
360 inspection_endpoint_codes,
361 }))
362 }
363}
364
365fn direct_readiness_error(detail: &'static str) -> ProxyError {
366 ProxyError::DaemonConfig {
367 detail: detail.to_string(),
368 }
369}
370
371#[allow(dead_code)]
373pub(crate) struct DirectRequest {
374 url: Url,
375 headers: HeaderMap,
376 body: Vec<u8>,
377 expected_response: WireFormat,
378 inspection_parsed_facts: Option<InspectionParsedFactsV1>,
379 signed_or_encrypted_surface: bool,
380}
381
382#[allow(dead_code)]
383impl DirectRequest {
384 pub(crate) fn from_proved(
386 url: Url,
387 headers: HeaderMap,
388 proved_body: ProvedRequestBody,
389 expected_response: WireFormat,
390 ) -> Result<Self, DirectProxyError> {
391 validate_direct_url(&url)?;
392 validate_outbound_headers(&headers)?;
393 if !proved_body.provenance().final_buffer_verified() {
394 return Err(ProxyErrorCode::InvalidProvenance.error(ProxyErrorPhase::RequestTransform));
395 }
396 let (body, inspection_parsed_facts, signed_or_encrypted_surface) =
397 proved_body.into_inspection_parts();
398 Ok(Self {
399 url,
400 headers,
401 body,
402 expected_response,
403 inspection_parsed_facts,
404 signed_or_encrypted_surface,
405 })
406 }
407}
408
409struct PreparedDirectRequest<'session> {
414 transaction: SessionTransaction<'session>,
415 request: DirectRequest,
416 lifecycle: DirectTransactionState,
417}
418
419impl<'session> PreparedDirectRequest<'session> {
420 fn new(transaction: SessionTransaction<'session>, request: DirectRequest) -> Self {
421 Self {
422 transaction,
423 request,
424 lifecycle: DirectTransactionState::new(),
425 }
426 }
427
428 fn commit(mut self) -> Result<CommittedDirectRequest, SessionTransactionError> {
429 let snapshot = self.transaction.commit()?;
430 self.lifecycle.state = DirectIoState::CommittedNoIo;
431 Ok(CommittedDirectRequest {
432 snapshot,
433 request: self.request,
434 lifecycle: self.lifecycle,
435 })
436 }
437}
438
439struct CommittedDirectRequest {
443 snapshot: CommittedSessionSnapshot,
444 request: DirectRequest,
445 lifecycle: DirectTransactionState,
446}
447
448impl fmt::Debug for CommittedDirectRequest {
449 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
450 formatter
451 .debug_struct("CommittedDirectRequest")
452 .field("lifecycle", &self.lifecycle)
453 .finish_non_exhaustive()
454 }
455}
456
457struct CommittedDirectResponse {
458 snapshot: CommittedSessionSnapshot,
459 response: DirectResponse,
460 lifecycle: DirectTransactionState,
461}
462
463struct CommittedDirectSse {
464 snapshot: CommittedSessionSnapshot,
465 response: reqwest::Response,
466 head: DirectResponseHead,
467 lifecycle: DirectTransactionState,
468 remaining_total_timeout: Duration,
469}
470
471impl fmt::Debug for CommittedDirectResponse {
472 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
473 formatter
474 .debug_struct("CommittedDirectResponse")
475 .field("status", &self.response.status())
476 .field("format", &self.response.format())
477 .field("lifecycle", &self.lifecycle)
478 .finish_non_exhaustive()
479 }
480}
481
482pub struct DirectResponseHead {
484 status: StatusCode,
485 format: WireFormat,
486 headers: HeaderMap,
487}
488
489impl fmt::Debug for DirectResponseHead {
490 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
491 formatter
492 .debug_struct("DirectResponseHead")
493 .field("status", &self.status)
494 .field("format", &self.format)
495 .field("header_count", &self.headers.len())
496 .finish()
497 }
498}
499
500impl DirectResponseHead {
501 #[must_use]
502 pub const fn status(&self) -> StatusCode {
503 self.status
504 }
505
506 #[must_use]
507 pub const fn format(&self) -> WireFormat {
508 self.format
509 }
510
511 #[must_use]
512 pub const fn headers(&self) -> &HeaderMap {
513 &self.headers
514 }
515}
516
517#[allow(dead_code)]
519pub(crate) struct DirectResponse {
520 head: DirectResponseHead,
521 body: Vec<u8>,
522}
523
524impl fmt::Debug for DirectResponse {
525 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
526 formatter
527 .debug_struct("DirectResponse")
528 .field("status", &self.head.status)
529 .field("format", &self.head.format)
530 .field("body_bytes", &self.body.len())
531 .finish()
532 }
533}
534
535#[allow(dead_code)]
536impl DirectResponse {
537 #[must_use]
538 pub(crate) const fn status(&self) -> StatusCode {
539 self.head.status()
540 }
541
542 #[must_use]
543 pub(crate) const fn format(&self) -> WireFormat {
544 self.head.format()
545 }
546
547 #[must_use]
548 pub(crate) const fn headers(&self) -> &HeaderMap {
549 self.head.headers()
550 }
551
552 #[must_use]
553 pub(crate) fn body(&self) -> &[u8] {
554 &self.body
555 }
556
557 #[must_use]
558 pub(crate) fn into_body(self) -> Vec<u8> {
559 self.body
560 }
561
562 fn into_parts(self) -> (DirectResponseHead, Vec<u8>) {
563 (self.head, self.body)
564 }
565}
566
567#[derive(Clone)]
569#[allow(dead_code)]
570pub(crate) struct DirectClient {
571 client: Client,
572 config: DirectClientConfig,
573}
574
575#[allow(dead_code)]
576impl DirectClient {
577 pub(crate) fn new(config: DirectClientConfig) -> Result<Self, DirectProxyError> {
579 let client = hardened_client_builder(config).build().map_err(|_| {
580 ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::UpstreamConfiguration)
581 })?;
582 Ok(Self { client, config })
583 }
584
585 async fn execute(
587 &self,
588 mut committed: CommittedDirectRequest,
589 ) -> Result<CommittedDirectResponse, DirectProxyError> {
590 committed.lifecycle.require_committed_no_io()?;
591 if committed.request.body.len() > self.config.max_request_bytes {
592 return Err(
593 ProxyErrorCode::RequestBodyLimitExceeded.error(ProxyErrorPhase::RequestValidation)
594 );
595 }
596 committed.lifecycle.mark_io_attempted()?;
597 match tokio::time::timeout(
598 self.config.total_timeout,
599 self.execute_inner(committed.request),
600 )
601 .await
602 {
603 Ok(Ok(response)) => Ok(CommittedDirectResponse {
604 snapshot: committed.snapshot,
605 response,
606 lifecycle: committed.lifecycle,
607 }),
608 Ok(Err(error)) => Err(error),
609 Err(_) => Err(ProxyErrorCode::TotalTimeout.error(ProxyErrorPhase::UpstreamBody)),
610 }
611 }
612
613 async fn execute_sse(
614 &self,
615 mut committed: CommittedDirectRequest,
616 ) -> Result<CommittedDirectSse, DirectProxyError> {
617 committed.lifecycle.require_committed_no_io()?;
618 if committed.request.body.len() > self.config.max_request_bytes {
619 return Err(
620 ProxyErrorCode::RequestBodyLimitExceeded.error(ProxyErrorPhase::RequestValidation)
621 );
622 }
623 if committed.request.expected_response != WireFormat::Sse {
624 return Err(ProxyErrorCode::InvalidStateTransition.error(ProxyErrorPhase::Framework));
625 }
626 committed.lifecycle.mark_io_attempted()?;
627 let started = tokio::time::Instant::now();
628 let opened = tokio::time::timeout(
629 self.config.total_timeout,
630 self.open_sse_response(committed.request),
631 )
632 .await
633 .map_err(|_| ProxyErrorCode::TotalTimeout.error(ProxyErrorPhase::UpstreamHeaders))??;
634 let remaining_total_timeout = self.config.total_timeout.saturating_sub(started.elapsed());
635 Ok(CommittedDirectSse {
636 snapshot: committed.snapshot,
637 response: opened.0,
638 head: opened.1,
639 lifecycle: committed.lifecycle,
640 remaining_total_timeout,
641 })
642 }
643
644 async fn open_sse_response(
645 &self,
646 request: DirectRequest,
647 ) -> Result<(reqwest::Response, DirectResponseHead), DirectProxyError> {
648 let response = self
649 .client
650 .post(request.url)
651 .headers(request.headers)
652 .header(reqwest::header::ACCEPT_ENCODING, "identity")
653 .body(request.body)
654 .timeout(self.config.request_timeout)
655 .send()
656 .await
657 .map_err(map_reqwest_error)?;
658 let status = response.status();
659 if status.is_redirection() {
660 return Err(ProxyErrorCode::UpstreamRedirect.error(ProxyErrorPhase::UpstreamHeaders));
661 }
662 match validate_direct_response_head_with_config(
663 status,
664 response.headers(),
665 WireFormat::Sse,
666 self.config,
667 ) {
668 Ok(head) => Ok((response, head)),
669 Err(error) if is_mapped_upstream_status(error.code()) => {
670 discard_error_body(response, self.config.max_error_body_bytes).await;
671 Err(error)
672 }
673 Err(error) => Err(error),
674 }
675 }
676
677 async fn execute_inner(
678 &self,
679 request: DirectRequest,
680 ) -> Result<DirectResponse, DirectProxyError> {
681 let mut outbound = self
682 .client
683 .post(request.url)
684 .headers(request.headers)
685 .header(reqwest::header::ACCEPT_ENCODING, "identity")
686 .body(request.body);
687 outbound = outbound.timeout(self.config.request_timeout);
688 let response = outbound.send().await.map_err(map_reqwest_error)?;
689 let status = response.status();
690
691 if status.is_redirection() {
692 return Err(ProxyErrorCode::UpstreamRedirect.error(ProxyErrorPhase::UpstreamHeaders));
693 }
694
695 let head = validate_direct_response_head_with_config(
696 status,
697 response.headers(),
698 request.expected_response,
699 self.config,
700 );
701 match head {
702 Ok(head) => {
703 let body = collect_response_body(response, self.config.max_response_bytes).await?;
704 validate_declared_body(
705 request.expected_response,
706 &body,
707 self.config.max_response_bytes,
708 )?;
709 Ok(DirectResponse { head, body })
710 }
711 Err(error) if is_mapped_upstream_status(error.code()) => {
712 discard_error_body(response, self.config.max_error_body_bytes).await;
713 Err(error)
714 }
715 Err(error) => Err(error),
716 }
717 }
718}
719
720#[allow(dead_code)]
721fn hardened_client_builder(config: DirectClientConfig) -> ClientBuilder {
722 Client::builder()
723 .redirect(reqwest::redirect::Policy::none())
724 .no_proxy()
725 .referer(false)
726 .retry(reqwest::retry::never())
727 .no_gzip()
728 .no_brotli()
729 .no_zstd()
730 .no_deflate()
731 .connect_timeout(config.connect_timeout)
732 .timeout(config.request_timeout)
733}
734
735#[allow(dead_code)]
736fn map_reqwest_error(error: reqwest::Error) -> DirectProxyError {
737 let code = if error.is_timeout() && error.is_connect() {
738 ProxyErrorCode::ConnectTimeout
739 } else if error.is_timeout() {
740 ProxyErrorCode::RequestTimeout
741 } else if error.is_connect() {
742 ProxyErrorCode::UpstreamUnreachable
743 } else {
744 ProxyErrorCode::UpstreamProtocol
745 };
746 code.error(ProxyErrorPhase::UpstreamConnect)
747}
748
749#[allow(dead_code)]
750fn map_response_body_error(error: reqwest::Error) -> DirectProxyError {
751 let code = if error.is_timeout() {
752 ProxyErrorCode::RequestTimeout
753 } else {
754 ProxyErrorCode::UpstreamProtocol
755 };
756 code.error(ProxyErrorPhase::UpstreamBody)
757}
758
759#[allow(dead_code)]
760async fn collect_response_body(
761 response: reqwest::Response,
762 limit: usize,
763) -> Result<Vec<u8>, DirectProxyError> {
764 let mut stream = response.bytes_stream();
765 let mut body = Vec::new();
766 while let Some(chunk) = stream.next().await {
767 let chunk = chunk.map_err(map_response_body_error)?;
768 let accumulated = body.len().checked_add(chunk.len()).ok_or_else(|| {
769 ProxyErrorCode::ResponseBodyLimitExceeded.error(ProxyErrorPhase::UpstreamBody)
770 })?;
771 if accumulated > limit {
772 return Err(
773 ProxyErrorCode::ResponseBodyLimitExceeded.error(ProxyErrorPhase::UpstreamBody)
774 );
775 }
776 body.extend_from_slice(&chunk);
777 }
778 Ok(body)
779}
780
781#[allow(dead_code)]
782async fn discard_error_body(response: reqwest::Response, limit: usize) {
783 let mut stream = response.bytes_stream();
784 let mut discarded = 0_usize;
785 while discarded < limit {
786 let Some(next) = stream.next().await else {
787 break;
788 };
789 let Ok(chunk) = next else {
790 break;
791 };
792 let remaining = limit - discarded;
793 discarded += chunk.len().min(remaining);
794 if chunk.len() > remaining {
795 break;
796 }
797 }
798}
799
800#[allow(dead_code)]
801fn validate_direct_url(url: &Url) -> Result<(), DirectProxyError> {
802 let safe_authority = url.username().is_empty()
803 && url.password().is_none()
804 && url.query().is_none()
805 && url.fragment().is_none()
806 && url.host().is_some();
807 let safe_scheme = match (url.scheme(), url.host()) {
808 ("https", Some(_)) => true,
809 ("http", Some(Host::Ipv4(address))) => IpAddr::V4(address).is_loopback(),
810 ("http", Some(Host::Ipv6(address))) => IpAddr::V6(address).is_loopback(),
811 _ => false,
812 };
813 if safe_authority && safe_scheme {
814 Ok(())
815 } else {
816 Err(ProxyErrorCode::InvalidUpstreamUrl.error(ProxyErrorPhase::UpstreamConfiguration))
817 }
818}
819
820fn validate_direct_origin(url: &Url) -> Result<(), DirectProxyError> {
821 if url.path() != "/" {
822 return Err(
823 ProxyErrorCode::InvalidUpstreamUrl.error(ProxyErrorPhase::UpstreamConfiguration)
824 );
825 }
826 validate_direct_url(url)
827}
828
829#[allow(dead_code)]
830fn validate_outbound_headers(headers: &HeaderMap) -> Result<(), DirectProxyError> {
831 const ALLOWED: &[&str] = &[
832 "content-type",
833 "x-api-key",
834 "anthropic-version",
835 "anthropic-beta",
836 ];
837 validate_header_space(headers, DirectClientConfig::default())?;
838 let mut singleton_names = HashSet::new();
839 for (name, value) in headers {
840 if !ALLOWED.contains(&name.as_str()) {
841 return Err(ProxyErrorCode::HeaderRejected
842 .error(ProxyErrorPhase::RequestValidation)
843 .with_header_rejection(name.as_str(), HeaderRejectionReason::NameNotAllowlisted));
844 }
845 if value.is_empty() {
846 return Err(ProxyErrorCode::HeaderRejected
847 .error(ProxyErrorPhase::RequestValidation)
848 .with_header_rejection(name.as_str(), HeaderRejectionReason::EmptyValue));
849 }
850 if matches!(
851 name.as_str(),
852 "x-api-key" | "anthropic-version" | "anthropic-beta"
853 ) && !singleton_names.insert(name.as_str())
854 {
855 return Err(ProxyErrorCode::HeaderRejected
856 .error(ProxyErrorPhase::RequestValidation)
857 .with_header_rejection(name.as_str(), HeaderRejectionReason::DuplicateSingleton));
858 }
859 }
860 Ok(())
861}
862
863pub fn classify_upstream_status(status: StatusCode) -> Result<(), DirectProxyError> {
865 let code = if status.is_success() {
866 return Ok(());
867 } else if status.is_redirection() {
868 ProxyErrorCode::UpstreamRedirect
869 } else {
870 match status.as_u16() {
871 400 => ProxyErrorCode::UpstreamBadRequest,
872 401 => ProxyErrorCode::UpstreamUnauthorized,
873 403 => ProxyErrorCode::UpstreamForbidden,
874 404 => ProxyErrorCode::UpstreamNotFound,
875 409 => ProxyErrorCode::UpstreamConflict,
876 413 => ProxyErrorCode::UpstreamPayloadTooLarge,
877 429 => ProxyErrorCode::UpstreamRateLimited,
878 502 | 503 | 504 | 529 => ProxyErrorCode::UpstreamUnavailable,
879 value if (400..=499).contains(&value) => ProxyErrorCode::UpstreamClientFailure,
880 500..=599 => ProxyErrorCode::UpstreamServerFailure,
881 _ => ProxyErrorCode::UpstreamServerFailure,
882 }
883 };
884 Err(code.error(ProxyErrorPhase::UpstreamHeaders))
885}
886
887#[allow(dead_code)]
888fn is_mapped_upstream_status(code: ProxyErrorCode) -> bool {
889 matches!(
890 code,
891 ProxyErrorCode::UpstreamBadRequest
892 | ProxyErrorCode::UpstreamUnauthorized
893 | ProxyErrorCode::UpstreamForbidden
894 | ProxyErrorCode::UpstreamNotFound
895 | ProxyErrorCode::UpstreamConflict
896 | ProxyErrorCode::UpstreamPayloadTooLarge
897 | ProxyErrorCode::UpstreamRateLimited
898 | ProxyErrorCode::UpstreamClientFailure
899 | ProxyErrorCode::UpstreamUnavailable
900 | ProxyErrorCode::UpstreamServerFailure
901 )
902}
903
904pub fn validate_direct_response_head(
909 status: StatusCode,
910 headers: &HeaderMap,
911 expected: WireFormat,
912) -> Result<DirectResponseHead, DirectProxyError> {
913 validate_direct_response_head_with_config(
914 status,
915 headers,
916 expected,
917 DirectClientConfig::default(),
918 )
919}
920
921fn validate_direct_response_head_with_config(
922 status: StatusCode,
923 headers: &HeaderMap,
924 expected: WireFormat,
925 config: DirectClientConfig,
926) -> Result<DirectResponseHead, DirectProxyError> {
927 if status.is_redirection() {
928 return Err(ProxyErrorCode::UpstreamRedirect.error(ProxyErrorPhase::UpstreamHeaders));
929 }
930 validate_header_space(headers, config)?;
931 validate_response_framing(headers)?;
932
933 if let Err(error) = classify_upstream_status(status) {
934 let retry_after_seconds = parse_retry_after(headers)?;
935 return Err(error.with_retry_after_seconds(retry_after_seconds));
936 }
937
938 let content_type = one_header(headers, reqwest::header::CONTENT_TYPE)?;
939 let actual = match content_type {
940 None if expected == WireFormat::Empty => WireFormat::Empty,
941 Some(value) => parse_wire_content_type(value)?,
942 None => {
943 return Err(ProxyErrorCode::InvalidUpstreamResponseFormat
944 .error(ProxyErrorPhase::ResponseValidation));
945 }
946 };
947 if actual != expected || (expected == WireFormat::Empty && content_type.is_some()) {
948 return Err(ProxyErrorCode::InvalidUpstreamResponseFormat
949 .error(ProxyErrorPhase::ResponseValidation));
950 }
951 let mut minimal = HeaderMap::new();
952 if expected != WireFormat::Empty {
953 minimal.insert(
954 reqwest::header::CONTENT_TYPE,
955 HeaderValue::from_static(canonical_content_type(expected)),
956 );
957 }
958 Ok(DirectResponseHead {
959 status,
960 format: expected,
961 headers: minimal,
962 })
963}
964
965pub fn validate_declared_body(
967 format: WireFormat,
968 body: &[u8],
969 limit: usize,
970) -> Result<(), DirectProxyError> {
971 if body.len() > limit {
972 return Err(
973 ProxyErrorCode::ResponseBodyLimitExceeded.error(ProxyErrorPhase::ResponseValidation)
974 );
975 }
976 match format {
977 WireFormat::Empty if body.is_empty() => Ok(()),
978 WireFormat::Empty => Err(ProxyErrorCode::InvalidUpstreamResponseFormat
979 .error(ProxyErrorPhase::ResponseValidation)),
980 WireFormat::Json => validate_unique_json(body),
981 WireFormat::Ndjson => {
982 let text = std::str::from_utf8(body).map_err(|_| {
983 ProxyErrorCode::InvalidUpstreamResponseFormat
984 .error(ProxyErrorPhase::ResponseValidation)
985 })?;
986 for line in text.lines().map(str::trim).filter(|line| !line.is_empty()) {
987 validate_unique_json(line.as_bytes())?;
988 }
989 Ok(())
990 }
991 WireFormat::Utf8Text => std::str::from_utf8(body).map(|_| ()).map_err(|_| {
992 ProxyErrorCode::InvalidUpstreamResponseFormat.error(ProxyErrorPhase::ResponseValidation)
993 }),
994 WireFormat::Sse => validate_sse_envelope(body),
995 }
996}
997
998fn validate_unique_json(body: &[u8]) -> Result<(), DirectProxyError> {
999 let parsed: UniqueJson = serde_json::from_slice(body).map_err(|_| {
1000 ProxyErrorCode::InvalidUpstreamResponseFormat.error(ProxyErrorPhase::ResponseValidation)
1001 })?;
1002 if parsed.duplicate {
1003 Err(ProxyErrorCode::DuplicateObjectKey.error(ProxyErrorPhase::ResponseValidation))
1004 } else {
1005 Ok(())
1006 }
1007}
1008
1009fn validate_sse_envelope(body: &[u8]) -> Result<(), DirectProxyError> {
1010 if body.starts_with(&[0xef, 0xbb, 0xbf]) || body.contains(&0) {
1011 return Err(ProxyErrorCode::InvalidSseLifecycle.error(ProxyErrorPhase::ResponseValidation));
1012 }
1013 std::str::from_utf8(body).map_err(|_| {
1014 ProxyErrorCode::InvalidUpstreamResponseFormat.error(ProxyErrorPhase::ResponseValidation)
1015 })?;
1016 if body.ends_with(b"\n\n") || body.ends_with(b"\r\r") || body.ends_with(b"\r\n\r\n") {
1017 Ok(())
1018 } else {
1019 Err(ProxyErrorCode::InvalidSseLifecycle.error(ProxyErrorPhase::ResponseValidation))
1020 }
1021}
1022
1023fn validate_header_space(
1024 headers: &HeaderMap,
1025 config: DirectClientConfig,
1026) -> Result<(), DirectProxyError> {
1027 let mut count = 0_usize;
1028 let mut aggregate = 0_usize;
1029 for (name, value) in headers {
1030 count = count.checked_add(1).ok_or_else(header_limit_error)?;
1031 aggregate = aggregate
1032 .checked_add(name.as_str().len())
1033 .and_then(|value_so_far| value_so_far.checked_add(value.as_bytes().len()))
1034 .ok_or_else(header_limit_error)?;
1035 if count > config.max_headers
1036 || name.as_str().len() > config.max_header_name_bytes
1037 || value.as_bytes().len() > config.max_header_value_bytes
1038 || aggregate > config.max_header_bytes
1039 {
1040 return Err(header_limit_error());
1041 }
1042 }
1043 Ok(())
1044}
1045
1046fn header_limit_error() -> DirectProxyError {
1047 ProxyErrorCode::HeaderLimitExceeded.error(ProxyErrorPhase::UpstreamHeaders)
1048}
1049
1050fn one_header(
1051 headers: &HeaderMap,
1052 name: HeaderName,
1053) -> Result<Option<&HeaderValue>, DirectProxyError> {
1054 let mut values = headers.get_all(name).iter();
1055 let first = values.next();
1056 if values.next().is_some() {
1057 return Err(ProxyErrorCode::InvalidUpstreamHeader.error(ProxyErrorPhase::UpstreamHeaders));
1058 }
1059 Ok(first)
1060}
1061
1062fn validate_response_framing(headers: &HeaderMap) -> Result<(), DirectProxyError> {
1063 let content_length = one_header(headers, reqwest::header::CONTENT_LENGTH)?;
1064 let transfer_encoding = one_header(headers, reqwest::header::TRANSFER_ENCODING)?;
1065 if content_length.is_some() && transfer_encoding.is_some() {
1066 return Err(ProxyErrorCode::InvalidFraming.error(ProxyErrorPhase::UpstreamHeaders));
1067 }
1068 if let Some(value) = content_length {
1069 let raw = value.as_bytes();
1070 if raw.is_empty() || !raw.iter().all(u8::is_ascii_digit) {
1071 return Err(ProxyErrorCode::InvalidFraming.error(ProxyErrorPhase::UpstreamHeaders));
1072 }
1073 }
1074 if let Some(value) = transfer_encoding {
1075 if !value.as_bytes().eq_ignore_ascii_case(b"chunked") {
1076 return Err(ProxyErrorCode::InvalidFraming.error(ProxyErrorPhase::UpstreamHeaders));
1077 }
1078 }
1079 if let Some(value) = one_header(headers, reqwest::header::CONTENT_ENCODING)? {
1080 if !trim_ascii(value.as_bytes()).eq_ignore_ascii_case(b"identity") {
1081 return Err(
1082 ProxyErrorCode::UnsupportedContentEncoding.error(ProxyErrorPhase::UpstreamHeaders)
1083 );
1084 }
1085 }
1086 Ok(())
1087}
1088
1089fn parse_retry_after(headers: &HeaderMap) -> Result<Option<u32>, DirectProxyError> {
1090 let Some(value) = one_header(headers, reqwest::header::RETRY_AFTER)? else {
1091 return Ok(None);
1092 };
1093 let bytes = value.as_bytes();
1094 if bytes.is_empty() || !bytes.iter().all(u8::is_ascii_digit) {
1095 return Err(ProxyErrorCode::InvalidUpstreamHeader.error(ProxyErrorPhase::UpstreamHeaders));
1096 }
1097 let seconds = std::str::from_utf8(bytes)
1098 .ok()
1099 .and_then(|text| text.parse::<u32>().ok())
1100 .ok_or_else(|| {
1101 ProxyErrorCode::InvalidUpstreamHeader.error(ProxyErrorPhase::UpstreamHeaders)
1102 })?;
1103 Ok(Some(seconds))
1104}
1105
1106fn parse_wire_content_type(value: &HeaderValue) -> Result<WireFormat, DirectProxyError> {
1107 let raw = value.as_bytes();
1108 if !raw.is_ascii() {
1109 return Err(ProxyErrorCode::InvalidUpstreamResponseFormat
1110 .error(ProxyErrorPhase::ResponseValidation));
1111 }
1112 let mut parts = raw.split(|byte| *byte == b';');
1113 let media_type = trim_ascii(parts.next().unwrap_or_default());
1114 let format = if media_type.eq_ignore_ascii_case(b"application/json") {
1115 WireFormat::Json
1116 } else if media_type.eq_ignore_ascii_case(b"application/x-ndjson") {
1117 WireFormat::Ndjson
1118 } else if media_type.eq_ignore_ascii_case(b"text/plain") {
1119 WireFormat::Utf8Text
1120 } else if media_type.eq_ignore_ascii_case(b"text/event-stream") {
1121 WireFormat::Sse
1122 } else {
1123 return Err(ProxyErrorCode::InvalidUpstreamResponseFormat
1124 .error(ProxyErrorPhase::ResponseValidation));
1125 };
1126 for parameter in parts {
1127 validate_mime_parameter(trim_ascii(parameter))?;
1128 }
1129 Ok(format)
1130}
1131
1132fn validate_mime_parameter(parameter: &[u8]) -> Result<(), DirectProxyError> {
1133 let Some(separator) = parameter.iter().position(|byte| *byte == b'=') else {
1134 return Err(ProxyErrorCode::InvalidUpstreamResponseFormat
1135 .error(ProxyErrorPhase::ResponseValidation));
1136 };
1137 let name = trim_ascii(¶meter[..separator]);
1138 let value = trim_ascii(¶meter[separator + 1..]);
1139 let valid_name = !name.is_empty() && name.iter().all(|byte| is_mime_token(*byte));
1140 let valid_value = if value.len() >= 2 && value[0] == b'"' && value[value.len() - 1] == b'"' {
1141 value[1..value.len() - 1]
1142 .iter()
1143 .all(|byte| matches!(*byte, b'\t' | 0x20..=0x7e) && *byte != b'"' && *byte != b'\\')
1144 } else {
1145 !value.is_empty() && value.iter().all(|byte| is_mime_token(*byte))
1146 };
1147 if valid_name && valid_value {
1148 Ok(())
1149 } else {
1150 Err(ProxyErrorCode::InvalidUpstreamResponseFormat
1151 .error(ProxyErrorPhase::ResponseValidation))
1152 }
1153}
1154
1155fn is_mime_token(byte: u8) -> bool {
1156 byte.is_ascii_alphanumeric()
1157 || matches!(
1158 byte,
1159 b'!' | b'#'
1160 | b'$'
1161 | b'%'
1162 | b'&'
1163 | b'\''
1164 | b'*'
1165 | b'+'
1166 | b'-'
1167 | b'.'
1168 | b'^'
1169 | b'_'
1170 | b'`'
1171 | b'|'
1172 | b'~'
1173 )
1174}
1175
1176fn trim_ascii(mut value: &[u8]) -> &[u8] {
1177 while matches!(value.first(), Some(b' ' | b'\t')) {
1178 value = &value[1..];
1179 }
1180 while matches!(value.last(), Some(b' ' | b'\t')) {
1181 value = &value[..value.len() - 1];
1182 }
1183 value
1184}
1185
1186fn canonical_content_type(format: WireFormat) -> &'static str {
1187 match format {
1188 WireFormat::Empty => "application/octet-stream",
1189 WireFormat::Json => "application/json",
1190 WireFormat::Ndjson => "application/x-ndjson",
1191 WireFormat::Utf8Text => "text/plain; charset=utf-8",
1192 WireFormat::Sse => "text/event-stream",
1193 }
1194}
1195
1196struct UniqueJson {
1197 value: Value,
1198 duplicate: bool,
1199}
1200
1201impl<'de> Deserialize<'de> for UniqueJson {
1202 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
1203 where
1204 D: Deserializer<'de>,
1205 {
1206 deserializer.deserialize_any(UniqueJsonVisitor)
1207 }
1208}
1209
1210struct UniqueJsonVisitor;
1211
1212impl<'de> Visitor<'de> for UniqueJsonVisitor {
1213 type Value = UniqueJson;
1214
1215 fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1216 formatter.write_str("one bounded JSON value")
1217 }
1218
1219 fn visit_bool<E>(self, value: bool) -> Result<Self::Value, E> {
1220 Ok(unique(Value::Bool(value)))
1221 }
1222
1223 fn visit_i64<E>(self, value: i64) -> Result<Self::Value, E> {
1224 Ok(unique(Value::Number(Number::from(value))))
1225 }
1226
1227 fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E> {
1228 Ok(unique(Value::Number(Number::from(value))))
1229 }
1230
1231 fn visit_f64<E>(self, value: f64) -> Result<Self::Value, E>
1232 where
1233 E: serde::de::Error,
1234 {
1235 Number::from_f64(value)
1236 .map(Value::Number)
1237 .map(unique)
1238 .ok_or_else(|| E::custom("non-finite number"))
1239 }
1240
1241 fn visit_str<E>(self, value: &str) -> Result<Self::Value, E> {
1242 Ok(unique(Value::String(value.to_owned())))
1243 }
1244
1245 fn visit_string<E>(self, value: String) -> Result<Self::Value, E> {
1246 Ok(unique(Value::String(value)))
1247 }
1248
1249 fn visit_none<E>(self) -> Result<Self::Value, E> {
1250 Ok(unique(Value::Null))
1251 }
1252
1253 fn visit_unit<E>(self) -> Result<Self::Value, E> {
1254 Ok(unique(Value::Null))
1255 }
1256
1257 fn visit_seq<A>(self, mut sequence: A) -> Result<Self::Value, A::Error>
1258 where
1259 A: SeqAccess<'de>,
1260 {
1261 let mut values = Vec::new();
1262 let mut duplicate = false;
1263 while let Some(value) = sequence.next_element::<UniqueJson>()? {
1264 duplicate |= value.duplicate;
1265 values.push(value.value);
1266 }
1267 Ok(UniqueJson {
1268 value: Value::Array(values),
1269 duplicate,
1270 })
1271 }
1272
1273 fn visit_map<A>(self, mut object: A) -> Result<Self::Value, A::Error>
1274 where
1275 A: MapAccess<'de>,
1276 {
1277 let mut values = Map::new();
1278 let mut seen = HashSet::new();
1279 let mut duplicate = false;
1280 while let Some(key) = object.next_key::<String>()? {
1281 duplicate |= !seen.insert(key.clone());
1282 let value = object.next_value::<UniqueJson>()?;
1283 duplicate |= value.duplicate;
1284 values.insert(key, value.value);
1285 }
1286 Ok(UniqueJson {
1287 value: Value::Object(values),
1288 duplicate,
1289 })
1290 }
1291}
1292
1293fn unique(value: Value) -> UniqueJson {
1294 UniqueJson {
1295 value,
1296 duplicate: false,
1297 }
1298}
1299
1300pub fn parse_request_stream_format(body: &[u8]) -> Result<WireFormat, DirectProxyError> {
1302 let parsed: UniqueJson = serde_json::from_slice(body).map_err(|_| {
1303 ProxyErrorCode::InvalidRequestFormat.error(ProxyErrorPhase::RequestValidation)
1304 })?;
1305 if parsed.duplicate {
1306 return Err(ProxyErrorCode::DuplicateObjectKey.error(ProxyErrorPhase::RequestValidation));
1307 }
1308 let Value::Object(object) = parsed.value else {
1309 return Err(ProxyErrorCode::InvalidRequestFormat.error(ProxyErrorPhase::RequestValidation));
1310 };
1311 match object.get("stream") {
1312 None | Some(Value::Bool(false)) => Ok(WireFormat::Json),
1313 Some(Value::Bool(true)) => Ok(WireFormat::Sse),
1314 _ => Err(ProxyErrorCode::InvalidRequestFormat.error(ProxyErrorPhase::RequestValidation)),
1315 }
1316}
1317
1318struct PipelineRequestPseudonymizer<'pipeline, 'context, 'transaction, 'session> {
1324 pipeline: &'pipeline Pipeline,
1325 transaction: &'transaction mut SessionTransaction<'session>,
1326 locale_chain: &'context [LocaleTag],
1327 dictionaries: &'context DictionaryBundle,
1328 refusal: Option<BoundaryRefusal>,
1330}
1331
1332impl<'pipeline, 'context, 'transaction, 'session>
1333 PipelineRequestPseudonymizer<'pipeline, 'context, 'transaction, 'session>
1334{
1335 fn new(
1336 pipeline: &'pipeline Pipeline,
1337 locale_chain: &'context [LocaleTag],
1338 dictionaries: &'context DictionaryBundle,
1339 transaction: &'transaction mut SessionTransaction<'session>,
1340 ) -> Self {
1341 Self {
1342 pipeline,
1343 transaction,
1344 locale_chain,
1345 dictionaries,
1346 refusal: None,
1347 }
1348 }
1349
1350 fn protect_unstaged_segment(
1352 &mut self,
1353 input: &str,
1354 _prefix_cache_write_mode: PrefixCacheWriteMode,
1355 ) -> Result<String, CodecErrorCode> {
1356 let result = self.pipeline.resolve_boundary_text_transaction(
1359 self.transaction,
1360 input,
1361 self.locale_chain,
1362 self.dictionaries,
1363 );
1364 self.refused(result)
1365 }
1366
1367 fn refused<T>(&mut self, result: Result<T, BoundaryRefusal>) -> Result<T, CodecErrorCode> {
1368 result.map_err(|refusal| {
1369 self.refusal = Some(refusal);
1370 CodecErrorCode::ProtectionFailedClosed
1371 })
1372 }
1373}
1374
1375impl RequestPseudonymizer for PipelineRequestPseudonymizer<'_, '_, '_, '_> {
1376 fn protect(&mut self, input: &str) -> Result<String, CodecErrorCode> {
1377 self.protect_with_prefix_cache_write_mode(input, PrefixCacheWriteMode::Allow)
1378 }
1379
1380 fn protect_with_prefix_cache_write_mode(
1381 &mut self,
1382 input: &str,
1383 prefix_cache_write_mode: PrefixCacheWriteMode,
1384 ) -> Result<String, CodecErrorCode> {
1385 let output = self.protect_preserving_tokens(input, prefix_cache_write_mode)?;
1386 let admitted = self.pipeline.admit_boundary_text_transaction(
1387 self.transaction,
1388 &output,
1389 self.locale_chain,
1390 self.dictionaries,
1391 );
1392 self.refused(admitted)?;
1393 Ok(output)
1394 }
1395
1396 fn validate_token_shapes(&mut self, input: &str) -> Result<(), CodecErrorCode> {
1397 self.transaction
1398 .validate_token_shapes(input)
1399 .map_err(|_| CodecErrorCode::RestoreFailedClosed)
1400 }
1401}
1402
1403impl PipelineRequestPseudonymizer<'_, '_, '_, '_> {
1404 fn protect_preserving_tokens(
1405 &mut self,
1406 input: &str,
1407 prefix_cache_write_mode: PrefixCacheWriteMode,
1408 ) -> Result<String, CodecErrorCode> {
1409 self.validate_token_shapes(input)?;
1410 let mut tokens = self.transaction.tokens();
1411 if tokens.is_empty() {
1412 return self.protect_unstaged_segment(input, prefix_cache_write_mode);
1413 }
1414 tokens.sort_unstable_by_key(|token| std::cmp::Reverse(token.len()));
1415
1416 let mut output = String::with_capacity(input.len());
1417 let mut cursor = 0;
1418 while cursor < input.len() {
1419 let next = tokens
1420 .iter()
1421 .filter_map(|token| {
1422 input[cursor..]
1423 .find(token)
1424 .map(|offset| (cursor + offset, token.as_str()))
1425 })
1426 .min_by(|left, right| {
1427 left.0
1428 .cmp(&right.0)
1429 .then_with(|| right.1.len().cmp(&left.1.len()))
1430 });
1431 let Some((start, token)) = next else {
1432 output.push_str(
1433 &self.protect_unstaged_segment(&input[cursor..], prefix_cache_write_mode)?,
1434 );
1435 break;
1436 };
1437 if start > cursor {
1438 output.push_str(
1439 &self
1440 .protect_unstaged_segment(&input[cursor..start], prefix_cache_write_mode)?,
1441 );
1442 }
1443 output.push_str(token);
1444 cursor = start + token.len();
1445 }
1446 Ok(output)
1447 }
1448}
1449
1450struct PipelineResponseResidualValidator<'pipeline, 'context> {
1451 pipeline: &'pipeline Pipeline,
1452 locale_chain: &'context [LocaleTag],
1453 dictionaries: &'context DictionaryBundle,
1454 validation_session: Session,
1455}
1456
1457impl<'pipeline, 'context> PipelineResponseResidualValidator<'pipeline, 'context> {
1458 fn new(
1459 pipeline: &'pipeline Pipeline,
1460 locale_chain: &'context [LocaleTag],
1461 dictionaries: &'context DictionaryBundle,
1462 ) -> Result<Self, DirectProxyError> {
1463 let validation_session = Session::new(Scope::Ephemeral).map_err(|_| {
1464 ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::ResponseValidation)
1465 })?;
1466 Ok(Self {
1467 pipeline,
1468 locale_chain,
1469 dictionaries,
1470 validation_session,
1471 })
1472 }
1473}
1474
1475impl ResponseResidualValidator for PipelineResponseResidualValidator<'_, '_> {
1476 fn validate(
1477 &mut self,
1478 value: &str,
1479 authorized_output_ranges: &[std::ops::Range<usize>],
1480 ) -> Result<(), CodecErrorCode> {
1481 let (clean, spans, _) = self
1482 .pipeline
1483 .clean_with_safety_net_detect_context(
1484 &self.validation_session,
1485 RawDocument::Text(value.to_owned()),
1486 self.locale_chain,
1487 self.dictionaries,
1488 )
1489 .map_err(|_| CodecErrorCode::ProviderOriginPii)?;
1490 let CleanDocument::Text(clean_text) = clean else {
1491 return Err(CodecErrorCode::ProviderOriginPii);
1492 };
1493 if !manifest_accounts_for_every_change(value, &clean_text, &spans) {
1501 return Err(CodecErrorCode::ProviderOriginPii);
1502 }
1503 for span in spans {
1504 if !authorized_output_ranges.iter().any(|authorized| {
1505 authorized.start <= span.raw_span.start && authorized.end >= span.raw_span.end
1506 }) {
1507 return Err(CodecErrorCode::ProviderOriginPii);
1508 }
1509 }
1510 Ok(())
1511 }
1512}
1513
1514fn manifest_accounts_for_every_change(raw: &str, clean: &str, spans: &[EmittedTokenSpan]) -> bool {
1523 let mut ordered: Vec<&EmittedTokenSpan> = spans.iter().collect();
1524 ordered.sort_by_key(|span| span.clean_span.start);
1525 let mut raw_cursor = 0usize;
1526 let mut clean_cursor = 0usize;
1527 for span in ordered {
1528 if span.raw_span.start < raw_cursor || span.clean_span.start < clean_cursor {
1529 return false;
1530 }
1531 let (Some(raw_gap), Some(clean_gap)) = (
1532 raw.get(raw_cursor..span.raw_span.start),
1533 clean.get(clean_cursor..span.clean_span.start),
1534 ) else {
1535 return false;
1536 };
1537 if raw_gap != clean_gap {
1538 return false;
1539 }
1540 if raw.get(span.raw_span.clone()).is_none() || clean.get(span.clean_span.clone()).is_none()
1543 {
1544 return false;
1545 }
1546 raw_cursor = span.raw_span.end;
1547 clean_cursor = span.clean_span.end;
1548 }
1549 matches!(
1550 (raw.get(raw_cursor..), clean.get(clean_cursor..)),
1551 (Some(raw_tail), Some(clean_tail)) if raw_tail == clean_tail
1552 )
1553}
1554
1555#[allow(clippy::too_many_arguments)]
1556fn prepare_and_commit_direct_request(
1557 pipeline: &Pipeline,
1558 locale_chain: &[LocaleTag],
1559 dictionaries: &DictionaryBundle,
1560 session: &Session,
1561 codec: &dyn BodyCodec,
1562 original_body: &[u8],
1563 endpoint: &Url,
1564 headers: &HeaderMap,
1565 expected_response: WireFormat,
1566 limits: CodecLimits,
1567 provider_request_projection_selected: bool,
1568 max_request_bytes: usize,
1569) -> Result<CommittedDirectRequest, DirectProxyError> {
1570 prepare_and_commit_direct_request_with_hook(
1571 pipeline,
1572 locale_chain,
1573 dictionaries,
1574 session,
1575 codec,
1576 original_body,
1577 endpoint,
1578 headers,
1579 expected_response,
1580 limits,
1581 provider_request_projection_selected,
1582 max_request_bytes,
1583 |_, _, _| Ok(()),
1584 )
1585}
1586
1587#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1588enum SessionCommitFailureKind {
1589 GenerationConflict,
1590 NonConflict,
1591}
1592
1593fn classify_session_commit_failure(error: &SessionTransactionError) -> SessionCommitFailureKind {
1594 match error {
1595 SessionTransactionError::GenerationConflict => SessionCommitFailureKind::GenerationConflict,
1596 _ => SessionCommitFailureKind::NonConflict,
1597 }
1598}
1599
1600fn map_session_commit_failure(kind: SessionCommitFailureKind) -> DirectProxyError {
1601 match kind {
1602 SessionCommitFailureKind::GenerationConflict => {
1603 ProxyErrorCode::SessionGenerationConflict.error(ProxyErrorPhase::Session)
1604 }
1605 SessionCommitFailureKind::NonConflict => {
1606 ProxyErrorCode::SessionCommitFailure.error(ProxyErrorPhase::Session)
1607 }
1608 }
1609}
1610
1611#[allow(clippy::too_many_arguments)]
1612fn prepare_and_commit_direct_request_with_hook(
1613 pipeline: &Pipeline,
1614 locale_chain: &[LocaleTag],
1615 dictionaries: &DictionaryBundle,
1616 session: &Session,
1617 codec: &dyn BodyCodec,
1618 original_body: &[u8],
1619 endpoint: &Url,
1620 headers: &HeaderMap,
1621 expected_response: WireFormat,
1622 limits: CodecLimits,
1623 provider_request_projection_selected: bool,
1624 max_request_bytes: usize,
1625 mut before_commit: impl FnMut(u8, &Session, &SessionTransaction<'_>) -> Result<(), DirectProxyError>,
1626) -> Result<CommittedDirectRequest, DirectProxyError> {
1627 for attempt in 0..=1 {
1628 let mut transaction = session.begin_transaction();
1629 let proved_body = {
1630 let mut pseudonymizer = PipelineRequestPseudonymizer::new(
1631 pipeline,
1632 locale_chain,
1633 dictionaries,
1634 &mut transaction,
1635 );
1636 let protected = {
1637 let mut context =
1638 RequestTransformContext::new(&mut pseudonymizer, WireFormat::Json, limits);
1639 if provider_request_projection_selected {
1640 context.request_inspection_projection();
1641 }
1642 codec.protect_request(original_body, &mut context)
1643 };
1644 protected.map_err(|error| match pseudonymizer.refusal.take() {
1645 Some(refusal) if error.code() == CodecErrorCode::ProtectionFailedClosed => {
1646 ProxyErrorCode::ProtectionRefused
1647 .error(ProxyErrorPhase::RequestTransform)
1648 .with_refusal(refusal)
1649 }
1650 _ => map_codec_error(error),
1651 })?
1652 };
1653 let request = DirectRequest::from_proved(
1654 endpoint.clone(),
1655 headers.clone(),
1656 proved_body,
1657 expected_response,
1658 )?;
1659 if request.body.len() > max_request_bytes {
1660 return Err(
1661 ProxyErrorCode::RequestBodyLimitExceeded.error(ProxyErrorPhase::RequestValidation)
1662 );
1663 }
1664 let prepared = PreparedDirectRequest::new(transaction, request);
1665 before_commit(attempt, session, &prepared.transaction)?;
1666 match prepared.commit() {
1667 Ok(committed) => return Ok(committed),
1668 Err(error) => match classify_session_commit_failure(&error) {
1669 SessionCommitFailureKind::GenerationConflict if attempt == 0 => {}
1670 kind => return Err(map_session_commit_failure(kind)),
1671 },
1672 }
1673 }
1674 Err(ProxyErrorCode::SessionGenerationConflict.error(ProxyErrorPhase::Session))
1675}
1676
1677#[derive(Debug, Eq, PartialEq)]
1679#[allow(dead_code)]
1680pub(crate) enum DirectIoState {
1681 Prepared,
1682 CommittedNoIo,
1683 IoAttempted,
1684 ResponseComplete,
1685}
1686
1687#[derive(Debug, Eq, PartialEq)]
1689#[allow(dead_code)]
1690pub(crate) struct DirectTransactionState {
1691 state: DirectIoState,
1692 generation_retries: u8,
1693}
1694
1695#[allow(dead_code)]
1696impl DirectTransactionState {
1697 #[must_use]
1698 pub(crate) const fn new() -> Self {
1699 Self {
1700 state: DirectIoState::Prepared,
1701 generation_retries: 0,
1702 }
1703 }
1704
1705 #[must_use]
1706 pub(crate) const fn state(&self) -> &DirectIoState {
1707 &self.state
1708 }
1709
1710 #[must_use]
1711 pub(crate) const fn generation_retries(&self) -> u8 {
1712 self.generation_retries
1713 }
1714
1715 #[must_use]
1716 pub(crate) const fn retains_committed_mapping(&self) -> bool {
1717 !matches!(self.state, DirectIoState::Prepared)
1718 }
1719
1720 pub(crate) fn retry_generation_conflict(&mut self) -> Result<(), DirectProxyError> {
1721 if self.state != DirectIoState::Prepared || self.generation_retries != 0 {
1722 return Err(ProxyErrorCode::SessionGenerationConflict.error(ProxyErrorPhase::Session));
1723 }
1724 self.generation_retries = 1;
1725 Ok(())
1726 }
1727
1728 pub(crate) fn mark_committed(&mut self) -> Result<(), DirectProxyError> {
1729 self.transition(DirectIoState::Prepared, DirectIoState::CommittedNoIo)
1730 }
1731
1732 pub(crate) fn mark_io_attempted(&mut self) -> Result<(), DirectProxyError> {
1733 self.transition(DirectIoState::CommittedNoIo, DirectIoState::IoAttempted)
1734 }
1735
1736 pub(crate) fn mark_response_complete(&mut self) -> Result<(), DirectProxyError> {
1737 self.transition(DirectIoState::IoAttempted, DirectIoState::ResponseComplete)
1738 }
1739
1740 fn transition(
1741 &mut self,
1742 expected: DirectIoState,
1743 next: DirectIoState,
1744 ) -> Result<(), DirectProxyError> {
1745 if self.state != expected {
1746 return Err(ProxyErrorCode::InvalidStateTransition.error(ProxyErrorPhase::Framework));
1747 }
1748 self.state = next;
1749 Ok(())
1750 }
1751
1752 fn require_committed_no_io(&self) -> Result<(), DirectProxyError> {
1753 if self.state == DirectIoState::CommittedNoIo {
1754 Ok(())
1755 } else {
1756 Err(ProxyErrorCode::InvalidStateTransition.error(ProxyErrorPhase::Framework))
1757 }
1758 }
1759}
1760
1761impl Default for DirectTransactionState {
1762 fn default() -> Self {
1763 Self::new()
1764 }
1765}
1766
1767#[derive(Clone)]
1768struct AppState {
1769 config: Arc<ProxyConfig>,
1770 pipeline: Arc<Pipeline>,
1771 client: Client,
1772 direct: Option<Arc<DirectRuntime>>,
1773 sessions: Arc<RwLock<HashMap<String, SessionEntry>>>,
1774 started_at: Instant,
1775}
1776
1777struct SessionEntry {
1778 session: Arc<LegacySession>,
1779 expires_at: Instant,
1780}
1781
1782struct LegacySession {
1785 session: Session,
1786 json_spelled_tokens: std::sync::Mutex<HashSet<String>>,
1790}
1791
1792impl LegacySession {
1793 fn new(session: Session) -> Self {
1794 Self {
1795 session,
1796 json_spelled_tokens: std::sync::Mutex::new(HashSet::new()),
1797 }
1798 }
1799
1800 fn json_spelled_tokens(&self) -> std::sync::MutexGuard<'_, HashSet<String>> {
1802 self.json_spelled_tokens
1803 .lock()
1804 .unwrap_or_else(std::sync::PoisonError::into_inner)
1805 }
1806}
1807
1808#[cfg(test)]
1809static SESSION_FOR_MISS_BARRIERS: std::sync::LazyLock<
1810 std::sync::Mutex<HashMap<String, Arc<tokio::sync::Barrier>>>,
1811> = std::sync::LazyLock::new(|| std::sync::Mutex::new(HashMap::new()));
1812
1813#[derive(Debug, Clone, Serialize)]
1814#[non_exhaustive]
1815pub struct HealthSnapshot {
1816 pub uptime_secs: u64,
1817 pub bind: String,
1818 pub adapters: Vec<AdapterSnapshot>,
1819}
1820
1821#[derive(Debug, Clone, Serialize)]
1822#[non_exhaustive]
1823pub struct AdapterSnapshot {
1824 pub name: &'static str,
1825 pub upstream: String,
1826}
1827
1828pub async fn serve(config: ProxyConfig, pipeline: Arc<Pipeline>) -> Result<(), ProxyError> {
1830 serve_with_ready(config, pipeline, || {}).await
1831}
1832
1833pub async fn serve_with_ready(
1835 config: ProxyConfig,
1836 pipeline: Arc<Pipeline>,
1837 ready: impl FnOnce(),
1838) -> Result<(), ProxyError> {
1839 let bind = config.bind;
1840 let app = build_app(config, pipeline)?;
1841 let listener = TcpListener::bind(bind)
1842 .await
1843 .map_err(|source| ProxyError::Server { source })?;
1844 ready();
1845 serve_app(listener, app).await
1846}
1847
1848pub async fn serve_with_listener(
1856 mut config: ProxyConfig,
1857 pipeline: Arc<Pipeline>,
1858 listener: TcpListener,
1859) -> Result<(), ProxyError> {
1860 let actual = listener
1861 .local_addr()
1862 .map_err(|source| ProxyError::Server { source })?;
1863 if config.bind.port() == 0 {
1864 config.bind.set_port(actual.port());
1865 }
1866 if config.bind != actual {
1867 return Err(ProxyError::Server {
1868 source: std::io::Error::new(
1869 std::io::ErrorKind::InvalidInput,
1870 "listener address does not match proxy configuration",
1871 ),
1872 });
1873 }
1874 let app = build_app(config, pipeline)?;
1875 serve_app(listener, app).await
1876}
1877
1878fn build_app(config: ProxyConfig, pipeline: Arc<Pipeline>) -> Result<Router, ProxyError> {
1879 let direct = DirectRuntime::from_config(&config)?.map(Arc::new);
1880 let state = AppState {
1881 config: Arc::new(config),
1882 pipeline,
1883 client: Client::new(),
1884 direct,
1885 sessions: Arc::new(RwLock::new(HashMap::new())),
1886 started_at: Instant::now(),
1887 };
1888 Ok(Router::new()
1889 .route("/_gaze_proxy/healthz", get(healthz))
1890 .fallback(proxy)
1891 .with_state(state))
1892}
1893
1894async fn serve_app(listener: TcpListener, app: Router) -> Result<(), ProxyError> {
1895 axum::serve(
1896 listener,
1897 app.into_make_service_with_connect_info::<SocketAddr>(),
1898 )
1899 .await
1900 .map_err(|source| ProxyError::Server { source })
1901}
1902
1903async fn healthz(State(state): State<AppState>) -> impl IntoResponse {
1904 axum::Json(health_snapshot(&state))
1905}
1906
1907fn health_snapshot(state: &AppState) -> HealthSnapshot {
1908 HealthSnapshot {
1909 uptime_secs: state.started_at.elapsed().as_secs(),
1910 bind: state.config.bind.to_string(),
1911 adapters: state
1912 .config
1913 .adapters
1914 .iter()
1915 .map(|adapter| AdapterSnapshot {
1916 name: adapter.name(),
1917 upstream: adapter.upstream_base().to_string(),
1918 })
1919 .collect(),
1920 }
1921}
1922
1923async fn proxy(
1924 State(state): State<AppState>,
1925 ConnectInfo(peer): ConnectInfo<SocketAddr>,
1926 request: Request<Body>,
1927) -> Response {
1928 let (parts, body) = request.into_parts();
1929 let direct_route_candidate =
1930 state.direct.is_some() && matches!(parts.uri.path(), "/v1/messages" | "/v1/messages/");
1931 let direct_ingress = if direct_route_candidate {
1932 match prepare_direct_ingress(&state, peer, &parts.method, &parts.uri, &parts.headers) {
1933 Ok(ingress) => Some(ingress),
1934 Err(error) => return direct_error_response(error),
1935 }
1936 } else {
1937 None
1938 };
1939 let is_direct = direct_ingress.is_some();
1940 let body_limit_bytes = if is_direct {
1941 state
1942 .direct
1943 .as_ref()
1944 .map_or(state.config.body_limit_bytes, |runtime| {
1945 state
1946 .config
1947 .body_limit_bytes
1948 .min(runtime.client_config.max_request_bytes as u64)
1949 })
1950 } else {
1951 state.config.body_limit_bytes
1952 };
1953 let body = match collect_inbound_body(body, &parts.headers, body_limit_bytes).await {
1954 Ok(body) => body,
1955 Err(_) if is_direct => {
1956 return direct_error_response(
1957 ProxyErrorCode::RequestBodyLimitExceeded.error(ProxyErrorPhase::RequestValidation),
1958 );
1959 }
1960 Err(error) => return proxy_error_response(error),
1961 };
1962 match proxy_inner(
1963 state,
1964 parts.method,
1965 parts.uri,
1966 parts.headers,
1967 body,
1968 direct_ingress,
1969 )
1970 .await
1971 {
1972 Ok(response) => response,
1973 Err(err) => proxy_error_response(err),
1974 }
1975}
1976
1977async fn collect_inbound_body(
1978 body: Body,
1979 headers: &HeaderMap,
1980 limit_bytes: u64,
1981) -> Result<Bytes, ProxyError> {
1982 if let Some(length) = headers
1983 .get(axum::http::header::CONTENT_LENGTH)
1984 .and_then(|value| std::str::from_utf8(value.as_bytes()).ok())
1985 .and_then(|value| value.parse::<u64>().ok())
1986 {
1987 if length > limit_bytes {
1988 return Err(ProxyError::BodyTooLarge { limit_bytes });
1989 }
1990 }
1991
1992 let mut stream = body.into_data_stream();
1993 let mut collected = Vec::new();
1994 let mut accumulated = 0_u64;
1995 while let Some(chunk) = stream.next().await {
1996 let chunk = chunk.map_err(|_| ProxyError::DaemonConfig {
1997 detail: "request_body_rejected".to_string(),
1998 })?;
1999 accumulated = accumulated
2000 .checked_add(chunk.len() as u64)
2001 .ok_or(ProxyError::BodyTooLarge { limit_bytes })?;
2002 if accumulated > limit_bytes {
2003 return Err(ProxyError::BodyTooLarge { limit_bytes });
2004 }
2005 collected.extend_from_slice(&chunk);
2006 }
2007 Ok(Bytes::from(collected))
2008}
2009
2010async fn proxy_inner(
2011 state: AppState,
2012 method: Method,
2013 uri: Uri,
2014 headers: HeaderMap,
2015 body: Bytes,
2016 direct_ingress: Option<PreparedDirectIngress>,
2017) -> Result<Response, ProxyError> {
2018 if body.len() as u64 > state.config.body_limit_bytes {
2019 return Err(ProxyError::BodyTooLarge {
2020 limit_bytes: state.config.body_limit_bytes,
2021 });
2022 }
2023 let path = uri.path();
2024 let adapter = state
2025 .config
2026 .adapters
2027 .iter()
2028 .find(|adapter| adapter.matches_path(&method, path))
2029 .cloned()
2030 .ok_or_else(|| ProxyError::AdapterNotFound {
2031 path: path.to_string(),
2032 method: method.clone(),
2033 })?;
2034 let contract = adapter.contract();
2035 if matches!(contract.protocol(), ProtocolContract::Codec(_)) {
2036 let runtime = state
2037 .direct
2038 .as_ref()
2039 .ok_or_else(|| direct_readiness_error("direct_runtime_missing"))?;
2040 let resolved_contract = runtime.adapter.provider().contract();
2041 let ProtocolContract::Codec(codec) = resolved_contract.protocol() else {
2042 return Err(direct_readiness_error("direct_codec_missing"));
2043 };
2044 return Ok(
2045 match direct_proxy_inner(
2046 &state,
2047 direct_ingress.ok_or_else(|| ProxyError::DaemonConfig {
2048 detail: "direct_ingress_missing".to_string(),
2049 })?,
2050 &body,
2051 codec,
2052 )
2053 .await
2054 {
2055 Ok(response) => response,
2056 Err(error) => direct_error_response(error),
2057 },
2058 );
2059 }
2060 let session = session_for(&state, &headers).await?;
2061 let mut json: Value =
2062 serde_json::from_slice(&body).map_err(|source| ProxyError::InvalidJson { source })?;
2063 let json_output = adapter.requests_json_output(&json);
2066 let locale_chain = state.config.locale_chain();
2070 let dictionaries = state.config.dictionaries();
2071 let redacted = redact_surfaces(
2072 &state.pipeline,
2073 &session.session,
2074 adapter.request_pii_surfaces(&mut json),
2075 locale_chain,
2076 dictionaries,
2077 )?;
2078 session
2079 .json_spelled_tokens()
2080 .extend(redacted.json_spelled_tokens);
2081 match contract.coverage_policy() {
2085 CoveragePolicy::LegacySurfaces => {
2088 residual_scan_request(
2089 &state.pipeline,
2090 &json,
2091 &redacted.field_paths,
2092 locale_chain,
2093 dictionaries,
2094 )?;
2095 }
2096 CoveragePolicy::CodecProved => return Err(ProxyError::UnprovenCoverage),
2099 }
2100
2101 let upstream_url = upstream_url(adapter.upstream_base(), &uri)?;
2102 let mut request = state
2103 .client
2104 .request(method.clone(), upstream_url.clone())
2105 .json(&json);
2106 request = forward_headers(request, &headers);
2107 let upstream = request
2108 .send()
2109 .await
2110 .map_err(|source| ProxyError::UpstreamUnreachable {
2111 url: upstream_url.clone(),
2112 source,
2113 })?;
2114
2115 let status = upstream.status();
2116 let upstream_headers = upstream.headers().clone();
2117 let content_type = upstream_headers
2118 .get(reqwest::header::CONTENT_TYPE)
2119 .and_then(|value| value.to_str().ok())
2120 .unwrap_or_default()
2121 .to_ascii_lowercase();
2122 let bytes = upstream
2123 .bytes()
2124 .await
2125 .map_err(|source| ProxyError::UpstreamUnreachable {
2126 url: upstream_url,
2127 source,
2128 })?;
2129
2130 let is_sse = content_type.contains("text/event-stream");
2131 let body = if is_sse {
2132 transform_sse(&adapter, &session, &bytes, json_output)?
2133 } else {
2134 let mut response_json: Value =
2135 serde_json::from_slice(&bytes).map_err(|source| ProxyError::InvalidJson { source })?;
2136 restore_surfaces(
2137 &session,
2138 adapter.response_pii_surfaces(&mut response_json),
2139 json_output,
2140 );
2141 serde_json::to_vec(&response_json).map_err(|source| ProxyError::InvalidJson { source })?
2142 };
2143
2144 let mut response = Response::builder().status(status);
2145 response = response.header(
2146 axum::http::header::CONTENT_TYPE,
2147 if is_sse {
2148 "text/event-stream"
2149 } else {
2150 "application/json"
2151 },
2152 );
2153 response
2154 .body(axum::body::Body::from(body))
2155 .map_err(|source| ProxyError::DaemonConfig {
2156 detail: source.to_string(),
2157 })
2158}
2159
2160fn prepare_direct_ingress(
2161 state: &AppState,
2162 peer: SocketAddr,
2163 method: &Method,
2164 uri: &Uri,
2165 headers: &HeaderMap,
2166) -> Result<PreparedDirectIngress, DirectProxyError> {
2167 let runtime = state
2168 .direct
2169 .as_ref()
2170 .ok_or_else(|| ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::Framework))?;
2171 if method != Method::POST || uri.path() != "/v1/messages" || uri.query().is_some() {
2172 return Err(ProxyErrorCode::RouteRejected.error(ProxyErrorPhase::Ingress));
2173 }
2174 let outbound_headers = validate_direct_request_headers(headers, runtime)?;
2175 let session = match runtime.session_policy {
2176 SessionPolicy::EphemeralSingleRequest => {
2177 if headers.get_all(SESSION_HEADER).iter().next().is_some() {
2178 return Err(ProxyErrorCode::UnexpectedSessionHeader
2179 .error(ProxyErrorPhase::RequestValidation));
2180 }
2181 DirectIngressSession::Ephemeral
2182 }
2183 SessionPolicy::OpaqueHeaderContinuity(_) => {
2184 let identifier = canonical_session_identifier(headers)?;
2185 let peer_scope = if peer.ip().is_loopback() {
2186 PeerScope::Loopback
2187 } else {
2188 PeerScope::NonLoopback
2189 };
2190 let context = PrincipalContext::new(runtime.listener_scope, peer_scope, None, None)
2191 .map_err(|_| ProxyErrorCode::PrincipalRejected.error(ProxyErrorPhase::Session))?;
2192 let principal = invoke_resolver(runtime.resolver.as_ref(), context)
2193 .map_err(|_| ProxyErrorCode::PrincipalRejected.error(ProxyErrorPhase::Session))?;
2194 let registry = runtime.registry.as_ref().ok_or_else(|| {
2195 ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::Session)
2196 })?;
2197 let lease = registry
2198 .acquire(principal, identifier, None)
2199 .map_err(map_registry_error)?;
2200 DirectIngressSession::Continuity(lease)
2201 }
2202 };
2203 Ok(PreparedDirectIngress {
2204 outbound_headers,
2205 session,
2206 })
2207}
2208
2209async fn direct_proxy_inner(
2210 state: &AppState,
2211 ingress: PreparedDirectIngress,
2212 body: &[u8],
2213 codec: &dyn BodyCodec,
2214) -> Result<Response, DirectProxyError> {
2215 let runtime = state
2216 .direct
2217 .as_ref()
2218 .ok_or_else(|| ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::Framework))?;
2219 if body.len() > runtime.client_config.max_request_bytes {
2220 return Err(
2221 ProxyErrorCode::RequestBodyLimitExceeded.error(ProxyErrorPhase::RequestValidation)
2222 );
2223 }
2224 let expected_response = parse_request_stream_format(body)?;
2225 let mut inspection = state
2226 .config
2227 .inspection
2228 .as_ref()
2229 .and_then(|producer| producer.begin_logical(runtime.inspection_endpoint_codes));
2230
2231 let mut continuity_lease = None;
2232 let mut committed = match ingress.session {
2233 DirectIngressSession::Ephemeral => {
2234 let session = Session::new(Scope::Ephemeral)
2235 .map_err(|_| ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::Session))?;
2236 prepare_and_commit_direct_request(
2237 &state.pipeline,
2238 state.config.locale_chain(),
2239 state.config.dictionaries(),
2240 &session,
2241 codec,
2242 body,
2243 &runtime.endpoint,
2244 &ingress.outbound_headers,
2245 expected_response,
2246 runtime.codec_limits,
2247 inspection
2248 .as_ref()
2249 .is_some_and(ProxyInspectionLogicalV1::provider_request_projection_selected),
2250 runtime.client_config.max_request_bytes,
2251 )?
2252 }
2253 DirectIngressSession::Continuity(lease) => {
2254 let request_guard = lease.lock_request().await;
2255 let committed = prepare_and_commit_direct_request(
2256 &state.pipeline,
2257 state.config.locale_chain(),
2258 state.config.dictionaries(),
2259 lease.session(),
2260 codec,
2261 body,
2262 &runtime.endpoint,
2263 &ingress.outbound_headers,
2264 expected_response,
2265 runtime.codec_limits,
2266 inspection
2267 .as_ref()
2268 .is_some_and(ProxyInspectionLogicalV1::provider_request_projection_selected),
2269 runtime.client_config.max_request_bytes,
2270 );
2271 drop(request_guard);
2272 let committed = committed?;
2273 continuity_lease = Some(lease);
2274 committed
2275 }
2276 };
2277
2278 if let Some(inspection) = inspection.as_mut() {
2279 let inspection_parsed_facts = committed.request.inspection_parsed_facts.take();
2280 inspection.emit_request_stages(
2281 body,
2282 &committed.request.body,
2283 inspection_parsed_facts,
2284 committed.request.signed_or_encrypted_surface,
2285 );
2286 }
2287
2288 if expected_response == WireFormat::Sse {
2289 let opened = runtime.client.execute_sse(committed).await?;
2290 return stage_and_replay_direct_sse(
2291 Arc::clone(&state.pipeline),
2292 state.config.locale_chain.clone(),
2293 state.config.dictionaries().clone(),
2294 opened,
2295 runtime.codec_limits,
2296 runtime.client_config.max_response_bytes,
2297 runtime.ping_interval,
2298 inspection,
2299 continuity_lease,
2300 );
2301 }
2302 let executed = runtime.client.execute(committed).await?;
2303 let response = restore_direct_response(
2304 &state.pipeline,
2305 state.config.locale_chain(),
2306 state.config.dictionaries(),
2307 codec,
2308 executed,
2309 runtime.codec_limits,
2310 inspection.as_mut(),
2311 );
2312 drop(continuity_lease);
2313 response
2314}
2315
2316#[allow(clippy::too_many_arguments)]
2317fn stage_and_replay_direct_sse(
2318 pipeline: Arc<Pipeline>,
2319 locale_chain: gaze::LocaleChain,
2320 dictionaries: DictionaryBundle,
2321 opened: CommittedDirectSse,
2322 limits: CodecLimits,
2323 max_response_bytes: usize,
2324 ping_interval: Duration,
2325 inspection: Option<ProxyInspectionLogicalV1>,
2326 continuity_lease: Option<SessionLease>,
2327) -> Result<Response, DirectProxyError> {
2328 let CommittedDirectSse {
2329 snapshot,
2330 response: upstream_response,
2331 head,
2332 lifecycle,
2333 remaining_total_timeout,
2334 } = opened;
2335 let status = head.status();
2336 let headers = head.headers().clone();
2337 let (sender, receiver) = tokio::sync::mpsc::channel::<Result<Bytes, std::io::Error>>(8);
2338 let mut response = Response::builder().status(status);
2339 for (name, value) in &headers {
2340 response = response.header(name, value);
2341 }
2342 let response = response
2343 .body(Body::from_stream(
2344 tokio_stream::wrappers::ReceiverStream::new(receiver),
2345 ))
2346 .map_err(|_| ProxyErrorCode::UpstreamProtocol.error(ProxyErrorPhase::Framework))?;
2347 tokio::spawn(async move {
2348 let _continuity_lease = continuity_lease;
2349 let mut inspection = inspection;
2350 let proof = async move {
2351 let body = collect_response_body(upstream_response, max_response_bytes).await?;
2352 validate_declared_body(WireFormat::Sse, &body, max_response_bytes)?;
2353 let mut residual = PipelineResponseResidualValidator::new(
2354 &pipeline,
2355 locale_chain.as_slice(),
2356 &dictionaries,
2357 )?;
2358 let mut context =
2359 ResponseTransformContext::new(&snapshot, &mut residual, WireFormat::Sse, limits);
2360 if inspection
2361 .as_ref()
2362 .is_some_and(ProxyInspectionLogicalV1::owner_restored_response_projection_selected)
2363 {
2364 context.request_inspection_projection();
2365 }
2366 let mut proved = crate::codecs::anthropic::AnthropicMessagesCodec
2367 .restore_response(&body, &mut context)
2368 .map_err(map_codec_error)?;
2369 if let Some(inspection) = inspection.as_mut() {
2370 let format = proved.format();
2371 let signed_or_encrypted_surface = proved.signed_or_encrypted_surface();
2372 let inspection_parsed_facts = proved.take_inspection_parsed_facts();
2373 inspection.emit_response_stages(
2374 &body,
2375 proved.bytes(),
2376 format,
2377 inspection_parsed_facts,
2378 signed_or_encrypted_surface,
2379 );
2380 }
2381 let mut lifecycle = lifecycle;
2382 lifecycle.mark_response_complete()?;
2383 Ok::<Vec<u8>, DirectProxyError>(proved.into_bytes())
2384 };
2385 let timed_proof = tokio::time::timeout(remaining_total_timeout, proof);
2386 tokio::pin!(timed_proof);
2387 let mut ping = tokio::time::interval(ping_interval);
2388 ping.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
2389 ping.tick().await;
2390 loop {
2391 tokio::select! {
2392 () = sender.closed() => {
2393 break;
2394 }
2395 result = &mut timed_proof => {
2396 let output = match result {
2397 Ok(Ok(bytes)) => Bytes::from(bytes),
2398 Ok(Err(_)) | Err(_) => Bytes::from_static(DIRECT_PROXY_ERROR_FRAME),
2399 };
2400 let _ = sender.send(Ok(output)).await;
2401 break;
2402 }
2403 _ = ping.tick() => {
2404 if sender
2405 .send(Ok(Bytes::from_static(DIRECT_PROXY_PING_FRAME)))
2406 .await
2407 .is_err()
2408 {
2409 break;
2410 }
2411 }
2412 }
2413 }
2414 });
2415 Ok(response)
2416}
2417
2418fn validate_direct_request_headers(
2419 headers: &HeaderMap,
2420 runtime: &DirectRuntime,
2421) -> Result<HeaderMap, DirectProxyError> {
2422 validate_header_space(headers, runtime.client_config)?;
2423 for singleton in [
2424 "content-type",
2425 "content-length",
2426 "transfer-encoding",
2427 "content-encoding",
2428 "connection",
2429 "x-api-key",
2430 "anthropic-version",
2431 "anthropic-beta",
2432 SESSION_HEADER,
2433 ] {
2434 reject_duplicate_request_header(headers, singleton)?;
2435 }
2436 if headers.get(axum::http::header::CONTENT_ENCODING).is_some() {
2437 return Err(
2438 ProxyErrorCode::UnsupportedContentEncoding.error(ProxyErrorPhase::RequestValidation)
2439 );
2440 }
2441 if let Some(connection) = one_request_header(headers, "connection", false)? {
2442 let connection = connection.to_str().map_err(|_| {
2443 ProxyErrorCode::InvalidFraming.error(ProxyErrorPhase::RequestValidation)
2444 })?;
2445 if !connection.eq_ignore_ascii_case("keep-alive")
2446 && !connection.eq_ignore_ascii_case("close")
2447 {
2448 return Err(ProxyErrorCode::InvalidFraming.error(ProxyErrorPhase::RequestValidation));
2449 }
2450 }
2451 if headers.get(axum::http::header::TRANSFER_ENCODING).is_some()
2452 && headers.get(axum::http::header::CONTENT_LENGTH).is_some()
2453 {
2454 return Err(ProxyErrorCode::InvalidFraming.error(ProxyErrorPhase::RequestValidation));
2455 }
2456 let content_type = one_request_header(headers, "content-type", true)?
2457 .ok_or_else(|| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2458 if parse_wire_content_type(content_type)? != WireFormat::Json {
2459 return Err(ProxyErrorCode::InvalidRequestFormat.error(ProxyErrorPhase::RequestValidation));
2460 }
2461 let api_key = one_request_header(headers, "x-api-key", true)?
2462 .ok_or_else(|| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2463 if !api_key
2464 .as_bytes()
2465 .iter()
2466 .all(|byte| (0x21..=0x7e).contains(byte) && *byte != b',')
2467 {
2468 return Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation));
2469 }
2470 let version = one_request_header(headers, "anthropic-version", true)?
2471 .ok_or_else(|| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2472 let version_text = std::str::from_utf8(version.as_bytes())
2473 .map_err(|_| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2474 if !runtime
2475 .allowed_versions
2476 .iter()
2477 .any(|allowed| allowed == version_text)
2478 {
2479 return Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation));
2480 }
2481 let beta = one_request_header(headers, "anthropic-beta", false)?;
2482 if let Some(beta) = beta {
2483 let beta_text = std::str::from_utf8(beta.as_bytes()).map_err(|_| {
2484 ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation)
2485 })?;
2486 if !runtime
2487 .allowed_betas
2488 .iter()
2489 .any(|allowed| allowed == beta_text)
2490 {
2491 return Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation));
2492 }
2493 }
2494
2495 let mut outbound = HeaderMap::new();
2496 outbound.insert(
2497 axum::http::header::CONTENT_TYPE,
2498 HeaderValue::from_static("application/json"),
2499 );
2500 outbound.insert(HeaderName::from_static("x-api-key"), api_key.clone());
2501 outbound.insert(
2502 HeaderName::from_static("anthropic-version"),
2503 version.clone(),
2504 );
2505 if let Some(beta) = beta {
2506 outbound.insert(HeaderName::from_static("anthropic-beta"), beta.clone());
2507 }
2508 Ok(outbound)
2509}
2510
2511fn one_request_header<'a>(
2512 headers: &'a HeaderMap,
2513 name: &'static str,
2514 required: bool,
2515) -> Result<Option<&'a HeaderValue>, DirectProxyError> {
2516 let mut values = headers.get_all(name).iter();
2517 let first = values.next();
2518 if values.next().is_some() || first.is_some_and(|value| value.is_empty()) {
2519 return Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation));
2520 }
2521 match (first, required) {
2522 (Some(value), _) => Ok(Some(value)),
2523 (None, false) => Ok(None),
2524 (None, true) => {
2525 Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))
2526 }
2527 }
2528}
2529
2530fn reject_duplicate_request_header(
2531 headers: &HeaderMap,
2532 name: &'static str,
2533) -> Result<(), DirectProxyError> {
2534 let mut values = headers.get_all(name).iter();
2535 let _ = values.next();
2536 if values.next().is_some() {
2537 Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))
2538 } else {
2539 Ok(())
2540 }
2541}
2542
2543fn canonical_session_identifier(headers: &HeaderMap) -> Result<&str, DirectProxyError> {
2544 let value = one_request_header(headers, SESSION_HEADER, false)?
2545 .ok_or_else(|| {
2546 ProxyErrorCode::SessionIdentityRequired.error(ProxyErrorPhase::RequestValidation)
2547 })?
2548 .to_str()
2549 .map_err(|_| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2550 let identifier = uuid::Uuid::parse_str(value)
2551 .map_err(|_| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2552 if identifier.get_version() != Some(uuid::Version::Random)
2553 || identifier.get_variant() != uuid::Variant::RFC4122
2554 || identifier.hyphenated().to_string() != value
2555 {
2556 return Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation));
2557 }
2558 Ok(value)
2559}
2560
2561fn restore_direct_response(
2562 pipeline: &Pipeline,
2563 locale_chain: &[LocaleTag],
2564 dictionaries: &DictionaryBundle,
2565 codec: &dyn BodyCodec,
2566 mut executed: CommittedDirectResponse,
2567 limits: CodecLimits,
2568 inspection: Option<&mut ProxyInspectionLogicalV1>,
2569) -> Result<Response, DirectProxyError> {
2570 let (head, body) = executed.response.into_parts();
2571 let mut residual =
2572 PipelineResponseResidualValidator::new(pipeline, locale_chain, dictionaries)?;
2573 let mut context =
2574 ResponseTransformContext::new(&executed.snapshot, &mut residual, head.format(), limits);
2575 if inspection
2576 .as_deref()
2577 .is_some_and(ProxyInspectionLogicalV1::owner_restored_response_projection_selected)
2578 {
2579 context.request_inspection_projection();
2580 }
2581 let mut proved = codec
2582 .restore_response(&body, &mut context)
2583 .map_err(map_codec_error)?;
2584 if let Some(inspection) = inspection {
2585 let format = proved.format();
2586 let signed_or_encrypted_surface = proved.signed_or_encrypted_surface();
2587 let inspection_parsed_facts = proved.take_inspection_parsed_facts();
2588 inspection.emit_response_stages(
2589 &body,
2590 proved.bytes(),
2591 format,
2592 inspection_parsed_facts,
2593 signed_or_encrypted_surface,
2594 );
2595 }
2596 executed.lifecycle.mark_response_complete()?;
2597 let mut response = Response::builder().status(head.status());
2598 for (name, value) in head.headers() {
2599 response = response.header(name, value);
2600 }
2601 response
2602 .body(Body::from(proved.into_bytes()))
2603 .map_err(|_| ProxyErrorCode::UpstreamProtocol.error(ProxyErrorPhase::Framework))
2604}
2605
2606fn map_registry_error(error: RegistryError) -> DirectProxyError {
2607 let code = match error {
2608 RegistryError::InvalidSessionIdentifier => ProxyErrorCode::HeaderRejected,
2609 RegistryError::SessionExpired => ProxyErrorCode::SessionExpired,
2610 RegistryError::SessionGenerationConflict => ProxyErrorCode::SessionGenerationConflict,
2611 RegistryError::ActiveCapacityReached => ProxyErrorCode::RegistryCapacity,
2612 RegistryError::EntropyUnavailable
2613 | RegistryError::SessionCreationFailed
2614 | RegistryError::EpochExhausted
2615 | RegistryError::InternalFailure => ProxyErrorCode::ProxyConfiguration,
2616 };
2617 code.error(ProxyErrorPhase::Session)
2618}
2619
2620fn map_codec_error(error: CodecError) -> DirectProxyError {
2621 let codec_phase = error.phase();
2622 let opaque_carrier = error.opaque_carrier();
2623 let request_direction = matches!(
2624 codec_phase,
2625 CodecPhase::RequestParse | CodecPhase::RequestProtection | CodecPhase::RequestProof
2626 );
2627 let code = match error.code() {
2628 CodecErrorCode::DuplicateObjectKey => ProxyErrorCode::DuplicateObjectKey,
2629 CodecErrorCode::InternalCoverageFailure => ProxyErrorCode::InternalCoverageFailure,
2630 CodecErrorCode::ControlWouldMutate => ProxyErrorCode::ControlWouldMutate,
2631 CodecErrorCode::OpaqueMediaUninspected => ProxyErrorCode::OpaqueMediaUninspected,
2632 CodecErrorCode::SignedMutationRequired => ProxyErrorCode::SignedMutationRequired,
2633 CodecErrorCode::SignedSurfaceMalformed => ProxyErrorCode::SignedSurfaceMalformed,
2634 CodecErrorCode::InvalidSseLifecycle | CodecErrorCode::IncompleteSseStream => {
2635 ProxyErrorCode::InvalidSseLifecycle
2636 }
2637 CodecErrorCode::InvalidContentBlockIndex | CodecErrorCode::InvalidContentBlockDelta => {
2638 ProxyErrorCode::InvalidContentBlockIndex
2639 }
2640 CodecErrorCode::DepthLimitExceeded
2641 | CodecErrorCode::NodeLimitExceeded
2642 | CodecErrorCode::StringLimitExceeded
2643 | CodecErrorCode::NumberLimitExceeded
2644 | CodecErrorCode::BodyLimitExceeded
2645 | CodecErrorCode::GrowthLimitExceeded
2646 | CodecErrorCode::SseLineLimitExceeded
2647 | CodecErrorCode::SseFrameLimitExceeded
2648 | CodecErrorCode::SseEventLimitExceeded
2649 | CodecErrorCode::SseIndexLimitExceeded
2650 | CodecErrorCode::SseAccumulatorLimitExceeded => {
2651 if request_direction {
2652 ProxyErrorCode::RequestBodyLimitExceeded
2653 } else {
2654 ProxyErrorCode::ResponseBodyLimitExceeded
2655 }
2656 }
2657 CodecErrorCode::ProtectionFailedClosed
2658 | CodecErrorCode::RestoreFailedClosed
2659 | CodecErrorCode::ProviderOriginPii
2660 | CodecErrorCode::ProviderOriginNumeric
2661 | CodecErrorCode::TransformedKeyCollision
2662 | CodecErrorCode::InvalidUtf8
2663 | CodecErrorCode::InvalidJson
2664 | CodecErrorCode::InvalidWireFormat
2665 | CodecErrorCode::UnsupportedRequestSurface
2666 | CodecErrorCode::UnsupportedResponseSurface
2667 | CodecErrorCode::InvalidSseField
2668 | CodecErrorCode::InvalidSseFrame
2669 | CodecErrorCode::InvalidSseEventPair
2670 | CodecErrorCode::UnsafeFutureEvent
2671 | CodecErrorCode::ProviderErrorEvent => ProxyErrorCode::InvalidToken,
2672 };
2673 let phase = match codec_phase {
2674 CodecPhase::RequestParse => ProxyErrorPhase::RequestValidation,
2675 CodecPhase::RequestProtection | CodecPhase::RequestProof => {
2676 ProxyErrorPhase::RequestTransform
2677 }
2678 CodecPhase::ResponseParse => ProxyErrorPhase::ResponseValidation,
2679 CodecPhase::ResponseRestore | CodecPhase::ResponseProof => {
2680 ProxyErrorPhase::ResponseTransform
2681 }
2682 CodecPhase::SseDecode | CodecPhase::SseLifecycle => ProxyErrorPhase::ResponseValidation,
2683 CodecPhase::SseReplay => ProxyErrorPhase::ResponseReplay,
2684 };
2685 let direct_error = code.error(phase);
2686 if let Some(opaque_carrier) = opaque_carrier {
2687 direct_error.with_opaque_carrier(opaque_carrier)
2688 } else {
2689 direct_error
2690 }
2691}
2692
2693fn refusal_json(refusal: &BoundaryRefusal) -> Value {
2696 serde_json::json!({
2697 "error": refusal.error.as_str(),
2698 "fallback_reason": refusal.fallback_reason.map(|reason| reason.as_str()),
2699 "suspect_classes": refusal
2700 .suspect_classes
2701 .iter()
2702 .map(gaze::PiiClass::to_canonical_str)
2703 .collect::<Vec<_>>(),
2704 })
2705}
2706
2707fn log_refusal(refusal: &BoundaryRefusal) {
2709 eprintln!("gaze-proxy: request refused: {}", refusal_json(refusal));
2710}
2711
2712fn direct_error_response(error: DirectProxyError) -> Response {
2713 let mut error_body = serde_json::json!({
2714 "type": "api_error",
2715 "message": "proxy_validation_failed",
2716 "code": format!("{:?}", error.code()),
2717 "phase": format!("{:?}", error.phase()),
2718 });
2719 if let (Some(name), Some(reason)) = (error.header_name(), error.header_rejection_reason()) {
2720 error_body["header"] = serde_json::json!({
2721 "name": name,
2722 "reason": reason.as_str(),
2723 });
2724 }
2725 if let Some(refusal) = error.refusal() {
2726 log_refusal(refusal);
2727 error_body["refusal"] = refusal_json(refusal);
2728 }
2729 if let Some(carrier) = error.opaque_carrier() {
2730 error_body["carrier"] = serde_json::json!({
2731 "scheme": carrier.scheme().as_str(),
2732 "span": {
2733 "byte_offset": carrier.byte_offset(),
2734 "byte_length": carrier.byte_length(),
2735 },
2736 "content_block_index": carrier.content_block_index(),
2737 });
2738 }
2739 let body = serde_json::json!({
2740 "type": "error",
2741 "error": error_body,
2742 });
2743 let mut response = Response::builder()
2744 .status(error.http_status())
2745 .header(axum::http::header::CONTENT_TYPE, "application/json");
2746 if let Some(seconds) = error.retry_after_seconds() {
2747 response = response.header(axum::http::header::RETRY_AFTER, seconds.to_string());
2748 }
2749 response
2750 .body(Body::from(body.to_string()))
2751 .unwrap_or_else(|_| Response::new(Body::empty()))
2752}
2753
2754async fn session_for(
2755 state: &AppState,
2756 headers: &HeaderMap,
2757) -> Result<Arc<LegacySession>, ProxyError> {
2758 let now = Instant::now();
2759 let id = headers
2760 .get(SESSION_HEADER)
2761 .and_then(|value| value.to_str().ok())
2762 .filter(|value| !value.trim().is_empty())
2763 .map(str::to_string)
2764 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
2765
2766 #[cfg(test)]
2767 {
2768 let barrier = SESSION_FOR_MISS_BARRIERS.lock().unwrap().get(&id).cloned();
2769 if let Some(barrier) = barrier {
2770 barrier.wait().await;
2771 }
2772 }
2773
2774 let mut sessions = state.sessions.write().await;
2775 sessions.retain(|_, entry| entry.expires_at > now);
2776 match sessions.entry(id) {
2777 std::collections::hash_map::Entry::Occupied(entry) => Ok(entry.get().session.clone()),
2778 std::collections::hash_map::Entry::Vacant(entry) => {
2779 let session = Arc::new(LegacySession::new(
2780 Session::new(Scope::Conversation(entry.key().clone()))
2781 .map_err(|source| ProxyError::Pipeline { source })?,
2782 ));
2783 entry.insert(SessionEntry {
2784 session: session.clone(),
2785 expires_at: now + state.config.session_ttl,
2786 });
2787 Ok(session)
2788 }
2789 }
2790}
2791
2792const NON_CARRIER_NUMERIC_KEYS: &[&str] = &[
2810 "budget_tokens",
2811 "candidateCount",
2812 "frequencyPenalty",
2813 "frequency_penalty",
2814 "logprobs",
2815 "maxOutputTokens",
2816 "max_completion_tokens",
2817 "max_tokens",
2818 "n",
2819 "presencePenalty",
2820 "presence_penalty",
2821 "seed",
2822 "temperature",
2823 "topK",
2824 "topP",
2825 "top_k",
2826 "top_logprobs",
2827 "top_p",
2828];
2829
2830const CARRIER_SUBTREE_KEYS: &[&str] = &[
2838 "functionDeclarations",
2839 "functions",
2840 "input_schema",
2841 "json_schema",
2842 "parameters",
2843 "properties",
2844 "response_format",
2845 "schema",
2846 "tools",
2847];
2848
2849fn redact_surfaces(
2863 pipeline: &Pipeline,
2864 session: &Session,
2865 surfaces: Vec<crate::adapter::PiiSurface<'_>>,
2866 locale_chain: &[LocaleTag],
2867 dictionaries: &DictionaryBundle,
2868) -> Result<RedactedSurfaces, ProxyError> {
2869 let mut redacted = RedactedSurfaces::default();
2870 for surface in surfaces {
2871 let text = pipeline
2875 .resolve_boundary_text(session, surface.text, locale_chain, dictionaries)
2876 .map_err(|refusal| ProxyError::Refused { refusal })?;
2877 pipeline
2878 .admit_boundary_text(session, &text, locale_chain, dictionaries)
2879 .map_err(|refusal| ProxyError::Refused { refusal })?;
2880 collect_json_spelled_tokens(surface.text, &text, &mut redacted.json_spelled_tokens);
2881 *surface.text = text;
2882 redacted.field_paths.insert(surface.field_path);
2883 }
2884 Ok(redacted)
2885}
2886
2887#[derive(Default)]
2888struct RedactedSurfaces {
2889 field_paths: HashSet<String>,
2891 json_spelled_tokens: HashSet<String>,
2893}
2894
2895fn collect_json_spelled_tokens(original: &str, clean: &str, tokens: &mut HashSet<String>) {
2901 if !matches!(clean.trim_start().as_bytes().first(), Some(b'{' | b'[')) {
2902 return;
2903 }
2904 let Ok(document) = serde_json::from_str::<Value>(clean) else {
2905 return;
2906 };
2907 let mut stack = vec![&document];
2908 while let Some(value) = stack.pop() {
2909 let mut record = |text: &str| {
2910 for token in token_shape::find_tokens(text) {
2911 if !original.contains(token) {
2912 tokens.insert(token.to_owned());
2913 }
2914 }
2915 };
2916 match value {
2917 Value::String(text) => record(text),
2918 Value::Array(items) => stack.extend(items),
2919 Value::Object(members) => {
2920 for (key, member) in members {
2921 record(key);
2922 stack.push(member);
2923 }
2924 }
2925 Value::Null | Value::Bool(_) | Value::Number(_) => {}
2926 }
2927 }
2928}
2929
2930fn residual_scan_request(
2946 pipeline: &Pipeline,
2947 body: &Value,
2948 surfaced: &HashSet<String>,
2949 locale_chain: &[LocaleTag],
2950 dictionaries: &DictionaryBundle,
2951) -> Result<(), ProxyError> {
2952 let session =
2953 Session::new(Scope::Ephemeral).map_err(|source| ProxyError::Pipeline { source })?;
2954 let scan = RequestResidualScan {
2955 pipeline,
2956 session,
2957 surfaced,
2958 locale_chain,
2959 dictionaries,
2960 };
2961 scan.walk("", "", body, false, true)
2962}
2963
2964struct RequestResidualScan<'a> {
2966 pipeline: &'a Pipeline,
2967 session: Session,
2968 surfaced: &'a HashSet<String>,
2969 locale_chain: &'a [LocaleTag],
2971 dictionaries: &'a DictionaryBundle,
2973}
2974
2975impl RequestResidualScan<'_> {
2976 fn probe(&self, field_path: &str, text: &str) -> Result<(), ProxyError> {
2977 let (clean, spans, report) = self
2978 .pipeline
2979 .clean_with_safety_net_detect_context(
2980 &self.session,
2981 RawDocument::Text(text.to_owned()),
2982 self.locale_chain,
2983 self.dictionaries,
2984 )
2985 .map_err(|source| ProxyError::Pipeline { source })?;
2986 let untouched = matches!(&clean, CleanDocument::Text(cleaned) if cleaned == text);
2992 if untouched && spans.is_empty() && report.suspects.is_empty() {
3000 return Ok(());
3001 }
3002 Err(ProxyError::UnsurfacedPii {
3003 field_path: field_path.to_string(),
3004 })
3005 }
3006
3007 fn walk(
3014 &self,
3015 path: &str,
3016 leaf_key: &str,
3017 value: &Value,
3018 in_carrier_subtree: bool,
3019 unambiguous: bool,
3020 ) -> Result<(), ProxyError> {
3021 match value {
3022 Value::String(text) => {
3023 if !(unambiguous && self.surfaced.contains(path)) {
3024 self.probe(path, text)?;
3025 }
3026 }
3027 Value::Number(number) => {
3028 if in_carrier_subtree || !NON_CARRIER_NUMERIC_KEYS.contains(&leaf_key) {
3029 self.probe(path, &number.to_string())?;
3030 }
3031 }
3032 Value::Array(items) => {
3033 for (index, item) in items.iter().enumerate() {
3034 self.walk(
3035 &format!("{path}[{index}]"),
3036 leaf_key,
3037 item,
3038 in_carrier_subtree,
3039 unambiguous,
3040 )?;
3041 }
3042 }
3043 Value::Object(map) => {
3044 for (key, child) in map {
3045 let child_path = if path.is_empty() {
3046 key.clone()
3047 } else {
3048 format!("{path}.{key}")
3049 };
3050 self.walk(
3051 &child_path,
3052 key,
3053 child,
3054 in_carrier_subtree || CARRIER_SUBTREE_KEYS.contains(&key.as_str()),
3055 unambiguous && !key.contains('.') && !key.contains('['),
3056 )?;
3057 }
3058 }
3059 Value::Null | Value::Bool(_) => {}
3060 }
3061 Ok(())
3062 }
3063}
3064
3065fn restore_surfaces(
3069 session: &LegacySession,
3070 surfaces: Vec<crate::adapter::PiiSurface<'_>>,
3071 json_output: bool,
3072) {
3073 let json_spelled_tokens = session.json_spelled_tokens();
3074 for surface in surfaces {
3075 let into_json_string = match surface.syntax {
3076 SurfaceSyntax::Text => false,
3077 SurfaceSyntax::Json => true,
3078 SurfaceSyntax::ModelOutput => json_output,
3079 };
3080 *surface.text = restore_text(
3081 &session.session,
3082 &json_spelled_tokens,
3083 surface.text,
3084 into_json_string,
3085 );
3086 }
3087}
3088
3089fn restore_text(
3100 session: &Session,
3101 json_spelled_tokens: &HashSet<String>,
3102 clean: &str,
3103 into_json_string: bool,
3104) -> String {
3105 let mut restored = String::new();
3106 let mut last = 0;
3107 for matched in token_shape::pattern().find_iter(clean) {
3108 restored.push_str(&clean[last..matched.start()]);
3109 let token = matched.as_str();
3110 match session.restore(token) {
3111 Some(raw)
3112 if into_json_string
3113 && !(json_spelled_tokens.contains(token) && is_json_string_content(&raw)) =>
3114 {
3115 push_json_string_contents(&mut restored, &raw);
3116 }
3117 Some(raw) => restored.push_str(&raw),
3118 None => restored.push_str(token),
3119 }
3120 last = matched.end();
3121 }
3122 restored.push_str(&clean[last..]);
3123 restored
3124}
3125
3126fn push_json_string_contents(out: &mut String, raw: &str) {
3128 let literal = Value::String(raw.to_owned()).to_string();
3129 out.push_str(&literal[1..literal.len() - 1]);
3130}
3131
3132fn is_json_string_content(text: &str) -> bool {
3134 serde_json::from_str::<String>(&format!("\"{text}\"")).is_ok()
3135}
3136
3137fn transform_sse(
3138 adapter: &Arc<dyn ProviderAdapter>,
3139 session: &LegacySession,
3140 bytes: &[u8],
3141 json_output: bool,
3142) -> Result<Vec<u8>, ProxyError> {
3143 let text = std::str::from_utf8(bytes).map_err(|err| ProxyError::SsePartialFrame {
3144 reason: err.to_string(),
3145 })?;
3146 let mut out = String::new();
3147 for frame in text.split("\n\n") {
3148 if frame.trim().is_empty() {
3149 continue;
3150 }
3151 let mut event_name = None;
3152 let mut data_lines = Vec::new();
3153 for line in frame.lines() {
3154 if let Some(rest) = line.strip_prefix("event:") {
3155 event_name = Some(rest.trim().to_string());
3156 } else if let Some(rest) = line.strip_prefix("data:") {
3157 data_lines.push(rest.trim_start());
3158 } else {
3159 out.push_str(line);
3160 out.push('\n');
3161 }
3162 }
3163 let data = data_lines.join("\n");
3164 if data == "[DONE]" {
3165 if let Some(name) = event_name {
3166 out.push_str("event: ");
3167 out.push_str(&name);
3168 out.push('\n');
3169 }
3170 out.push_str("data: [DONE]\n\n");
3171 continue;
3172 }
3173 let mut event = SseEvent {
3174 event: event_name.clone(),
3175 data: serde_json::from_str(&data)
3176 .map_err(|source| ProxyError::InvalidJson { source })?,
3177 };
3178 restore_surfaces(
3179 session,
3180 adapter.sse_event_pii_surfaces(&mut event),
3181 json_output,
3182 );
3183 if let Some(name) = event_name {
3184 out.push_str("event: ");
3185 out.push_str(&name);
3186 out.push('\n');
3187 }
3188 out.push_str("data: ");
3189 out.push_str(
3190 &serde_json::to_string(&event.data)
3191 .map_err(|source| ProxyError::InvalidJson { source })?,
3192 );
3193 out.push_str("\n\n");
3194 }
3195 Ok(out.into_bytes())
3196}
3197
3198fn upstream_url(base: &Url, uri: &Uri) -> Result<Url, ProxyError> {
3199 let path = uri.path_and_query().map(|pq| pq.as_str()).unwrap_or("/");
3200 base.join(path).map_err(|err| ProxyError::DaemonConfig {
3201 detail: err.to_string(),
3202 })
3203}
3204
3205fn forward_headers(
3206 mut request: reqwest::RequestBuilder,
3207 headers: &HeaderMap,
3208) -> reqwest::RequestBuilder {
3209 const FORWARDED: &[&str] = &[
3210 "authorization",
3211 "x-api-key",
3212 "anthropic-version",
3213 "openai-beta",
3214 "content-type",
3215 ];
3216 for name in FORWARDED {
3217 let Ok(header_name) = HeaderName::from_bytes(name.as_bytes()) else {
3218 continue;
3219 };
3220 if let Some(value) = headers.get(&header_name) {
3221 request = request.header(name.to_string(), value.clone());
3222 }
3223 }
3224 request
3225}
3226
3227fn proxy_error_response(err: ProxyError) -> Response {
3228 let status = match err {
3229 ProxyError::AdapterNotFound { .. } => StatusCode::NOT_FOUND,
3230 ProxyError::BodyTooLarge { .. } => StatusCode::PAYLOAD_TOO_LARGE,
3231 ProxyError::InvalidJson { .. } => StatusCode::BAD_REQUEST,
3232 ProxyError::UnsurfacedPii { .. } | ProxyError::Refused { .. } => {
3233 StatusCode::UNPROCESSABLE_ENTITY
3234 }
3235 ProxyError::UpstreamUnreachable { .. } => StatusCode::BAD_GATEWAY,
3236 _ => StatusCode::INTERNAL_SERVER_ERROR,
3237 };
3238 let mut body = serde_json::json!({
3239 "error": proxy_error_name(&err),
3240 });
3241 if let ProxyError::Refused { refusal } = &err {
3242 log_refusal(refusal);
3243 body["refusal"] = refusal_json(refusal);
3244 }
3245 (
3246 status,
3247 [(
3248 axum::http::header::CONTENT_TYPE,
3249 HeaderValue::from_static("application/json"),
3250 )],
3251 body.to_string(),
3252 )
3253 .into_response()
3254}
3255
3256fn proxy_error_name(err: &ProxyError) -> &'static str {
3257 match err {
3258 ProxyError::UpstreamUnreachable { .. } => "UpstreamUnreachable",
3259 ProxyError::AdapterNotFound { .. } => "AdapterNotFound",
3260 ProxyError::BodyTooLarge { .. } => "BodyTooLarge",
3261 ProxyError::InvalidJson { .. } => "InvalidJson",
3262 ProxyError::SsePartialFrame { .. } => "SsePartialFrame",
3263 ProxyError::UnsurfacedPii { .. } => "UnsurfacedPii",
3264 ProxyError::UnprovenCoverage => "UnprovenCoverage",
3265 ProxyError::DirectAdapterMismatch => "DirectAdapterMismatch",
3266 ProxyError::Refused { .. } => "Refused",
3267 ProxyError::Pipeline { .. } => "Pipeline",
3268 ProxyError::Server { .. } => "Server",
3269 ProxyError::DaemonAlreadyRunning { .. } => "DaemonAlreadyRunning",
3270 ProxyError::DaemonStopTimeout { .. } => "DaemonStopTimeout",
3271 ProxyError::DaemonNotRunning => "DaemonNotRunning",
3272 ProxyError::DaemonPidfileStale { .. } => "DaemonPidfileStale",
3273 ProxyError::DaemonIo { .. } => "DaemonIo",
3274 ProxyError::DaemonConfig { .. } => "DaemonConfig",
3275 ProxyError::DaemonExitedEarly { .. } => "DaemonExitedEarly",
3276 }
3277}
3278
3279#[cfg(test)]
3280mod tests {
3281 use super::*;
3282 use std::process::Command;
3283 use std::sync::atomic::{AtomicUsize, Ordering};
3284 use std::sync::Mutex;
3285
3286 use axum::body::to_bytes;
3287 use axum::routing::post;
3288 use gaze::{Action, ClassRule, DefaultRule, PiiClass};
3289 use gaze_recognizers::RegexDetector;
3290
3291 use crate::codec::OutputProvenance;
3292
3293 const SYNTHETIC_EMAIL: &str = "alice@example.invalid";
3294 const SYNTHETIC_PHONE: &str = "+1-555-0199";
3295
3296 struct SyntheticPrincipalResolver;
3297
3298 impl PrincipalResolver for SyntheticPrincipalResolver {
3299 fn resolve(
3300 &self,
3301 _context: PrincipalContext<'_>,
3302 ) -> Result<crate::AuthenticatedPrincipal, crate::principal::PrincipalResolveError>
3303 {
3304 crate::AuthenticatedPrincipal::try_from_canonical_bytes(b"synthetic-principal".to_vec())
3305 }
3306 }
3307
3308 #[derive(Clone, Default)]
3309 struct Capture {
3310 hits: Arc<AtomicUsize>,
3311 headers: Arc<Mutex<Option<HeaderMap>>>,
3312 body: Arc<Mutex<Vec<u8>>>,
3313 }
3314
3315 async fn capture_ok(
3316 State(capture): State<Capture>,
3317 headers: HeaderMap,
3318 body: Bytes,
3319 ) -> Response {
3320 capture.hits.fetch_add(1, Ordering::SeqCst);
3321 *capture.headers.lock().unwrap() = Some(headers);
3322 *capture.body.lock().unwrap() = body.to_vec();
3323 Response::builder()
3324 .status(StatusCode::OK)
3325 .header("content-type", "application/json")
3326 .header("set-cookie", "synthetic=drop")
3327 .header("x-request-id", "drop-me")
3328 .body(Body::from(br#"{"ok":true}"#.as_slice()))
3329 .unwrap()
3330 }
3331
3332 async fn spawn_test_server(app: Router) -> (Url, tokio::task::JoinHandle<()>) {
3333 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3334 let address = listener.local_addr().unwrap();
3335 let handle = tokio::spawn(async move {
3336 axum::serve(listener, app).await.unwrap();
3337 });
3338 (
3339 Url::parse(&format!("http://{address}/v1/messages")).unwrap(),
3340 handle,
3341 )
3342 }
3343
3344 fn proved(bytes: &[u8]) -> ProvedRequestBody {
3345 let mut provenance = OutputProvenance::default();
3346 provenance.mark_final_buffer_verified();
3347 ProvedRequestBody::new(bytes.to_vec(), WireFormat::Json, provenance)
3348 }
3349
3350 fn direct_request(
3351 url: Url,
3352 expected: WireFormat,
3353 bytes: &[u8],
3354 ) -> Result<DirectRequest, DirectProxyError> {
3355 let mut headers = HeaderMap::new();
3356 headers.insert("content-type", HeaderValue::from_static("application/json"));
3357 headers.insert("x-api-key", HeaderValue::from_static("synthetic-key"));
3358 headers.insert("anthropic-version", HeaderValue::from_static("2023-06-01"));
3359 DirectRequest::from_proved(url, headers, proved(bytes), expected)
3360 }
3361
3362 fn committed_direct_request(
3363 url: Url,
3364 expected: WireFormat,
3365 bytes: &[u8],
3366 ) -> CommittedDirectRequest {
3367 let session = Session::new(Scope::Ephemeral).unwrap();
3368 let transaction = session.begin_transaction();
3369 PreparedDirectRequest::new(transaction, direct_request(url, expected, bytes).unwrap())
3370 .commit()
3371 .unwrap()
3372 }
3373
3374 fn email_pipeline() -> Pipeline {
3375 Pipeline::builder()
3376 .detector(RegexDetector::emails().unwrap())
3377 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
3378 .rule(DefaultRule::new(Action::Preserve))
3379 .build()
3380 .unwrap()
3381 }
3382
3383 #[test]
3384 fn net_admission_failure_discards_direct_staging_but_retains_legacy_live_mappings() {
3385 struct RejectNet;
3386 impl gaze::SafetyNet for RejectNet {
3387 fn id(&self) -> &str {
3388 "synthetic-reject"
3389 }
3390 fn supported_locales(&self) -> &[LocaleTag] {
3391 &[LocaleTag::Global]
3392 }
3393 fn check(
3394 &self,
3395 text: &str,
3396 _: gaze::SafetyNetContext<'_>,
3397 ) -> Result<Vec<gaze::LeakSuspect>, gaze::SafetyNetError> {
3398 if text.contains("reject-this") {
3399 Err(gaze::SafetyNetError::Runtime {
3400 message: "synthetic failure".into(),
3401 })
3402 } else {
3403 Ok(vec![])
3404 }
3405 }
3406 }
3407 let pipeline = Pipeline::builder()
3408 .detector(RegexDetector::emails().unwrap())
3409 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
3410 .rule(DefaultRule::new(Action::Preserve))
3411 .register_safety_net(RejectNet)
3412 .enable_prefix_cache()
3413 .build()
3414 .unwrap();
3415 let session = Session::new(Scope::Ephemeral).unwrap();
3416 let body = format!(
3417 r#"{{"model":"claude-test","max_tokens":32,"system":"{SYNTHETIC_EMAIL}","messages":[{{"role":"user","content":"reject-this"}}]}}"#
3418 );
3419 let request = direct_request(
3420 Url::parse("http://127.0.0.1:1/v1/messages").unwrap(),
3421 WireFormat::Json,
3422 body.as_bytes(),
3423 )
3424 .unwrap();
3425 assert!(prepare_and_commit_direct_request(
3426 &pipeline,
3427 &[LocaleTag::Global],
3428 &DictionaryBundle::default(),
3429 &session,
3430 &crate::codecs::anthropic::AnthropicMessagesCodec,
3431 body.as_bytes(),
3432 &request.url,
3433 &request.headers,
3434 WireFormat::Json,
3435 CodecLimits::default(),
3436 false,
3437 DEFAULT_MAX_REQUEST_BYTES,
3438 )
3439 .is_err());
3440 assert!(session.tokens().is_empty());
3441 assert_eq!(session.prefix_cache_entry_count(), 0);
3442
3443 let mut first = SYNTHETIC_EMAIL.to_owned();
3444 let mut second = "reject-this".to_owned();
3445 assert!(redact_surfaces(
3446 &pipeline,
3447 &session,
3448 vec![
3449 crate::adapter::PiiSurface {
3450 field_path: "$.first".into(),
3451 text: &mut first,
3452 syntax: SurfaceSyntax::Text,
3453 },
3454 crate::adapter::PiiSurface {
3455 field_path: "$.second".into(),
3456 text: &mut second,
3457 syntax: SurfaceSyntax::Text,
3458 },
3459 ],
3460 &[LocaleTag::Global],
3461 &DictionaryBundle::default()
3462 )
3463 .is_err());
3464 assert_eq!(session.tokens().len(), 1);
3465 assert_eq!(
3466 session.restore_strict_text(&first).unwrap(),
3467 SYNTHETIC_EMAIL
3468 );
3469 }
3470
3471 const DIRECT_REQUEST_BODY: &[u8] = br#"{"model":"claude-test","max_tokens":32,"messages":[{"role":"user","content":"synthetic request"}]}"#;
3472 const DIRECT_SSE_REQUEST_BODY: &[u8] = br#"{"model":"claude-test","max_tokens":32,"messages":[{"role":"user","content":"synthetic request"}],"stream":true}"#;
3473 const DIRECT_JSON_RESPONSE: &[u8] = br#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-test","content":[{"type":"text","text":"synthetic response"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":1,"output_tokens":1}}"#;
3474 const DIRECT_SSE_RESPONSE: &[u8] = b"event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-test\",\"content\":[],\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"synthetic response\"}}\n\nevent: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\",\"stop_sequence\":null},\"usage\":{\"output_tokens\":1}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
3475 const NON_STREAM_SESSION_ID: &str = "00000000-0000-4000-8000-000000000003";
3476
3477 #[derive(Clone)]
3478 struct ResponseControl {
3479 hits: Arc<AtomicUsize>,
3480 body_started: Arc<tokio::sync::Notify>,
3481 release: Arc<tokio::sync::Notify>,
3482 stream_error: bool,
3483 }
3484
3485 impl ResponseControl {
3486 fn new(stream_error: bool) -> Self {
3487 Self {
3488 hits: Arc::new(AtomicUsize::new(0)),
3489 body_started: Arc::new(tokio::sync::Notify::new()),
3490 release: Arc::new(tokio::sync::Notify::new()),
3491 stream_error,
3492 }
3493 }
3494 }
3495
3496 async fn delayed_json_response(
3497 State(control): State<ResponseControl>,
3498 _body: Bytes,
3499 ) -> Response {
3500 control.hits.fetch_add(1, Ordering::SeqCst);
3501 control.body_started.notify_one();
3502 control.release.notified().await;
3503 let body = if control.stream_error {
3504 b"{".as_slice()
3505 } else {
3506 DIRECT_JSON_RESPONSE
3507 };
3508 Response::builder()
3509 .status(StatusCode::OK)
3510 .header("content-type", "application/json")
3511 .body(Body::from(body))
3512 .unwrap()
3513 }
3514
3515 async fn controlled_sse_response(
3516 State(control): State<ResponseControl>,
3517 _body: Bytes,
3518 ) -> Response {
3519 control.hits.fetch_add(1, Ordering::SeqCst);
3520 let body_control = control.clone();
3521 let chunks = futures_util::stream::once(async move {
3522 body_control.body_started.notify_one();
3523 body_control.release.notified().await;
3524 if body_control.stream_error {
3525 Err(std::io::Error::other("synthetic stream failure"))
3526 } else {
3527 Ok(Bytes::from_static(DIRECT_SSE_RESPONSE))
3528 }
3529 });
3530 Response::builder()
3531 .status(StatusCode::OK)
3532 .header("content-type", "text/event-stream")
3533 .body(Body::from_stream(chunks))
3534 .unwrap()
3535 }
3536
3537 fn synthetic_principal() -> crate::AuthenticatedPrincipal {
3538 crate::AuthenticatedPrincipal::try_from_canonical_bytes(b"synthetic-principal".to_vec())
3539 .unwrap()
3540 }
3541
3542 fn direct_test_state(
3543 endpoint: Url,
3544 pipeline: Pipeline,
3545 session_policy: SessionPolicy,
3546 registry: Option<SessionRegistry>,
3547 ) -> AppState {
3548 let client_config = DirectClientConfig::default();
3549 let runtime = DirectRuntime {
3550 adapter: ResolvedDirectAdapter::Compatibility(Arc::new(AnthropicAdapter::new(
3551 endpoint.clone(),
3552 ))),
3553 client: DirectClient::new(client_config).unwrap(),
3554 endpoint: endpoint.clone(),
3555 session_policy,
3556 registry,
3557 resolver: Arc::new(SyntheticPrincipalResolver),
3558 listener_scope: ListenerScope::Loopback,
3559 allowed_versions: Arc::from([DEFAULT_ANTHROPIC_VERSION.to_string()]),
3560 allowed_betas: Arc::from([]),
3561 codec_limits: CodecLimits::default(),
3562 client_config,
3563 ping_interval: Duration::from_secs(10),
3564 inspection_endpoint_codes: ProxyInspectionEndpointCodesV1::from_validated_origin(
3565 &endpoint,
3566 ),
3567 };
3568 AppState {
3569 config: Arc::new(ProxyConfig::new("127.0.0.1:0".parse().unwrap(), Vec::new())),
3570 pipeline: Arc::new(pipeline),
3571 client: Client::new(),
3572 direct: Some(Arc::new(runtime)),
3573 sessions: Arc::new(RwLock::new(HashMap::new())),
3574 started_at: Instant::now(),
3575 }
3576 }
3577
3578 #[tokio::test(flavor = "multi_thread", worker_threads = 8)]
3579 async fn concurrent_first_requests_share_one_session() {
3580 const TASK_COUNT: usize = 8;
3581 const SESSION_ID: &str = "synthetic-concurrent-session";
3582
3583 let state = Arc::new(AppState {
3584 config: Arc::new(ProxyConfig::new("127.0.0.1:0".parse().unwrap(), Vec::new())),
3585 pipeline: Arc::new(Pipeline::builder().build().unwrap()),
3586 client: Client::new(),
3587 direct: None,
3588 sessions: Arc::new(RwLock::new(HashMap::new())),
3589 started_at: Instant::now(),
3590 });
3591 SESSION_FOR_MISS_BARRIERS.lock().unwrap().insert(
3592 SESSION_ID.to_string(),
3593 Arc::new(tokio::sync::Barrier::new(TASK_COUNT)),
3594 );
3595
3596 let mut tasks = Vec::with_capacity(TASK_COUNT);
3597 for _ in 0..TASK_COUNT {
3598 let state = Arc::clone(&state);
3599 tasks.push(tokio::spawn(async move {
3600 let mut headers = HeaderMap::new();
3601 headers.insert(SESSION_HEADER, HeaderValue::from_static(SESSION_ID));
3602 session_for(&state, &headers).await.unwrap()
3603 }));
3604 }
3605
3606 let mut sessions = Vec::with_capacity(TASK_COUNT);
3607 for task in tasks {
3608 sessions.push(task.await.unwrap());
3609 }
3610 SESSION_FOR_MISS_BARRIERS.lock().unwrap().remove(SESSION_ID);
3611
3612 let first = &sessions[0];
3613 assert!(
3614 sessions.iter().all(|session| Arc::ptr_eq(first, session)),
3615 "concurrent first requests returned orphaned Sessions"
3616 );
3617 }
3618
3619 fn direct_test_ingress(
3620 endpoint: Url,
3621 session: DirectIngressSession,
3622 expected: WireFormat,
3623 body: &[u8],
3624 ) -> PreparedDirectIngress {
3625 let request = direct_request(endpoint, expected, body).unwrap();
3626 PreparedDirectIngress {
3627 outbound_headers: request.headers,
3628 session,
3629 }
3630 }
3631
3632 async fn wait_for_lease_count(registry: &SessionRegistry, expected: usize) {
3633 tokio::time::timeout(Duration::from_secs(60), async {
3634 loop {
3635 if registry.stats().total_leases() == expected {
3636 return;
3637 }
3638 tokio::task::yield_now().await;
3639 }
3640 })
3641 .await
3642 .unwrap();
3643 }
3644
3645 #[derive(Clone, Copy)]
3646 enum LimitFailureDirection {
3647 Request,
3648 Response,
3649 }
3650
3651 struct LimitFailureCodec(LimitFailureDirection);
3652
3653 impl BodyCodec for LimitFailureCodec {
3654 fn protect_request(
3655 &self,
3656 input: &[u8],
3657 _ctx: &mut RequestTransformContext<'_>,
3658 ) -> Result<ProvedRequestBody, CodecError> {
3659 match self.0 {
3660 LimitFailureDirection::Request => Err(CodecError::new(
3661 CodecErrorCode::DepthLimitExceeded,
3662 CodecPhase::RequestParse,
3663 )),
3664 LimitFailureDirection::Response => Ok(proved(input)),
3665 }
3666 }
3667
3668 fn restore_response(
3669 &self,
3670 _input: &[u8],
3671 _ctx: &mut ResponseTransformContext<'_>,
3672 ) -> Result<crate::codec::ProvedResponseBody, CodecError> {
3673 Err(CodecError::new(
3674 CodecErrorCode::DepthLimitExceeded,
3675 CodecPhase::ResponseParse,
3676 ))
3677 }
3678 }
3679
3680 struct CustomCodecAdapter {
3681 anthropic: AnthropicAdapter,
3682 codec: LimitFailureCodec,
3683 }
3684
3685 #[async_trait::async_trait]
3686 impl ProviderAdapter for CustomCodecAdapter {
3687 fn contract(&self) -> crate::AdapterContract<'_> {
3688 crate::AdapterContract::codec(
3689 &self.codec,
3690 SessionPolicy::EphemeralSingleRequest,
3691 CodecLimits::default(),
3692 )
3693 }
3694 fn name(&self) -> &'static str {
3695 "custom-codec"
3696 }
3697 fn matches_path(&self, method: &Method, path: &str) -> bool {
3698 self.anthropic.matches_path(method, path)
3699 }
3700 fn upstream_base(&self) -> &Url {
3701 self.anthropic.upstream_base()
3702 }
3703 fn request_pii_surfaces<'a>(&self, body: &'a mut Value) -> Vec<crate::PiiSurface<'a>> {
3704 self.anthropic.request_pii_surfaces(body)
3705 }
3706 fn response_pii_surfaces<'a>(&self, body: &'a mut Value) -> Vec<crate::PiiSurface<'a>> {
3707 self.anthropic.response_pii_surfaces(body)
3708 }
3709 fn sse_event_pii_surfaces<'a>(
3710 &self,
3711 event: &'a mut SseEvent,
3712 ) -> Vec<crate::PiiSurface<'a>> {
3713 self.anthropic.sse_event_pii_surfaces(event)
3714 }
3715 }
3716
3717 #[tokio::test]
3718 async fn compatibility_constructor_uses_its_custom_codec_before_upstream_io() {
3719 let capture = Capture::default();
3720 let (mut endpoint, upstream) = spawn_test_server(
3721 Router::new()
3722 .route("/v1/messages", post(capture_ok))
3723 .with_state(capture.clone()),
3724 )
3725 .await;
3726 endpoint.set_path("/");
3727 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3728 let bind = listener.local_addr().unwrap();
3729 let config = ProxyConfig::new(
3730 bind,
3731 vec![Arc::new(CustomCodecAdapter {
3732 anthropic: AnthropicAdapter::new(endpoint),
3733 codec: LimitFailureCodec(LimitFailureDirection::Request),
3734 })],
3735 );
3736 let server = serve_with_listener(config, Arc::new(email_pipeline()), listener);
3737 let request = async {
3738 let response = Client::new()
3739 .post(format!("http://{bind}/v1/messages"))
3740 .header("x-api-key", "synthetic-provider-key")
3741 .header("anthropic-version", DEFAULT_ANTHROPIC_VERSION)
3742 .header("content-type", "application/json")
3743 .body(DIRECT_REQUEST_BODY)
3744 .send()
3745 .await
3746 .unwrap();
3747 assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
3748 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
3749 };
3750 tokio::select! {
3751 result = server => panic!("compatibility server stopped: {result:?}"),
3752 result = tokio::time::timeout(Duration::from_secs(60), request) => result.unwrap(),
3753 }
3754 upstream.abort();
3755 }
3756
3757 async fn begin_delayed_non_stream(
3758 invalid_response: bool,
3759 ) -> (
3760 tokio::task::JoinHandle<Result<Response, DirectProxyError>>,
3761 ResponseControl,
3762 SessionRegistry,
3763 Instant,
3764 tokio::task::JoinHandle<()>,
3765 ) {
3766 let control = ResponseControl::new(invalid_response);
3767 let (endpoint, server) = spawn_test_server(
3768 Router::new()
3769 .route("/v1/messages", post(delayed_json_response))
3770 .with_state(control.clone()),
3771 )
3772 .await;
3773 let registry_config = crate::SessionRegistryConfig::default();
3774 let registry = SessionRegistry::with_secret(registry_config, [0x7c; 32]);
3775 let request_lease = registry
3776 .acquire(synthetic_principal(), NON_STREAM_SESSION_ID, None)
3777 .unwrap();
3778 let state = direct_test_state(
3779 endpoint.clone(),
3780 email_pipeline(),
3781 SessionPolicy::OpaqueHeaderContinuity(registry_config),
3782 Some(registry.clone()),
3783 );
3784 let ingress = direct_test_ingress(
3785 endpoint,
3786 DirectIngressSession::Continuity(request_lease),
3787 WireFormat::Json,
3788 DIRECT_REQUEST_BODY,
3789 );
3790 let request = tokio::spawn(async move {
3791 direct_proxy_inner(
3792 &state,
3793 ingress,
3794 DIRECT_REQUEST_BODY,
3795 &crate::codecs::anthropic::AnthropicMessagesCodec,
3796 )
3797 .await
3798 });
3799 control.body_started.notified().await;
3800 let expired_at = Instant::now() + Duration::from_secs(60 * 60);
3801 (request, control, registry, expired_at, server)
3802 }
3803
3804 async fn begin_controlled_sse(
3805 stream_error: bool,
3806 ) -> (
3807 Response,
3808 ResponseControl,
3809 SessionRegistry,
3810 Instant,
3811 tokio::task::JoinHandle<()>,
3812 ) {
3813 const SESSION_ID: &str = "00000000-0000-4000-8000-000000000002";
3814
3815 let control = ResponseControl::new(stream_error);
3816 let (endpoint, server) = spawn_test_server(
3817 Router::new()
3818 .route("/v1/messages", post(controlled_sse_response))
3819 .with_state(control.clone()),
3820 )
3821 .await;
3822 let registry_config = crate::SessionRegistryConfig::default();
3823 let registry = SessionRegistry::with_secret(registry_config, [0x6b; 32]);
3824 let request_lease = registry
3825 .acquire(synthetic_principal(), SESSION_ID, None)
3826 .unwrap();
3827 let state = direct_test_state(
3828 endpoint.clone(),
3829 email_pipeline(),
3830 SessionPolicy::OpaqueHeaderContinuity(registry_config),
3831 Some(registry.clone()),
3832 );
3833 let ingress = direct_test_ingress(
3834 endpoint,
3835 DirectIngressSession::Continuity(request_lease),
3836 WireFormat::Sse,
3837 DIRECT_SSE_REQUEST_BODY,
3838 );
3839
3840 let response = direct_proxy_inner(
3841 &state,
3842 ingress,
3843 DIRECT_SSE_REQUEST_BODY,
3844 &crate::codecs::anthropic::AnthropicMessagesCodec,
3845 )
3846 .await
3847 .unwrap();
3848 control.body_started.notified().await;
3849 let expired_at = Instant::now() + Duration::from_secs(60 * 60);
3850 (response, control, registry, expired_at, server)
3851 }
3852
3853 async fn assert_owned_listener_health(config: ProxyConfig, listener: TcpListener) {
3854 let actual = listener.local_addr().unwrap();
3855 let destination = SocketAddr::from(([127, 0, 0, 1], actual.port()));
3856 let server = serve_with_listener(config, Arc::new(email_pipeline()), listener);
3857 let health = async {
3858 let response = Client::new()
3859 .get(format!("http://{destination}/_gaze_proxy/healthz"))
3860 .send()
3861 .await
3862 .unwrap()
3863 .error_for_status()
3864 .unwrap()
3865 .json::<serde_json::Value>()
3866 .await
3867 .unwrap();
3868 assert_eq!(response["bind"], actual.to_string());
3869 assert_eq!(response["adapters"][0]["name"], "anthropic");
3870 };
3871 tokio::select! {
3872 result = server => panic!("owned listener server stopped: {result:?}"),
3873 result = tokio::time::timeout(Duration::from_secs(60), health) => result.unwrap(),
3874 }
3875 }
3877
3878 #[tokio::test]
3879 async fn reserved_port_rebind_fails_but_owned_listener_handoff_serves() {
3880 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3881 let bind = listener.local_addr().unwrap();
3882 let config = ProxyConfig::anthropic_direct(
3883 bind,
3884 AnthropicAdapter::new(Url::parse("http://127.0.0.1:1").unwrap()),
3885 );
3886 let error = serve(config.clone(), Arc::new(email_pipeline()))
3887 .await
3888 .unwrap_err();
3889 assert!(matches!(error, ProxyError::Server { source }
3890 if source.kind() == std::io::ErrorKind::AddrInUse));
3891 assert_owned_listener_health(config, listener).await;
3892 }
3893
3894 #[tokio::test]
3895 async fn owned_listener_resolves_ephemeral_port_in_health() {
3896 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3897 let config = ProxyConfig::anthropic_direct(
3898 "127.0.0.1:0".parse().unwrap(),
3899 AnthropicAdapter::new(Url::parse("http://127.0.0.1:1").unwrap()),
3900 );
3901 assert_owned_listener_health(config, listener).await;
3902 }
3903
3904 #[tokio::test]
3905 async fn owned_listener_rejects_wrong_ip_wildcard_and_port() {
3906 for (listen, configured) in [
3907 ("0.0.0.0:0", "127.0.0.1:0"),
3908 ("127.0.0.1:0", "0.0.0.0:0"),
3909 ("127.0.0.1:0", "127.0.0.2:0"),
3910 ("127.0.0.1:0", "127.0.0.1:1"),
3911 ] {
3912 let listener = TcpListener::bind(listen).await.unwrap();
3913 let config = ProxyConfig::anthropic_direct(
3914 configured.parse().unwrap(),
3915 AnthropicAdapter::new(Url::parse("http://127.0.0.1:1").unwrap()),
3916 );
3917 let error = serve_with_listener(config, Arc::new(email_pipeline()), listener)
3918 .await
3919 .unwrap_err();
3920 assert!(matches!(error, ProxyError::Server { source }
3921 if source.kind() == std::io::ErrorKind::InvalidInput));
3922 }
3923 }
3924
3925 #[tokio::test]
3926 async fn owned_non_loopback_listener_requires_explicit_resolver() {
3927 let listener = TcpListener::bind("0.0.0.0:0").await.unwrap();
3928 let config = ProxyConfig::anthropic_direct(
3929 "0.0.0.0:0".parse().unwrap(),
3930 AnthropicAdapter::new(Url::parse("http://127.0.0.1:1").unwrap()),
3931 );
3932 let error = serve_with_listener(config, Arc::new(email_pipeline()), listener)
3933 .await
3934 .unwrap_err();
3935 assert!(matches!(error, ProxyError::DaemonConfig { detail }
3936 if detail == "direct_principal_required"));
3937
3938 let listener = TcpListener::bind("0.0.0.0:0").await.unwrap();
3939 let config = ProxyConfig::anthropic_direct(
3940 "0.0.0.0:0".parse().unwrap(),
3941 AnthropicAdapter::builder(Url::parse("http://127.0.0.1:1").unwrap())
3942 .principal_resolver(Arc::new(SyntheticPrincipalResolver))
3943 .build()
3944 .unwrap(),
3945 );
3946 assert_owned_listener_health(config, listener).await;
3947 }
3948
3949 #[tokio::test]
3950 async fn serve_validates_before_rebind_and_owned_listener_validates_before_serving() {
3951 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3952 let config = ProxyConfig::anthropic_direct(
3953 listener.local_addr().unwrap(),
3954 AnthropicAdapter::new(Url::parse("http://synthetic:credential@127.0.0.1").unwrap()),
3955 );
3956 for error in [
3957 serve(config.clone(), Arc::new(email_pipeline()))
3958 .await
3959 .unwrap_err(),
3960 serve_with_listener(config, Arc::new(email_pipeline()), listener)
3961 .await
3962 .unwrap_err(),
3963 ] {
3964 assert!(matches!(error, ProxyError::DaemonConfig { detail }
3965 if detail == "direct_upstream_invalid"));
3966 }
3967 }
3968
3969 #[test]
3970 fn non_loopback_readiness_requires_an_explicit_principal_resolver() {
3971 let bind = "0.0.0.0:8787".parse().unwrap();
3972 let upstream = Url::parse("https://api.anthropic.com").unwrap();
3973 let missing = ProxyConfig::anthropic_direct(bind, AnthropicAdapter::new(upstream.clone()));
3974 assert!(DirectRuntime::from_config(&missing).is_err());
3975
3976 let configured = ProxyConfig::anthropic_direct(
3977 bind,
3978 AnthropicAdapter::builder(upstream)
3979 .principal_resolver(Arc::new(SyntheticPrincipalResolver))
3980 .build()
3981 .unwrap(),
3982 );
3983 assert!(DirectRuntime::from_config(&configured).unwrap().is_some());
3984 }
3985
3986 #[test]
3987 fn invalid_direct_upstream_fails_readiness_before_any_request() {
3988 let config = ProxyConfig::anthropic_direct(
3989 "127.0.0.1:8787".parse().unwrap(),
3990 AnthropicAdapter::new(Url::parse("https://synthetic:credential@127.0.0.1").unwrap()),
3991 );
3992 assert!(DirectRuntime::from_config(&config).is_err());
3993 }
3994
3995 #[tokio::test]
3996 async fn direct_send_requires_commit_is_single_use_and_captures_exact_proof() {
3997 let capture = Capture::default();
3998 let (url, server) = spawn_test_server(
3999 Router::new()
4000 .route("/v1/messages", post(capture_ok))
4001 .with_state(capture.clone()),
4002 )
4003 .await;
4004 let exact = br#"{"messages":[{"role":"user","content":"<Email_1>"}]}"#;
4005 let client = DirectClient::new(DirectClientConfig::default()).unwrap();
4006
4007 let session = Session::new(Scope::Ephemeral).unwrap();
4008 let prepared = PreparedDirectRequest::new(
4009 session.begin_transaction(),
4010 direct_request(url.clone(), WireFormat::Json, exact).unwrap(),
4011 );
4012 drop(prepared);
4013 assert!(session.tokens().is_empty());
4014 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4015
4016 let mut lifecycle = DirectTransactionState::new();
4017 lifecycle.retry_generation_conflict().unwrap();
4018 assert_eq!(lifecycle.generation_retries(), 1);
4019 assert_eq!(
4020 lifecycle.retry_generation_conflict().unwrap_err().code(),
4021 ProxyErrorCode::SessionGenerationConflict
4022 );
4023 let response = client
4024 .execute(committed_direct_request(url, WireFormat::Json, exact))
4025 .await
4026 .unwrap();
4027 assert_eq!(response.lifecycle.state(), &DirectIoState::IoAttempted);
4028 assert_eq!(capture.hits.load(Ordering::SeqCst), 1);
4029 assert_eq!(&*capture.body.lock().unwrap(), exact);
4030 assert_eq!(response.response.body(), br#"{"ok":true}"#);
4031 assert_eq!(response.response.headers().len(), 1);
4032 assert!(response.response.headers().get("set-cookie").is_none());
4033 server.abort();
4034 }
4035
4036 #[test]
4037 fn signed_shadow_probe_reuses_only_the_previously_staged_mapping() {
4038 let pipeline = Pipeline::builder()
4039 .detector(RegexDetector::emails().unwrap())
4040 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
4041 .rule(DefaultRule::new(Action::Preserve))
4042 .enable_prefix_cache()
4043 .build()
4044 .unwrap();
4045 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4046 let session = Session::new(Scope::Ephemeral).unwrap();
4047 let expected_token = format!("<{}:Email_1>", session.session_hex());
4048 let body = format!(
4049 r#"{{"model":"claude-test","max_tokens":32,"system":"{SYNTHETIC_EMAIL}","messages":[{{"role":"assistant","content":[{{"type":"thinking","thinking":"{expected_token}","signature":"opaque-signature"}}]}}]}}"#
4050 );
4051 let request = direct_request(
4052 Url::parse("https://api.anthropic.com/v1/messages").unwrap(),
4053 WireFormat::Json,
4054 body.as_bytes(),
4055 )
4056 .unwrap();
4057
4058 let committed = prepare_and_commit_direct_request_with_hook(
4059 &pipeline,
4060 &[LocaleTag::Global],
4061 &DictionaryBundle::default(),
4062 &session,
4063 &codec,
4064 body.as_bytes(),
4065 &request.url,
4066 &request.headers,
4067 WireFormat::Json,
4068 CodecLimits::default(),
4069 false,
4070 DEFAULT_MAX_REQUEST_BYTES,
4071 |attempt, live, staged| {
4072 assert_eq!(attempt, 0);
4073 assert_eq!(staged.prefix_cache_entry_count(), 0);
4074 assert_eq!(live.prefix_cache_entry_count(), 0);
4075 Ok(())
4076 },
4077 )
4078 .unwrap();
4079
4080 assert_eq!(session.tokens(), vec![expected_token.clone()]);
4081 assert_eq!(committed.snapshot.tokens(), vec![expected_token.clone()]);
4082 assert_eq!(committed.snapshot.prefix_cache_entry_count(), 0);
4083 assert_eq!(session.prefix_cache_entry_count(), 0);
4084 let protected = std::str::from_utf8(&committed.request.body).unwrap();
4085 assert!(!protected.contains(SYNTHETIC_EMAIL));
4086 assert_eq!(protected.matches(&expected_token).count(), 2);
4087 assert!(protected.contains(&format!(
4088 r#""thinking":"{expected_token}","signature":"opaque-signature""#
4089 )));
4090 }
4091
4092 #[test]
4093 fn final_request_reproof_is_idempotent_for_ordinary_and_format_preserving_tokens() {
4094 let endpoint = Url::parse("https://api.anthropic.com/v1/messages").unwrap();
4095 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4096
4097 for action in [Action::Tokenize, Action::FormatPreserve] {
4098 let pipeline = Pipeline::builder()
4099 .detector(RegexDetector::emails().unwrap())
4100 .rule(ClassRule::new(PiiClass::Email, action))
4101 .rule(DefaultRule::new(Action::Preserve))
4102 .build()
4103 .unwrap();
4104 let session = Session::new(Scope::Ephemeral).unwrap();
4105 let original = format!(
4106 r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":"{SYNTHETIC_EMAIL}"}}]}}"#
4107 );
4108 let headers = direct_request(endpoint.clone(), WireFormat::Json, original.as_bytes())
4109 .unwrap()
4110 .headers;
4111 let first = prepare_and_commit_direct_request(
4112 &pipeline,
4113 &[LocaleTag::Global],
4114 &DictionaryBundle::default(),
4115 &session,
4116 &codec,
4117 original.as_bytes(),
4118 &endpoint,
4119 &headers,
4120 WireFormat::Json,
4121 CodecLimits::default(),
4122 false,
4123 DEFAULT_MAX_REQUEST_BYTES,
4124 )
4125 .unwrap_or_else(|error| panic!("{action:?} first proof failed: {error:?}"));
4126 let proved_once = first.request.body.clone();
4127 assert_eq!(session.tokens().len(), 1);
4128
4129 let second = prepare_and_commit_direct_request(
4130 &pipeline,
4131 &[LocaleTag::Global],
4132 &DictionaryBundle::default(),
4133 &session,
4134 &codec,
4135 &proved_once,
4136 &endpoint,
4137 &headers,
4138 WireFormat::Json,
4139 CodecLimits::default(),
4140 false,
4141 DEFAULT_MAX_REQUEST_BYTES,
4142 )
4143 .unwrap();
4144 assert_eq!(second.request.body, proved_once);
4145 assert_eq!(session.tokens().len(), 1);
4146 }
4147 }
4148
4149 #[tokio::test]
4150 async fn late_signed_failure_discards_all_staged_state_and_does_zero_io() {
4151 let capture = Capture::default();
4152 let (url, server) = spawn_test_server(
4153 Router::new()
4154 .route("/v1/messages", post(capture_ok))
4155 .with_state(capture.clone()),
4156 )
4157 .await;
4158 let pipeline = Pipeline::builder()
4159 .detector(RegexDetector::emails().unwrap())
4160 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
4161 .rule(DefaultRule::new(Action::Preserve))
4162 .enable_prefix_cache()
4163 .build()
4164 .unwrap();
4165 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4166 let session = Session::new(Scope::Ephemeral).unwrap();
4167 let body = format!(
4168 r#"{{"model":"claude-test","max_tokens":32,"system":"{SYNTHETIC_EMAIL}","messages":[{{"role":"assistant","content":[{{"type":"thinking","thinking":"{SYNTHETIC_EMAIL}","signature":"opaque-signature"}}]}}]}}"#
4169 );
4170 let request = direct_request(url, WireFormat::Json, body.as_bytes()).unwrap();
4171 let error = prepare_and_commit_direct_request(
4172 &pipeline,
4173 &[LocaleTag::Global],
4174 &DictionaryBundle::default(),
4175 &session,
4176 &codec,
4177 body.as_bytes(),
4178 &request.url,
4179 &request.headers,
4180 WireFormat::Json,
4181 CodecLimits::default(),
4182 false,
4183 DEFAULT_MAX_REQUEST_BYTES,
4184 )
4185 .unwrap_err();
4186
4187 assert_eq!(error.code(), ProxyErrorCode::SignedMutationRequired);
4188 assert!(session.tokens().is_empty());
4189 assert_eq!(session.prefix_cache_entry_count(), 0);
4190 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4191
4192 let valid = format!(
4193 r#"{{"model":"claude-test","max_tokens":32,"system":"{SYNTHETIC_EMAIL} reports","messages":[{{"role":"user","content":"synthetic"}}]}}"#
4194 );
4195 prepare_and_commit_direct_request(
4196 &pipeline,
4197 &[LocaleTag::Global],
4198 &DictionaryBundle::default(),
4199 &session,
4200 &codec,
4201 valid.as_bytes(),
4202 &request.url,
4203 &request.headers,
4204 WireFormat::Json,
4205 CodecLimits::default(),
4206 false,
4207 DEFAULT_MAX_REQUEST_BYTES,
4208 )
4209 .unwrap();
4210 assert_eq!(session.tokens().len(), 1);
4211 server.abort();
4212 }
4213
4214 #[tokio::test]
4215 async fn logical_domain_failure_discards_state_before_commit_and_does_zero_io() {
4216 let capture = Capture::default();
4217 let (url, server) = spawn_test_server(
4218 Router::new()
4219 .route("/v1/messages", post(capture_ok))
4220 .with_state(capture.clone()),
4221 )
4222 .await;
4223 let pipeline = Pipeline::builder()
4224 .detector(RegexDetector::emails().unwrap())
4225 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
4226 .rule(DefaultRule::new(Action::Preserve))
4227 .enable_prefix_cache()
4228 .build()
4229 .unwrap();
4230 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4231 let session = Session::new(Scope::Ephemeral).unwrap();
4232 let split = 6;
4233 let (left, right) = SYNTHETIC_EMAIL.split_at(split);
4234
4235 let mut prefix_transaction = session.begin_transaction();
4236 pipeline
4237 .pseudonymize_transaction_with_detect_context(
4238 &mut prefix_transaction,
4239 RawDocument::Text(left.to_owned()),
4240 &[LocaleTag::Global],
4241 &DictionaryBundle::default(),
4242 )
4243 .unwrap();
4244 prefix_transaction.commit().unwrap();
4245 let baseline_entries = session.snapshot_entries();
4246 let baseline_tokens = session.tokens();
4247 let baseline_cache_entries = session.prefix_cache_entry_count();
4248 assert_eq!(baseline_cache_entries, 0);
4249
4250 let body = format!(
4251 r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":[{{"type":"text","text":"{left}"}},{{"type":"text","text":"{right}"}}]}}]}}"#
4252 );
4253 let request = direct_request(url, WireFormat::Json, body.as_bytes()).unwrap();
4254 let mut before_commit_calls = 0usize;
4255 let error = match prepare_and_commit_direct_request_with_hook(
4256 &pipeline,
4257 &[LocaleTag::Global],
4258 &DictionaryBundle::default(),
4259 &session,
4260 &codec,
4261 body.as_bytes(),
4262 &request.url,
4263 &request.headers,
4264 WireFormat::Json,
4265 CodecLimits::default(),
4266 false,
4267 DEFAULT_MAX_REQUEST_BYTES,
4268 |_, _, _| {
4269 before_commit_calls += 1;
4270 Ok(())
4271 },
4272 ) {
4273 Ok(_) => panic!("expected closed request proof failure"),
4274 Err(error) => error,
4275 };
4276
4277 assert_eq!(error.code(), ProxyErrorCode::ControlWouldMutate);
4278 assert_eq!(error.phase(), ProxyErrorPhase::RequestTransform);
4279 assert_eq!(error.http_status(), StatusCode::UNPROCESSABLE_ENTITY);
4280 assert_eq!(before_commit_calls, 0);
4281 assert_eq!(session.snapshot_entries(), baseline_entries);
4282 assert_eq!(session.tokens(), baseline_tokens);
4283 assert_eq!(session.prefix_cache_entry_count(), baseline_cache_entries);
4284 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4285 server.abort();
4286 }
4287
4288 #[test]
4289 fn suppression_only_request_sites_do_not_populate_the_prefix_cache() {
4290 let pipeline = Pipeline::builder()
4291 .detector(RegexDetector::emails().unwrap())
4292 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
4293 .rule(DefaultRule::new(Action::Preserve))
4294 .enable_prefix_cache()
4295 .build()
4296 .unwrap();
4297 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4298 let session = Session::new(Scope::Ephemeral).unwrap();
4299 let body = br#"{"model":"claude-test","max_tokens":32,"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"synthetic reasoning","signature":"opaque-signature","cache_control":{"type":"ephemeral"}}]}],"tools":[{"name":"synthetic_tool","input_schema":{"type":"object","$defs":{"synthetic_definition":{"type":"string"}}}}]}"#;
4300 let request = direct_request(
4301 Url::parse("https://api.anthropic.com/v1/messages").unwrap(),
4302 WireFormat::Json,
4303 body,
4304 )
4305 .unwrap();
4306
4307 let committed = prepare_and_commit_direct_request_with_hook(
4308 &pipeline,
4309 &[LocaleTag::Global],
4310 &DictionaryBundle::default(),
4311 &session,
4312 &codec,
4313 body,
4314 &request.url,
4315 &request.headers,
4316 WireFormat::Json,
4317 CodecLimits::default(),
4318 false,
4319 DEFAULT_MAX_REQUEST_BYTES,
4320 |attempt, live, staged| {
4321 assert_eq!(attempt, 0);
4322 assert_eq!(live.prefix_cache_entry_count(), 0);
4323 assert_eq!(staged.prefix_cache_entry_count(), 0);
4324 Ok(())
4325 },
4326 )
4327 .unwrap();
4328
4329 assert!(session.tokens().is_empty());
4330 assert_eq!(committed.snapshot.prefix_cache_entry_count(), 0);
4331 assert_eq!(session.prefix_cache_entry_count(), 0);
4332 }
4333
4334 #[test]
4335 fn every_logical_domain_probe_family_leaves_prefix_storage_empty() {
4336 let pipeline = Pipeline::builder()
4337 .detector(RegexDetector::emails().unwrap())
4338 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
4339 .rule(DefaultRule::new(Action::Preserve))
4340 .enable_prefix_cache()
4341 .build()
4342 .unwrap();
4343 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4344 let session = Session::new(Scope::Ephemeral).unwrap();
4345 let body = br#"{"model":"claude-test","max_tokens":32,"system":"system-safe","tools":[{"name":"tool-safe-name","description":"tool.safe-description","input_schema":{"type":"object","description":"schema-safe-description","properties":{"property-safe":{"type":"string"}},"required":["property-safe"]}}],"metadata":{"metadata-safe-key":"metadata-safe-value"},"messages":[{"role":"user","content":"message-safe"}]}"#;
4346 let request = direct_request(
4347 Url::parse("https://api.anthropic.com/v1/messages").unwrap(),
4348 WireFormat::Json,
4349 body,
4350 )
4351 .unwrap();
4352
4353 let committed = prepare_and_commit_direct_request_with_hook(
4354 &pipeline,
4355 &[LocaleTag::Global],
4356 &DictionaryBundle::default(),
4357 &session,
4358 &codec,
4359 body,
4360 &request.url,
4361 &request.headers,
4362 WireFormat::Json,
4363 CodecLimits::default(),
4364 false,
4365 DEFAULT_MAX_REQUEST_BYTES,
4366 |attempt, live, staged| {
4367 assert_eq!(attempt, 0);
4368 assert_eq!(live.prefix_cache_entry_count(), 0);
4369 assert_eq!(staged.prefix_cache_entry_count(), 0);
4370 Ok(())
4371 },
4372 )
4373 .unwrap();
4374
4375 assert_eq!(committed.snapshot.prefix_cache_entry_count(), 0);
4376 assert_eq!(session.prefix_cache_entry_count(), 0);
4377 }
4378
4379 struct RejectOnSecondAttemptCodec {
4380 protect_calls: AtomicUsize,
4381 }
4382
4383 impl BodyCodec for RejectOnSecondAttemptCodec {
4384 fn protect_request(
4385 &self,
4386 input: &[u8],
4387 ctx: &mut RequestTransformContext<'_>,
4388 ) -> Result<ProvedRequestBody, CodecError> {
4389 if self.protect_calls.fetch_add(1, Ordering::SeqCst) == 0 {
4390 crate::codecs::anthropic::AnthropicMessagesCodec.protect_request(input, ctx)
4391 } else {
4392 Err(CodecError::new(
4393 CodecErrorCode::ControlWouldMutate,
4394 CodecPhase::RequestProof,
4395 ))
4396 }
4397 }
4398
4399 fn restore_response(
4400 &self,
4401 input: &[u8],
4402 ctx: &mut ResponseTransformContext<'_>,
4403 ) -> Result<crate::codec::ProvedResponseBody, CodecError> {
4404 crate::codecs::anthropic::AnthropicMessagesCodec.restore_response(input, ctx)
4405 }
4406 }
4407
4408 #[tokio::test]
4409 async fn generation_conflict_reproofs_once_and_sends_only_the_winning_body() {
4410 let capture = Capture::default();
4411 let (url, server) = spawn_test_server(
4412 Router::new()
4413 .route("/v1/messages", post(capture_ok))
4414 .with_state(capture.clone()),
4415 )
4416 .await;
4417 let pipeline = email_pipeline();
4418 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4419 let session = Session::new(Scope::Ephemeral).unwrap();
4420 let body = format!(
4421 r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":"{SYNTHETIC_EMAIL}"}}]}}"#
4422 );
4423 let request = direct_request(url.clone(), WireFormat::Json, body.as_bytes()).unwrap();
4424 let mut hook_calls = 0;
4425 let committed = prepare_and_commit_direct_request_with_hook(
4426 &pipeline,
4427 &[LocaleTag::Global],
4428 &DictionaryBundle::default(),
4429 &session,
4430 &codec,
4431 body.as_bytes(),
4432 &request.url,
4433 &request.headers,
4434 WireFormat::Json,
4435 CodecLimits::default(),
4436 false,
4437 DEFAULT_MAX_REQUEST_BYTES,
4438 |attempt, live, _staged| {
4439 hook_calls += 1;
4440 if attempt == 0 {
4441 live.tokenize(&PiiClass::Name, "Dr. Schmidt").unwrap();
4442 }
4443 Ok(())
4444 },
4445 )
4446 .unwrap();
4447 assert_eq!(hook_calls, 2);
4448 assert_eq!(session.tokens().len(), 2);
4449
4450 DirectClient::new(DirectClientConfig::default())
4451 .unwrap()
4452 .execute(committed)
4453 .await
4454 .unwrap();
4455 assert_eq!(capture.hits.load(Ordering::SeqCst), 1);
4456 let captured = String::from_utf8(capture.body.lock().unwrap().clone()).unwrap();
4457 assert!(!captured.contains(SYNTHETIC_EMAIL));
4458 assert_eq!(captured.matches(":Email_1>").count(), 1);
4459 server.abort();
4460 }
4461
4462 #[tokio::test]
4463 async fn second_attempt_request_proof_failure_rejects_without_upstream_io() {
4464 let capture = Capture::default();
4465 let (url, server) = spawn_test_server(
4466 Router::new()
4467 .route("/v1/messages", post(capture_ok))
4468 .with_state(capture.clone()),
4469 )
4470 .await;
4471 let pipeline = Pipeline::builder()
4472 .rule(DefaultRule::new(Action::Preserve))
4473 .build()
4474 .unwrap();
4475 let codec = RejectOnSecondAttemptCodec {
4476 protect_calls: AtomicUsize::new(0),
4477 };
4478 let session = Session::new(Scope::Ephemeral).unwrap();
4479 let body = br#"{"model":"claude-test","max_tokens":32,"messages":[{"role":"user","content":"synthetic"}]}"#;
4480 let request = direct_request(url, WireFormat::Json, body).unwrap();
4481 let mut hook_calls = 0usize;
4482
4483 let error = prepare_and_commit_direct_request_with_hook(
4484 &pipeline,
4485 &[LocaleTag::Global],
4486 &DictionaryBundle::default(),
4487 &session,
4488 &codec,
4489 body,
4490 &request.url,
4491 &request.headers,
4492 WireFormat::Json,
4493 CodecLimits::default(),
4494 false,
4495 DEFAULT_MAX_REQUEST_BYTES,
4496 |attempt, live, _staged| {
4497 hook_calls += 1;
4498 assert_eq!(attempt, 0);
4499 live.tokenize(&PiiClass::Name, "Dr. Schmidt").unwrap();
4500 Ok(())
4501 },
4502 )
4503 .unwrap_err();
4504
4505 assert_eq!(codec.protect_calls.load(Ordering::SeqCst), 2);
4506 assert_eq!(hook_calls, 1);
4507 assert_eq!(error.code(), ProxyErrorCode::ControlWouldMutate);
4508 assert_eq!(error.phase(), ProxyErrorPhase::RequestTransform);
4509 assert_eq!(error.http_status(), StatusCode::UNPROCESSABLE_ENTITY);
4510 assert_eq!(session.tokens().len(), 1);
4511 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4512 server.abort();
4513 }
4514
4515 #[tokio::test]
4516 async fn second_generation_conflict_fails_closed_without_upstream_io() {
4517 let capture = Capture::default();
4518 let (url, server) = spawn_test_server(
4519 Router::new()
4520 .route("/v1/messages", post(capture_ok))
4521 .with_state(capture.clone()),
4522 )
4523 .await;
4524 let pipeline = Pipeline::builder()
4525 .detector(RegexDetector::emails().unwrap())
4526 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
4527 .rule(DefaultRule::new(Action::Preserve))
4528 .enable_prefix_cache()
4529 .build()
4530 .unwrap();
4531 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4532 let session = Session::new(Scope::Ephemeral).unwrap();
4533 let body = format!(
4534 r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":"{SYNTHETIC_EMAIL}"}}]}}"#
4535 );
4536 let request = direct_request(url, WireFormat::Json, body.as_bytes()).unwrap();
4537 let error = prepare_and_commit_direct_request_with_hook(
4538 &pipeline,
4539 &[LocaleTag::Global],
4540 &DictionaryBundle::default(),
4541 &session,
4542 &codec,
4543 body.as_bytes(),
4544 &request.url,
4545 &request.headers,
4546 WireFormat::Json,
4547 CodecLimits::default(),
4548 false,
4549 DEFAULT_MAX_REQUEST_BYTES,
4550 |attempt, live, staged| {
4551 assert_eq!(staged.prefix_cache_entry_count(), 0);
4552 assert_eq!(live.prefix_cache_entry_count(), 0);
4553 live.tokenize(&PiiClass::Name, &format!("Dr. Schmidt {attempt}"))
4554 .unwrap();
4555 Ok(())
4556 },
4557 )
4558 .unwrap_err();
4559 assert_eq!(error.code(), ProxyErrorCode::SessionGenerationConflict);
4560 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4561 assert_eq!(session.prefix_cache_entry_count(), 0);
4562 assert!(session
4563 .tokens()
4564 .iter()
4565 .all(|token| !token.contains(":Email_")));
4566 server.abort();
4567 }
4568
4569 #[tokio::test]
4570 async fn cancellation_publishes_only_after_commit_and_never_rolls_it_back() {
4571 let session = Session::new(Scope::Ephemeral).unwrap();
4572 let mut staged = session.begin_transaction();
4573 staged.tokenize(&PiiClass::Email, SYNTHETIC_EMAIL).unwrap();
4574 let prepared = PreparedDirectRequest::new(
4575 staged,
4576 direct_request(
4577 Url::parse("https://api.anthropic.com/v1/messages").unwrap(),
4578 WireFormat::Json,
4579 br#"{}"#,
4580 )
4581 .unwrap(),
4582 );
4583 drop(prepared);
4584 assert!(session.tokens().is_empty());
4585
4586 let (url, server) = spawn_test_server(Router::new().route(
4587 "/v1/messages",
4588 post(|| async {
4589 tokio::time::sleep(Duration::from_secs(1)).await;
4590 ([("content-type", "application/json")], "{}")
4591 }),
4592 ))
4593 .await;
4594 let mut staged = session.begin_transaction();
4595 staged.tokenize(&PiiClass::Email, SYNTHETIC_EMAIL).unwrap();
4596 let committed = PreparedDirectRequest::new(
4597 staged,
4598 direct_request(url, WireFormat::Json, br#"{}"#).unwrap(),
4599 )
4600 .commit()
4601 .unwrap();
4602 assert_eq!(session.tokens().len(), 1);
4603 let client = DirectClient::new(DirectClientConfig::default()).unwrap();
4604 assert!(
4605 tokio::time::timeout(Duration::from_millis(20), client.execute(committed))
4606 .await
4607 .is_err()
4608 );
4609 assert_eq!(session.tokens().len(), 1);
4610 server.abort();
4611 }
4612
4613 #[tokio::test]
4614 async fn response_restore_uses_the_exact_committed_snapshot_after_live_mutation() {
4615 let pipeline = email_pipeline();
4616 let session = Session::new(Scope::Ephemeral).unwrap();
4617 let mut transaction = session.begin_transaction();
4618 let committed_token = transaction
4619 .tokenize(&PiiClass::Email, SYNTHETIC_EMAIL)
4620 .unwrap();
4621 let snapshot = transaction.commit().unwrap();
4622 session.tokenize(&PiiClass::Name, "Dr. Schmidt").unwrap();
4623 assert_eq!(session.tokens().len(), 2);
4624
4625 let body = format!(
4626 r#"{{"id":"msg_1","type":"message","role":"assistant","model":"claude-test","content":[{{"type":"text","text":"{committed_token}"}}],"stop_reason":"end_turn","stop_sequence":null,"usage":{{"input_tokens":1,"output_tokens":1}}}}"#
4627 );
4628 let mut headers = HeaderMap::new();
4629 headers.insert(
4630 axum::http::header::CONTENT_TYPE,
4631 HeaderValue::from_static("application/json"),
4632 );
4633 let executed = CommittedDirectResponse {
4634 snapshot,
4635 response: DirectResponse {
4636 head: DirectResponseHead {
4637 status: StatusCode::OK,
4638 format: WireFormat::Json,
4639 headers,
4640 },
4641 body: body.into_bytes(),
4642 },
4643 lifecycle: DirectTransactionState {
4644 state: DirectIoState::IoAttempted,
4645 generation_retries: 0,
4646 },
4647 };
4648 let response = restore_direct_response(
4649 &pipeline,
4650 &[LocaleTag::Global],
4651 &DictionaryBundle::default(),
4652 &crate::codecs::anthropic::AnthropicMessagesCodec,
4653 executed,
4654 CodecLimits::default(),
4655 None,
4656 )
4657 .unwrap();
4658 let restored = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
4659 assert!(std::str::from_utf8(&restored)
4660 .unwrap()
4661 .contains(SYNTHETIC_EMAIL));
4662 assert!(!std::str::from_utf8(&restored)
4663 .unwrap()
4664 .contains("Dr. Schmidt"));
4665 }
4666
4667 #[tokio::test]
4668 async fn transformed_request_overlimit_rejects_before_commit_and_zero_io() {
4669 const SESSION_ID: &str = "00000000-0000-4000-8000-000000000001";
4670
4671 let capture = Capture::default();
4672 let (url, server) = spawn_test_server(
4673 Router::new()
4674 .route("/v1/messages", post(capture_ok))
4675 .with_state(capture.clone()),
4676 )
4677 .await;
4678 let phone_class = PiiClass::Custom("phone".to_string());
4679 let pipeline = Pipeline::builder()
4680 .detector(RegexDetector::new(r"\+1-555-0199", phone_class.clone()).unwrap())
4681 .rule(ClassRule::new(phone_class, Action::Tokenize))
4682 .rule(DefaultRule::new(Action::Preserve))
4683 .enable_prefix_cache()
4684 .build()
4685 .unwrap();
4686 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4687 let body = format!(
4688 r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":"{SYNTHETIC_PHONE}"}}]}}"#
4689 );
4690 let request_limit = body.len() + 1;
4691 let probe_session = Session::new(Scope::Ephemeral).unwrap();
4692 let probe_request = direct_request(url.clone(), WireFormat::Json, body.as_bytes()).unwrap();
4693 let transformed = prepare_and_commit_direct_request(
4694 &pipeline,
4695 &[LocaleTag::Global],
4696 &DictionaryBundle::default(),
4697 &probe_session,
4698 &codec,
4699 body.as_bytes(),
4700 &probe_request.url,
4701 &probe_request.headers,
4702 WireFormat::Json,
4703 CodecLimits::default(),
4704 false,
4705 DEFAULT_MAX_REQUEST_BYTES,
4706 )
4707 .unwrap();
4708 assert!(body.len() < request_limit);
4709 assert!(transformed.request.body.len() > request_limit);
4710
4711 let registry_config = crate::SessionRegistryConfig::default();
4712 let registry = SessionRegistry::with_secret(registry_config, [0x5a; 32]);
4713 let observed_lease = registry
4714 .acquire(
4715 crate::AuthenticatedPrincipal::try_from_canonical_bytes(
4716 b"synthetic-principal".to_vec(),
4717 )
4718 .unwrap(),
4719 SESSION_ID,
4720 None,
4721 )
4722 .unwrap();
4723 let request_lease = registry
4724 .acquire(
4725 crate::AuthenticatedPrincipal::try_from_canonical_bytes(
4726 b"synthetic-principal".to_vec(),
4727 )
4728 .unwrap(),
4729 SESSION_ID,
4730 None,
4731 )
4732 .unwrap();
4733 let client_config = DirectClientConfig::default()
4734 .try_with_max_request_bytes(request_limit)
4735 .unwrap();
4736 let runtime = DirectRuntime {
4737 adapter: ResolvedDirectAdapter::Compatibility(Arc::new(AnthropicAdapter::new(
4738 url.clone(),
4739 ))),
4740 client: DirectClient::new(client_config).unwrap(),
4741 endpoint: url.clone(),
4742 session_policy: SessionPolicy::OpaqueHeaderContinuity(registry_config),
4743 registry: Some(registry),
4744 resolver: Arc::new(SyntheticPrincipalResolver),
4745 listener_scope: ListenerScope::Loopback,
4746 allowed_versions: Arc::from([DEFAULT_ANTHROPIC_VERSION.to_string()]),
4747 allowed_betas: Arc::from([]),
4748 codec_limits: CodecLimits::default(),
4749 client_config,
4750 ping_interval: Duration::from_secs(10),
4751 inspection_endpoint_codes: ProxyInspectionEndpointCodesV1::from_validated_origin(&url),
4752 };
4753 let state = AppState {
4754 config: Arc::new(ProxyConfig::new("127.0.0.1:0".parse().unwrap(), Vec::new())),
4755 pipeline: Arc::new(pipeline),
4756 client: Client::new(),
4757 direct: Some(Arc::new(runtime)),
4758 sessions: Arc::new(RwLock::new(HashMap::new())),
4759 started_at: Instant::now(),
4760 };
4761 let ingress = PreparedDirectIngress {
4762 outbound_headers: probe_request.headers.clone(),
4763 session: DirectIngressSession::Continuity(request_lease),
4764 };
4765
4766 let error = direct_proxy_inner(&state, ingress, body.as_bytes(), &codec)
4767 .await
4768 .unwrap_err();
4769
4770 assert_eq!(error.code(), ProxyErrorCode::RequestBodyLimitExceeded);
4771 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4772 assert!(observed_lease.session().tokens().is_empty());
4773 assert!(observed_lease.session().snapshot_entries().is_empty());
4774
4775 let extended_body = format!(
4776 r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":"{SYNTHETIC_PHONE} reports"}}]}}"#
4777 );
4778 let recovered = prepare_and_commit_direct_request(
4779 &state.pipeline,
4780 state.config.locale_chain(),
4781 state.config.dictionaries(),
4782 observed_lease.session(),
4783 &codec,
4784 extended_body.as_bytes(),
4785 &url,
4786 &probe_request.headers,
4787 WireFormat::Json,
4788 CodecLimits::default(),
4789 false,
4790 DEFAULT_MAX_REQUEST_BYTES,
4791 )
4792 .unwrap();
4793 let first_token = format!(
4794 "<{}:Custom:phone_1>",
4795 observed_lease.session().session_hex()
4796 );
4797 assert_eq!(observed_lease.session().tokens(), vec![first_token.clone()]);
4798 assert!(std::str::from_utf8(&recovered.request.body)
4799 .unwrap()
4800 .contains(&first_token));
4801 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4802 server.abort();
4803 }
4804
4805 #[tokio::test]
4806 async fn overlimit_proved_request_and_duplicate_credentials_do_zero_io() {
4807 let capture = Capture::default();
4808 let (url, server) = spawn_test_server(
4809 Router::new()
4810 .route("/v1/messages", post(capture_ok))
4811 .with_state(capture.clone()),
4812 )
4813 .await;
4814 let config = DirectClientConfig::default()
4815 .try_with_max_request_bytes(4)
4816 .unwrap();
4817 let client = DirectClient::new(config).unwrap();
4818 let error = client
4819 .execute(committed_direct_request(
4820 url.clone(),
4821 WireFormat::Json,
4822 br#"{"over":true}"#,
4823 ))
4824 .await
4825 .unwrap_err();
4826 assert_eq!(error.code(), ProxyErrorCode::RequestBodyLimitExceeded);
4827 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4828
4829 for singleton in ["x-api-key", "anthropic-version", "anthropic-beta"] {
4830 let mut duplicate = HeaderMap::new();
4831 duplicate.insert("content-type", HeaderValue::from_static("application/json"));
4832 duplicate.append(singleton, HeaderValue::from_static("synthetic-value-1"));
4833 duplicate.append(singleton, HeaderValue::from_static("synthetic-value-2"));
4834 let error = DirectRequest::from_proved(
4835 url.clone(),
4836 duplicate,
4837 proved(br#"{}"#),
4838 WireFormat::Json,
4839 )
4840 .err()
4841 .unwrap();
4842 assert_eq!(error.code(), ProxyErrorCode::HeaderRejected);
4843 }
4844
4845 let mut unverified_headers = HeaderMap::new();
4846 unverified_headers.insert("content-type", HeaderValue::from_static("application/json"));
4847 let unverified = ProvedRequestBody::new(
4848 br#"{}"#.to_vec(),
4849 WireFormat::Json,
4850 OutputProvenance::default(),
4851 );
4852 let error =
4853 DirectRequest::from_proved(url, unverified_headers, unverified, WireFormat::Json)
4854 .err()
4855 .unwrap();
4856 assert_eq!(error.code(), ProxyErrorCode::InvalidProvenance);
4857 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4858 server.abort();
4859 }
4860
4861 #[tokio::test]
4862 async fn header_rejections_name_only_the_header_and_disambiguate_the_cause() {
4863 const SECRET_ONE: &str = "synthetic-credential-one";
4864 const SECRET_TWO: &str = "synthetic-credential-two";
4865
4866 let mut not_allowlisted = HeaderMap::new();
4867 not_allowlisted.insert("x-synthetic-extra", HeaderValue::from_static(SECRET_ONE));
4868
4869 let mut empty = HeaderMap::new();
4870 empty.insert("anthropic-beta", HeaderValue::from_static(""));
4871
4872 let mut duplicate = HeaderMap::new();
4873 duplicate.append("x-api-key", HeaderValue::from_static(SECRET_ONE));
4874 duplicate.append("x-api-key", HeaderValue::from_static(SECRET_TWO));
4875
4876 let cases = [
4877 (
4878 not_allowlisted,
4879 "x-synthetic-extra",
4880 HeaderRejectionReason::NameNotAllowlisted,
4881 ),
4882 (empty, "anthropic-beta", HeaderRejectionReason::EmptyValue),
4883 (
4884 duplicate,
4885 "x-api-key",
4886 HeaderRejectionReason::DuplicateSingleton,
4887 ),
4888 ];
4889
4890 for (headers, expected_name, expected_reason) in cases {
4891 let error = validate_outbound_headers(&headers).unwrap_err();
4892 assert_eq!(error.code(), ProxyErrorCode::HeaderRejected);
4893 assert_eq!(error.phase(), ProxyErrorPhase::RequestValidation);
4894 assert_eq!(error.header_name().as_deref(), Some(expected_name));
4895 assert_eq!(error.header_rejection_reason(), Some(expected_reason));
4896
4897 let debug = format!("{error:?}");
4898 assert!(debug.contains(expected_name));
4899 assert!(!debug.contains(SECRET_ONE));
4900 assert!(!debug.contains(SECRET_TWO));
4901
4902 let response = direct_error_response(error);
4903 let bytes = to_bytes(response.into_body(), 4 * 1024).await.unwrap();
4904 let body: Value = serde_json::from_slice(&bytes).unwrap();
4905 assert_eq!(body["error"]["phase"], "RequestValidation");
4906 assert_eq!(body["error"]["header"]["name"], expected_name);
4907 assert_eq!(body["error"]["header"]["reason"], expected_reason.as_str());
4908 let rendered = String::from_utf8(bytes.to_vec()).unwrap();
4909 assert!(!rendered.contains(SECRET_ONE));
4910 assert!(!rendered.contains(SECRET_TWO));
4911 }
4912
4913 let maximum_name = format!("x-{}", "a".repeat(126));
4914 let mut maximum_length = HeaderMap::new();
4915 maximum_length.insert(
4916 HeaderName::from_bytes(maximum_name.as_bytes()).unwrap(),
4917 HeaderValue::from_static(SECRET_ONE),
4918 );
4919 let error = validate_outbound_headers(&maximum_length).unwrap_err();
4920 assert_eq!(error.header_name().as_deref(), Some(maximum_name.as_str()));
4921 assert_eq!(
4922 error.header_rejection_reason(),
4923 Some(HeaderRejectionReason::NameNotAllowlisted)
4924 );
4925 }
4926
4927 #[test]
4928 fn actionable_header_reporting_preserves_the_rejected_request_set() {
4929 let mut valid = HeaderMap::new();
4930 valid.insert("content-type", HeaderValue::from_static("application/json"));
4931 valid.insert("x-api-key", HeaderValue::from_static("synthetic-key"));
4932 valid.insert("anthropic-version", HeaderValue::from_static("2023-06-01"));
4933
4934 let mut not_allowlisted = valid.clone();
4935 not_allowlisted.insert(
4936 "x-synthetic-extra",
4937 HeaderValue::from_static("synthetic-value"),
4938 );
4939
4940 let mut empty = valid.clone();
4941 empty.insert("anthropic-version", HeaderValue::from_static(""));
4942
4943 let mut duplicate = valid.clone();
4944 duplicate.append("x-api-key", HeaderValue::from_static("synthetic-key-two"));
4945
4946 let rejected = [valid, not_allowlisted, empty, duplicate]
4947 .iter()
4948 .map(|headers| validate_outbound_headers(headers).is_err())
4949 .collect::<Vec<_>>();
4950 assert_eq!(rejected, vec![false, true, true, true]);
4951 }
4952
4953 #[tokio::test]
4954 async fn redirects_retries_and_response_overflow_remain_disabled() {
4955 let target_hits = Arc::new(AtomicUsize::new(0));
4956 let target_hits_for_handler = Arc::clone(&target_hits);
4957 let (target_url, target) = spawn_test_server(Router::new().route(
4958 "/v1/messages",
4959 post(move || {
4960 let hits = Arc::clone(&target_hits_for_handler);
4961 async move {
4962 hits.fetch_add(1, Ordering::SeqCst);
4963 StatusCode::OK
4964 }
4965 }),
4966 ))
4967 .await;
4968 let origin_hits = Arc::new(AtomicUsize::new(0));
4969 let origin_hits_for_handler = Arc::clone(&origin_hits);
4970 let location = target_url.to_string();
4971 let (origin_url, origin) = spawn_test_server(Router::new().route(
4972 "/v1/messages",
4973 post(move || {
4974 let hits = Arc::clone(&origin_hits_for_handler);
4975 let location = location.clone();
4976 async move {
4977 hits.fetch_add(1, Ordering::SeqCst);
4978 Response::builder()
4979 .status(StatusCode::TEMPORARY_REDIRECT)
4980 .header("location", location)
4981 .body(Body::from("synthetic redirect body"))
4982 .unwrap()
4983 }
4984 }),
4985 ))
4986 .await;
4987 let client = DirectClient::new(DirectClientConfig::default()).unwrap();
4988 let error = client
4989 .execute(committed_direct_request(
4990 origin_url,
4991 WireFormat::Json,
4992 br#"{}"#,
4993 ))
4994 .await
4995 .unwrap_err();
4996 assert_eq!(error.code(), ProxyErrorCode::UpstreamRedirect);
4997 assert_eq!(origin_hits.load(Ordering::SeqCst), 1);
4998 assert_eq!(target_hits.load(Ordering::SeqCst), 0);
4999 origin.abort();
5000 target.abort();
5001
5002 let attempts = Arc::new(AtomicUsize::new(0));
5003 let attempts_for_handler = Arc::clone(&attempts);
5004 let (url, server) = spawn_test_server(Router::new().route(
5005 "/v1/messages",
5006 post(move || {
5007 let attempts = Arc::clone(&attempts_for_handler);
5008 async move {
5009 attempts.fetch_add(1, Ordering::SeqCst);
5010 StatusCode::SERVICE_UNAVAILABLE
5011 }
5012 }),
5013 ))
5014 .await;
5015 let error = client
5016 .execute(committed_direct_request(url, WireFormat::Json, br#"{}"#))
5017 .await
5018 .unwrap_err();
5019 assert_eq!(error.code(), ProxyErrorCode::UpstreamUnavailable);
5020 assert_eq!(attempts.load(Ordering::SeqCst), 1);
5021 server.abort();
5022
5023 let chunks = futures_util::stream::iter([
5024 Ok::<_, std::convert::Infallible>(Bytes::from_static(b"12345")),
5025 Ok(Bytes::from_static(b"67890")),
5026 ]);
5027 let body = Arc::new(Mutex::new(Some(Body::from_stream(chunks))));
5028 let body_for_handler = Arc::clone(&body);
5029 let (url, server) = spawn_test_server(Router::new().route(
5030 "/v1/messages",
5031 post(move || {
5032 let body = body_for_handler.lock().unwrap().take().unwrap();
5033 async move {
5034 Response::builder()
5035 .status(StatusCode::OK)
5036 .header("content-type", "application/json")
5037 .body(body)
5038 .unwrap()
5039 }
5040 }),
5041 ))
5042 .await;
5043 let config = DirectClientConfig::default()
5044 .try_with_max_response_bytes(8)
5045 .unwrap();
5046 let client = DirectClient::new(config).unwrap();
5047 let error = client
5048 .execute(committed_direct_request(url, WireFormat::Json, br#"{}"#))
5049 .await
5050 .unwrap_err();
5051 assert_eq!(error.code(), ProxyErrorCode::ResponseBodyLimitExceeded);
5052 server.abort();
5053 }
5054
5055 #[tokio::test]
5056 async fn timeout_and_body_stream_error_retain_io_attempted() {
5057 let (slow_url, slow_server) = spawn_test_server(Router::new().route(
5058 "/v1/messages",
5059 post(|| async {
5060 tokio::time::sleep(Duration::from_millis(100)).await;
5061 ([("content-type", "application/json")], "{}")
5062 }),
5063 ))
5064 .await;
5065 let config = DirectClientConfig::default()
5066 .try_with_timeouts(
5067 Duration::from_millis(20),
5068 Duration::from_millis(80),
5069 Duration::from_millis(30),
5070 )
5071 .unwrap();
5072 let client = DirectClient::new(config).unwrap();
5073 let error = client
5074 .execute(committed_direct_request(
5075 slow_url,
5076 WireFormat::Json,
5077 br#"{}"#,
5078 ))
5079 .await
5080 .unwrap_err();
5081 assert_eq!(error.code(), ProxyErrorCode::TotalTimeout);
5082 slow_server.abort();
5083
5084 let chunks =
5085 futures_util::stream::once(async { Ok::<_, std::io::Error>(Bytes::from_static(b"{")) })
5086 .chain(futures_util::stream::once(async {
5087 tokio::time::sleep(Duration::from_millis(20)).await;
5088 Err(std::io::Error::other("synthetic stream failure"))
5089 }));
5090 let body = Body::from_stream(chunks);
5091 let body = Arc::new(Mutex::new(Some(body)));
5092 let body_for_handler = Arc::clone(&body);
5093 let (url, server) = spawn_test_server(Router::new().route(
5094 "/v1/messages",
5095 post(move || {
5096 let body = body_for_handler.lock().unwrap().take().unwrap();
5097 async move {
5098 Response::builder()
5099 .status(StatusCode::OK)
5100 .header("content-type", "application/json")
5101 .body(body)
5102 .unwrap()
5103 }
5104 }),
5105 ))
5106 .await;
5107 let client = DirectClient::new(DirectClientConfig::default()).unwrap();
5108 let error = client
5109 .execute(committed_direct_request(url, WireFormat::Json, br#"{}"#))
5110 .await
5111 .unwrap_err();
5112 assert_eq!(error.phase(), ProxyErrorPhase::UpstreamBody);
5113 server.abort();
5114 }
5115
5116 #[tokio::test]
5117 async fn codec_limit_direction_pins_code_phase_status_body_and_io() {
5118 let capture = Capture::default();
5119 let (endpoint, server) = spawn_test_server(
5120 Router::new()
5121 .route("/v1/messages", post(capture_ok))
5122 .with_state(capture.clone()),
5123 )
5124 .await;
5125 let state = direct_test_state(
5126 endpoint.clone(),
5127 email_pipeline(),
5128 SessionPolicy::EphemeralSingleRequest,
5129 None,
5130 );
5131
5132 let request_error = direct_proxy_inner(
5133 &state,
5134 direct_test_ingress(
5135 endpoint.clone(),
5136 DirectIngressSession::Ephemeral,
5137 WireFormat::Json,
5138 DIRECT_REQUEST_BODY,
5139 ),
5140 DIRECT_REQUEST_BODY,
5141 &LimitFailureCodec(LimitFailureDirection::Request),
5142 )
5143 .await
5144 .unwrap_err();
5145 assert_eq!(
5146 request_error.code(),
5147 ProxyErrorCode::RequestBodyLimitExceeded
5148 );
5149 assert_eq!(request_error.phase(), ProxyErrorPhase::RequestValidation);
5150 assert_eq!(request_error.http_status(), StatusCode::PAYLOAD_TOO_LARGE);
5151 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
5152 let request_response = direct_error_response(request_error);
5153 assert_eq!(request_response.status(), StatusCode::PAYLOAD_TOO_LARGE);
5154 let request_body = to_bytes(request_response.into_body(), 4 * 1024)
5155 .await
5156 .unwrap();
5157 assert_eq!(
5158 serde_json::from_slice::<serde_json::Value>(&request_body).unwrap(),
5159 serde_json::json!({
5160 "type": "error",
5161 "error": {
5162 "type": "api_error",
5163 "message": "proxy_validation_failed",
5164 "code": "RequestBodyLimitExceeded",
5165 "phase": "RequestValidation"
5166 }
5167 })
5168 );
5169
5170 let response_error = direct_proxy_inner(
5171 &state,
5172 direct_test_ingress(
5173 endpoint,
5174 DirectIngressSession::Ephemeral,
5175 WireFormat::Json,
5176 DIRECT_REQUEST_BODY,
5177 ),
5178 DIRECT_REQUEST_BODY,
5179 &LimitFailureCodec(LimitFailureDirection::Response),
5180 )
5181 .await
5182 .unwrap_err();
5183 assert_eq!(
5184 response_error.code(),
5185 ProxyErrorCode::ResponseBodyLimitExceeded
5186 );
5187 assert_eq!(response_error.phase(), ProxyErrorPhase::ResponseValidation);
5188 assert_eq!(response_error.http_status(), StatusCode::BAD_GATEWAY);
5189 assert_eq!(capture.hits.load(Ordering::SeqCst), 1);
5190 let response = direct_error_response(response_error);
5191 assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
5192 let response_body = to_bytes(response.into_body(), 4 * 1024).await.unwrap();
5193 assert_eq!(
5194 serde_json::from_slice::<serde_json::Value>(&response_body).unwrap(),
5195 serde_json::json!({
5196 "type": "error",
5197 "error": {
5198 "type": "api_error",
5199 "message": "proxy_validation_failed",
5200 "code": "ResponseBodyLimitExceeded",
5201 "phase": "ResponseValidation"
5202 }
5203 })
5204 );
5205 server.abort();
5206 }
5207
5208 #[test]
5209 fn future_non_conflict_session_commit_failure_is_never_a_conflict() {
5210 let generation_conflict =
5211 map_session_commit_failure(SessionCommitFailureKind::GenerationConflict);
5212 assert_eq!(
5213 generation_conflict.code(),
5214 ProxyErrorCode::SessionGenerationConflict
5215 );
5216 assert_eq!(generation_conflict.phase(), ProxyErrorPhase::Session);
5217 assert_eq!(generation_conflict.http_status(), StatusCode::CONFLICT);
5218
5219 let non_conflict = map_session_commit_failure(SessionCommitFailureKind::NonConflict);
5220 assert_eq!(non_conflict.code(), ProxyErrorCode::SessionCommitFailure);
5221 assert_eq!(non_conflict.phase(), ProxyErrorPhase::Session);
5222 assert_eq!(
5223 non_conflict.http_status(),
5224 StatusCode::INTERNAL_SERVER_ERROR
5225 );
5226 assert_ne!(non_conflict.code(), generation_conflict.code());
5227 }
5228
5229 #[tokio::test]
5230 async fn continuity_lease_spans_non_stream_response_without_holding_request_mutex() {
5231 let (request, control, registry, expired_at, server) =
5232 begin_delayed_non_stream(false).await;
5233
5234 assert_eq!(control.hits.load(Ordering::SeqCst), 1);
5235 assert_eq!(registry.stats().total_leases(), 1);
5236 let concurrent_lease = registry
5237 .acquire(synthetic_principal(), NON_STREAM_SESSION_ID, None)
5238 .unwrap();
5239 let concurrent_request_guard = concurrent_lease
5240 .try_lock_request()
5241 .expect("request mutex must be free before upstream I/O");
5242 drop(concurrent_request_guard);
5243 drop(concurrent_lease);
5244 assert_eq!(registry.stats().total_leases(), 1);
5245
5246 registry.cleanup_at(expired_at);
5247 assert_eq!(registry.stats().active_entries(), 1);
5248 assert_eq!(registry.stats().tombstones(), 0);
5249
5250 control.release.notify_one();
5251 let response = request.await.unwrap().unwrap();
5252 assert_eq!(response.status(), StatusCode::OK);
5253 wait_for_lease_count(®istry, 0).await;
5254 registry.cleanup_at(expired_at);
5255 assert_eq!(registry.stats().active_entries(), 0);
5256 assert_eq!(registry.stats().tombstones(), 1);
5257 server.abort();
5258 }
5259
5260 #[tokio::test]
5261 async fn continuity_lease_releases_on_non_stream_error() {
5262 let (request, control, registry, expired_at, server) = begin_delayed_non_stream(true).await;
5263 assert_eq!(registry.stats().total_leases(), 1);
5264 registry.cleanup_at(expired_at);
5265 assert_eq!(registry.stats().active_entries(), 1);
5266
5267 control.release.notify_one();
5268 let error = request.await.unwrap().unwrap_err();
5269 assert_eq!(error.phase(), ProxyErrorPhase::ResponseValidation);
5270 assert_eq!(error.http_status(), StatusCode::BAD_GATEWAY);
5271 wait_for_lease_count(®istry, 0).await;
5272 registry.cleanup_at(expired_at);
5273 assert_eq!(registry.stats().active_entries(), 0);
5274 assert_eq!(registry.stats().tombstones(), 1);
5275 server.abort();
5276 }
5277
5278 #[tokio::test]
5279 async fn continuity_lease_releases_on_non_stream_cancellation() {
5280 let (request, _control, registry, expired_at, server) =
5281 begin_delayed_non_stream(false).await;
5282 assert_eq!(registry.stats().total_leases(), 1);
5283 registry.cleanup_at(expired_at);
5284 assert_eq!(registry.stats().active_entries(), 1);
5285
5286 request.abort();
5287 assert!(request.await.unwrap_err().is_cancelled());
5288 wait_for_lease_count(®istry, 0).await;
5289 registry.cleanup_at(expired_at);
5290 assert_eq!(registry.stats().active_entries(), 0);
5291 assert_eq!(registry.stats().tombstones(), 1);
5292 server.abort();
5293 }
5294
5295 #[tokio::test]
5296 async fn continuity_lease_spans_sse_message_stop_then_becomes_expirable() {
5297 let (response, control, registry, expired_at, server) = begin_controlled_sse(false).await;
5298 assert_eq!(control.hits.load(Ordering::SeqCst), 1);
5299 assert_eq!(registry.stats().total_leases(), 1);
5300 registry.cleanup_at(expired_at);
5301 assert_eq!(registry.stats().active_entries(), 1);
5302 assert_eq!(registry.stats().tombstones(), 0);
5303
5304 control.release.notify_one();
5305 let body = to_bytes(response.into_body(), DEFAULT_MAX_RESPONSE_BYTES)
5306 .await
5307 .unwrap();
5308 assert_eq!(body.as_ref(), DIRECT_SSE_RESPONSE);
5309 wait_for_lease_count(®istry, 0).await;
5310 registry.cleanup_at(expired_at);
5311 assert_eq!(registry.stats().active_entries(), 0);
5312 assert_eq!(registry.stats().tombstones(), 1);
5313 server.abort();
5314 }
5315
5316 #[tokio::test]
5317 async fn continuity_lease_releases_once_on_sse_error() {
5318 let (response, control, registry, expired_at, server) = begin_controlled_sse(true).await;
5319 assert_eq!(registry.stats().total_leases(), 1);
5320 registry.cleanup_at(expired_at);
5321 assert_eq!(registry.stats().active_entries(), 1);
5322
5323 control.release.notify_one();
5324 let body = to_bytes(response.into_body(), DEFAULT_MAX_RESPONSE_BYTES)
5325 .await
5326 .unwrap();
5327 assert_eq!(body.as_ref(), DIRECT_PROXY_ERROR_FRAME);
5328 wait_for_lease_count(®istry, 0).await;
5329 registry.cleanup_at(expired_at);
5330 assert_eq!(registry.stats().active_entries(), 0);
5331 assert_eq!(registry.stats().tombstones(), 1);
5332 server.abort();
5333 }
5334
5335 #[tokio::test]
5336 async fn continuity_lease_releases_once_on_sse_downstream_cancellation() {
5337 let (response, _control, registry, expired_at, server) = begin_controlled_sse(false).await;
5338 assert_eq!(registry.stats().total_leases(), 1);
5339 registry.cleanup_at(expired_at);
5340 assert_eq!(registry.stats().active_entries(), 1);
5341
5342 drop(response);
5343 wait_for_lease_count(®istry, 0).await;
5344 registry.cleanup_at(expired_at);
5345 assert_eq!(registry.stats().active_entries(), 0);
5346 assert_eq!(registry.stats().tombstones(), 1);
5347 server.abort();
5348 }
5349
5350 #[test]
5351 fn proxy_environment_is_ignored_in_isolated_helper() {
5352 let output = Command::new(std::env::current_exe().unwrap())
5353 .args([
5354 "--exact",
5355 "server::tests::proxy_environment_helper",
5356 "--nocapture",
5357 ])
5358 .env("GAZE_PROXY_ENV_HELPER", "1")
5359 .output()
5360 .unwrap();
5361 assert!(
5362 output.status.success(),
5363 "isolated proxy helper failed: {}",
5364 String::from_utf8_lossy(&output.stderr)
5365 );
5366 }
5367
5368 #[test]
5369 fn proxy_environment_helper() {
5370 if std::env::var_os("GAZE_PROXY_ENV_HELPER").is_none() {
5371 return;
5372 }
5373 tokio::runtime::Runtime::new().unwrap().block_on(async {
5374 let capture = Capture::default();
5375 let (url, server) = spawn_test_server(
5376 Router::new()
5377 .route("/v1/messages", post(capture_ok))
5378 .with_state(capture.clone()),
5379 )
5380 .await;
5381 let trap = TcpListener::bind("127.0.0.1:0").await.unwrap();
5382 let proxy = format!("http://{}", trap.local_addr().unwrap());
5383 for name in [
5384 "HTTP_PROXY",
5385 "HTTPS_PROXY",
5386 "ALL_PROXY",
5387 "http_proxy",
5388 "https_proxy",
5389 ] {
5390 std::env::set_var(name, &proxy);
5391 }
5392 std::env::remove_var("NO_PROXY");
5393 std::env::remove_var("no_proxy");
5394
5395 let client = DirectClient::new(DirectClientConfig::default()).unwrap();
5396 client
5397 .execute(committed_direct_request(url, WireFormat::Json, br#"{}"#))
5398 .await
5399 .unwrap();
5400 assert_eq!(capture.hits.load(Ordering::SeqCst), 1);
5401 assert!(
5402 tokio::time::timeout(Duration::from_millis(100), trap.accept())
5403 .await
5404 .is_err()
5405 );
5406 server.abort();
5407 });
5408 }
5409
5410 #[test]
5411 fn direct_frames_match_anthropic_codec_bytes() {
5412 assert_eq!(
5413 DIRECT_PROXY_PING_FRAME,
5414 crate::codecs::anthropic::ANTHROPIC_PROXY_PING_FRAME
5415 );
5416 assert_eq!(
5417 DIRECT_PROXY_ERROR_FRAME,
5418 crate::codecs::anthropic::ANTHROPIC_PROXY_ERROR_FRAME
5419 );
5420 }
5421
5422 #[test]
5423 fn restore_text_json_escapes_raw_values_only_into_json_syntax() {
5424 let session = Session::new(Scope::Conversation("test".to_string())).unwrap();
5425 let raw = "q\"b\\s\u{1}\u{1f}\u{7f}\u{2028}\u{e9}\r\n";
5426 let mut transaction = session.begin_transaction();
5427 let token = transaction
5428 .tokenize(&PiiClass::Custom("secret".to_string()), raw)
5429 .unwrap();
5430 transaction.commit().unwrap();
5431
5432 let none = HashSet::new();
5433 let restored = restore_text(&session, &none, &format!("{{\"v\":\"x{token}y\"}}"), true);
5434 let parsed: Value = serde_json::from_str(&restored).unwrap();
5435 assert_eq!(parsed["v"].as_str(), Some(format!("x{raw}y").as_str()));
5436 assert_eq!(
5437 restore_text(&session, &none, &format!("x{token}y"), false),
5438 format!("x{raw}y")
5439 );
5440 let json_spelled = HashSet::from([token.clone()]);
5442 let restored = restore_text(
5443 &session,
5444 &json_spelled,
5445 &format!("{{\"v\":\"{token}\"}}"),
5446 true,
5447 );
5448 assert_eq!(
5449 serde_json::from_str::<Value>(&restored).unwrap()["v"].as_str(),
5450 Some(raw)
5451 );
5452 }
5453
5454 #[test]
5455 fn json_spelled_tokens_are_the_ones_minted_inside_json_string_literals() {
5456 let mut tokens = HashSet::new();
5457 collect_json_spelled_tokens(
5458 r#"{"a":"Robert \"Bob\" Smith","b":"<Email_9>"}"#,
5459 r#"{"a":"<Custom:name_1>","<Custom:key_1>":["<Custom:name_2>"],"b":"<Email_9>"}"#,
5460 &mut tokens,
5461 );
5462 assert_eq!(
5463 tokens,
5464 HashSet::from([
5465 "<Custom:name_1>".to_string(),
5466 "<Custom:key_1>".to_string(),
5467 "<Custom:name_2>".to_string(),
5468 ])
5469 );
5470
5471 let mut none = HashSet::new();
5472 collect_json_spelled_tokens("prose", "prose <Custom:name_1>", &mut none);
5473 collect_json_spelled_tokens("{broken", "{broken <Custom:name_1>", &mut none);
5474 assert!(none.is_empty());
5475 }
5476
5477 #[test]
5478 fn restore_text_leaves_unknown_tokens() {
5479 let session = Session::new(Scope::Conversation("test".to_string())).unwrap();
5480 let none = HashSet::new();
5481 assert_eq!(
5482 restore_text(&session, &none, "hello <Email_1>", false),
5483 "hello <Email_1>"
5484 );
5485 assert_eq!(
5486 restore_text(&session, &none, "{\"to\":\"<Email_1>\"}", true),
5487 "{\"to\":\"<Email_1>\"}"
5488 );
5489 }
5490
5491 #[test]
5492 fn upstream_url_preserves_path_and_query() {
5493 let base = Url::parse("https://api.openai.com").unwrap();
5494 let uri: Uri = "/v1/chat/completions?x=1".parse().unwrap();
5495 assert_eq!(
5496 upstream_url(&base, &uri).unwrap().as_str(),
5497 "https://api.openai.com/v1/chat/completions?x=1"
5498 );
5499 }
5500
5501 #[test]
5502 fn health_lists_adapters() {
5503 let config = ProxyConfig::new(
5504 "127.0.0.1:0".parse().unwrap(),
5505 vec![Arc::new(crate::adapters::OpenAiAdapter::new(
5506 Url::parse("https://api.openai.com").unwrap(),
5507 ))],
5508 );
5509 let state = AppState {
5510 config: Arc::new(config),
5511 pipeline: Arc::new(Pipeline::builder().build().unwrap()),
5512 client: Client::new(),
5513 sessions: Arc::new(RwLock::new(HashMap::new())),
5514 direct: None,
5515 started_at: Instant::now(),
5516 };
5517 assert_eq!(health_snapshot(&state).adapters[0].name, "openai");
5518 }
5519
5520 fn span(clean: std::ops::Range<usize>, raw: std::ops::Range<usize>) -> EmittedTokenSpan {
5521 EmittedTokenSpan::new(clean, raw, gaze::PiiClass::Email)
5522 }
5523
5524 #[test]
5525 fn manifest_accounting_accepts_an_untouched_pair_and_a_described_substitution() {
5526 assert!(manifest_accounts_for_every_change(
5527 "no pii here",
5528 "no pii here",
5529 &[]
5530 ));
5531 assert!(manifest_accounts_for_every_change(
5533 "lead a@b tail",
5534 "lead <T> tail",
5535 &[span(5..8, 5..8)],
5536 ));
5537 assert!(manifest_accounts_for_every_change(
5540 "a@b and c@d",
5541 "<TOKEN1> and <TOKEN2>",
5542 &[span(0..8, 0..3), span(13..21, 8..11)],
5543 ));
5544 }
5545
5546 #[test]
5547 fn manifest_accounting_refuses_a_deletion_no_entry_describes() {
5548 assert!(!manifest_accounts_for_every_change(
5551 "lead deleted@example.invalid tail",
5552 "lead tail",
5553 &[],
5554 ));
5555 assert!(!manifest_accounts_for_every_change(
5557 "a@b and deleted@example.invalid",
5558 "<T> and ",
5559 &[span(0..3, 0..3)],
5560 ));
5561 }
5562
5563 #[test]
5564 fn manifest_accounting_refuses_entries_that_cannot_describe_the_pair() {
5565 assert!(!manifest_accounts_for_every_change(
5567 "short",
5568 "short",
5569 &[span(0..5, 0..99)]
5570 ));
5571 assert!(!manifest_accounts_for_every_change(
5573 "a@b c@d",
5574 "<T><T>",
5575 &[span(0..3, 0..4), span(3..6, 2..7)],
5576 ));
5577 assert!(!manifest_accounts_for_every_change(
5579 "a@b tail",
5580 "<T> other",
5581 &[span(0..3, 0..3)]
5582 ));
5583 assert!(!manifest_accounts_for_every_change(
5585 "\u{00e4}x",
5586 "\u{00e4}x",
5587 &[span(0..1, 0..1)]
5588 ));
5589 }
5590
5591 #[test]
5592 fn manifest_accounting_refuses_a_deletion_that_falls_between_two_entries() {
5593 assert!(!manifest_accounts_for_every_change(
5606 "a@b GONE c@d",
5607 "<T1> <T2>",
5608 &[span(0..4, 0..3), span(5..9, 9..12)],
5609 ));
5610 }
5611}