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