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, CleanDocument, CommittedSessionSnapshot, DictionaryBundle, EmittedTokenSpan,
16 LocaleTag, Pipeline, PrefixCacheWriteMode, RawDocument, Scope, Session, SessionTransaction,
17 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}
1290
1291impl<'pipeline, 'context, 'transaction, 'session>
1292 PipelineRequestPseudonymizer<'pipeline, 'context, 'transaction, 'session>
1293{
1294 fn new(
1295 pipeline: &'pipeline Pipeline,
1296 locale_chain: &'context [LocaleTag],
1297 dictionaries: &'context DictionaryBundle,
1298 transaction: &'transaction mut SessionTransaction<'session>,
1299 ) -> Self {
1300 Self {
1301 pipeline,
1302 transaction,
1303 locale_chain,
1304 dictionaries,
1305 }
1306 }
1307
1308 fn protect_unstaged_segment(
1309 &mut self,
1310 input: &str,
1311 prefix_cache_write_mode: PrefixCacheWriteMode,
1312 ) -> Result<String, CodecErrorCode> {
1313 let clean = self
1314 .pipeline
1315 .pseudonymize_transaction_with_detect_context_and_prefix_cache_write_mode(
1316 self.transaction,
1317 RawDocument::Text(input.to_owned()),
1318 self.locale_chain,
1319 self.dictionaries,
1320 prefix_cache_write_mode,
1321 )
1322 .map_err(|_| CodecErrorCode::ProtectionFailedClosed)?;
1323 match clean {
1324 CleanDocument::Text(text) => Ok(text),
1325 _ => Err(CodecErrorCode::ProtectionFailedClosed),
1326 }
1327 }
1328}
1329
1330impl RequestPseudonymizer for PipelineRequestPseudonymizer<'_, '_, '_, '_> {
1331 fn protect(&mut self, input: &str) -> Result<String, CodecErrorCode> {
1332 self.protect_with_prefix_cache_write_mode(input, PrefixCacheWriteMode::Allow)
1333 }
1334
1335 fn protect_with_prefix_cache_write_mode(
1336 &mut self,
1337 input: &str,
1338 prefix_cache_write_mode: PrefixCacheWriteMode,
1339 ) -> Result<String, CodecErrorCode> {
1340 let output = self.protect_preserving_tokens(input, prefix_cache_write_mode)?;
1341 self.pipeline
1342 .admit_safety_nets_transaction(
1343 self.transaction,
1344 &output,
1345 self.locale_chain,
1346 self.dictionaries,
1347 )
1348 .map_err(|_| CodecErrorCode::ProtectionFailedClosed)?;
1349 Ok(output)
1350 }
1351
1352 fn validate_token_shapes(&mut self, input: &str) -> Result<(), CodecErrorCode> {
1353 self.transaction
1354 .validate_token_shapes(input)
1355 .map_err(|_| CodecErrorCode::RestoreFailedClosed)
1356 }
1357}
1358
1359impl PipelineRequestPseudonymizer<'_, '_, '_, '_> {
1360 fn protect_preserving_tokens(
1361 &mut self,
1362 input: &str,
1363 prefix_cache_write_mode: PrefixCacheWriteMode,
1364 ) -> Result<String, CodecErrorCode> {
1365 self.validate_token_shapes(input)?;
1366 let mut tokens = self.transaction.tokens();
1367 if tokens.is_empty() {
1368 return self.protect_unstaged_segment(input, prefix_cache_write_mode);
1369 }
1370 tokens.sort_unstable_by_key(|token| std::cmp::Reverse(token.len()));
1371
1372 let mut output = String::with_capacity(input.len());
1373 let mut cursor = 0;
1374 while cursor < input.len() {
1375 let next = tokens
1376 .iter()
1377 .filter_map(|token| {
1378 input[cursor..]
1379 .find(token)
1380 .map(|offset| (cursor + offset, token.as_str()))
1381 })
1382 .min_by(|left, right| {
1383 left.0
1384 .cmp(&right.0)
1385 .then_with(|| right.1.len().cmp(&left.1.len()))
1386 });
1387 let Some((start, token)) = next else {
1388 output.push_str(
1389 &self.protect_unstaged_segment(&input[cursor..], prefix_cache_write_mode)?,
1390 );
1391 break;
1392 };
1393 if start > cursor {
1394 output.push_str(
1395 &self
1396 .protect_unstaged_segment(&input[cursor..start], prefix_cache_write_mode)?,
1397 );
1398 }
1399 output.push_str(token);
1400 cursor = start + token.len();
1401 }
1402 Ok(output)
1403 }
1404}
1405
1406struct PipelineResponseResidualValidator<'pipeline, 'context> {
1407 pipeline: &'pipeline Pipeline,
1408 locale_chain: &'context [LocaleTag],
1409 dictionaries: &'context DictionaryBundle,
1410 validation_session: Session,
1411}
1412
1413impl<'pipeline, 'context> PipelineResponseResidualValidator<'pipeline, 'context> {
1414 fn new(
1415 pipeline: &'pipeline Pipeline,
1416 locale_chain: &'context [LocaleTag],
1417 dictionaries: &'context DictionaryBundle,
1418 ) -> Result<Self, DirectProxyError> {
1419 let validation_session = Session::new(Scope::Ephemeral).map_err(|_| {
1420 ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::ResponseValidation)
1421 })?;
1422 Ok(Self {
1423 pipeline,
1424 locale_chain,
1425 dictionaries,
1426 validation_session,
1427 })
1428 }
1429}
1430
1431impl ResponseResidualValidator for PipelineResponseResidualValidator<'_, '_> {
1432 fn validate(
1433 &mut self,
1434 value: &str,
1435 authorized_output_ranges: &[std::ops::Range<usize>],
1436 ) -> Result<(), CodecErrorCode> {
1437 let (clean, spans, _) = self
1438 .pipeline
1439 .clean_with_safety_net_detect_context(
1440 &self.validation_session,
1441 RawDocument::Text(value.to_owned()),
1442 self.locale_chain,
1443 self.dictionaries,
1444 )
1445 .map_err(|_| CodecErrorCode::ProviderOriginPii)?;
1446 let CleanDocument::Text(clean_text) = clean else {
1447 return Err(CodecErrorCode::ProviderOriginPii);
1448 };
1449 if !manifest_accounts_for_every_change(value, &clean_text, &spans) {
1457 return Err(CodecErrorCode::ProviderOriginPii);
1458 }
1459 for span in spans {
1460 if !authorized_output_ranges.iter().any(|authorized| {
1461 authorized.start <= span.raw_span.start && authorized.end >= span.raw_span.end
1462 }) {
1463 return Err(CodecErrorCode::ProviderOriginPii);
1464 }
1465 }
1466 Ok(())
1467 }
1468}
1469
1470fn manifest_accounts_for_every_change(raw: &str, clean: &str, spans: &[EmittedTokenSpan]) -> bool {
1479 let mut ordered: Vec<&EmittedTokenSpan> = spans.iter().collect();
1480 ordered.sort_by_key(|span| span.clean_span.start);
1481 let mut raw_cursor = 0usize;
1482 let mut clean_cursor = 0usize;
1483 for span in ordered {
1484 if span.raw_span.start < raw_cursor || span.clean_span.start < clean_cursor {
1485 return false;
1486 }
1487 let (Some(raw_gap), Some(clean_gap)) = (
1488 raw.get(raw_cursor..span.raw_span.start),
1489 clean.get(clean_cursor..span.clean_span.start),
1490 ) else {
1491 return false;
1492 };
1493 if raw_gap != clean_gap {
1494 return false;
1495 }
1496 if raw.get(span.raw_span.clone()).is_none() || clean.get(span.clean_span.clone()).is_none()
1499 {
1500 return false;
1501 }
1502 raw_cursor = span.raw_span.end;
1503 clean_cursor = span.clean_span.end;
1504 }
1505 matches!(
1506 (raw.get(raw_cursor..), clean.get(clean_cursor..)),
1507 (Some(raw_tail), Some(clean_tail)) if raw_tail == clean_tail
1508 )
1509}
1510
1511#[allow(clippy::too_many_arguments)]
1512fn prepare_and_commit_direct_request(
1513 pipeline: &Pipeline,
1514 locale_chain: &[LocaleTag],
1515 dictionaries: &DictionaryBundle,
1516 session: &Session,
1517 codec: &dyn BodyCodec,
1518 original_body: &[u8],
1519 endpoint: &Url,
1520 headers: &HeaderMap,
1521 expected_response: WireFormat,
1522 limits: CodecLimits,
1523 provider_request_projection_selected: bool,
1524 max_request_bytes: usize,
1525) -> Result<CommittedDirectRequest, DirectProxyError> {
1526 prepare_and_commit_direct_request_with_hook(
1527 pipeline,
1528 locale_chain,
1529 dictionaries,
1530 session,
1531 codec,
1532 original_body,
1533 endpoint,
1534 headers,
1535 expected_response,
1536 limits,
1537 provider_request_projection_selected,
1538 max_request_bytes,
1539 |_, _, _| Ok(()),
1540 )
1541}
1542
1543#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1544enum SessionCommitFailureKind {
1545 GenerationConflict,
1546 NonConflict,
1547}
1548
1549fn classify_session_commit_failure(error: &SessionTransactionError) -> SessionCommitFailureKind {
1550 match error {
1551 SessionTransactionError::GenerationConflict => SessionCommitFailureKind::GenerationConflict,
1552 _ => SessionCommitFailureKind::NonConflict,
1553 }
1554}
1555
1556fn map_session_commit_failure(kind: SessionCommitFailureKind) -> DirectProxyError {
1557 match kind {
1558 SessionCommitFailureKind::GenerationConflict => {
1559 ProxyErrorCode::SessionGenerationConflict.error(ProxyErrorPhase::Session)
1560 }
1561 SessionCommitFailureKind::NonConflict => {
1562 ProxyErrorCode::SessionCommitFailure.error(ProxyErrorPhase::Session)
1563 }
1564 }
1565}
1566
1567#[allow(clippy::too_many_arguments)]
1568fn prepare_and_commit_direct_request_with_hook(
1569 pipeline: &Pipeline,
1570 locale_chain: &[LocaleTag],
1571 dictionaries: &DictionaryBundle,
1572 session: &Session,
1573 codec: &dyn BodyCodec,
1574 original_body: &[u8],
1575 endpoint: &Url,
1576 headers: &HeaderMap,
1577 expected_response: WireFormat,
1578 limits: CodecLimits,
1579 provider_request_projection_selected: bool,
1580 max_request_bytes: usize,
1581 mut before_commit: impl FnMut(u8, &Session, &SessionTransaction<'_>) -> Result<(), DirectProxyError>,
1582) -> Result<CommittedDirectRequest, DirectProxyError> {
1583 for attempt in 0..=1 {
1584 let mut transaction = session.begin_transaction();
1585 let proved_body = {
1586 let mut pseudonymizer = PipelineRequestPseudonymizer::new(
1587 pipeline,
1588 locale_chain,
1589 dictionaries,
1590 &mut transaction,
1591 );
1592 let mut context =
1593 RequestTransformContext::new(&mut pseudonymizer, WireFormat::Json, limits);
1594 if provider_request_projection_selected {
1595 context.request_inspection_projection();
1596 }
1597 codec
1598 .protect_request(original_body, &mut context)
1599 .map_err(map_codec_error)?
1600 };
1601 let request = DirectRequest::from_proved(
1602 endpoint.clone(),
1603 headers.clone(),
1604 proved_body,
1605 expected_response,
1606 )?;
1607 if request.body.len() > max_request_bytes {
1608 return Err(
1609 ProxyErrorCode::RequestBodyLimitExceeded.error(ProxyErrorPhase::RequestValidation)
1610 );
1611 }
1612 let prepared = PreparedDirectRequest::new(transaction, request);
1613 before_commit(attempt, session, &prepared.transaction)?;
1614 match prepared.commit() {
1615 Ok(committed) => return Ok(committed),
1616 Err(error) => match classify_session_commit_failure(&error) {
1617 SessionCommitFailureKind::GenerationConflict if attempt == 0 => {}
1618 kind => return Err(map_session_commit_failure(kind)),
1619 },
1620 }
1621 }
1622 Err(ProxyErrorCode::SessionGenerationConflict.error(ProxyErrorPhase::Session))
1623}
1624
1625#[derive(Debug, Eq, PartialEq)]
1627#[allow(dead_code)]
1628pub(crate) enum DirectIoState {
1629 Prepared,
1630 CommittedNoIo,
1631 IoAttempted,
1632 ResponseComplete,
1633}
1634
1635#[derive(Debug, Eq, PartialEq)]
1637#[allow(dead_code)]
1638pub(crate) struct DirectTransactionState {
1639 state: DirectIoState,
1640 generation_retries: u8,
1641}
1642
1643#[allow(dead_code)]
1644impl DirectTransactionState {
1645 #[must_use]
1646 pub(crate) const fn new() -> Self {
1647 Self {
1648 state: DirectIoState::Prepared,
1649 generation_retries: 0,
1650 }
1651 }
1652
1653 #[must_use]
1654 pub(crate) const fn state(&self) -> &DirectIoState {
1655 &self.state
1656 }
1657
1658 #[must_use]
1659 pub(crate) const fn generation_retries(&self) -> u8 {
1660 self.generation_retries
1661 }
1662
1663 #[must_use]
1664 pub(crate) const fn retains_committed_mapping(&self) -> bool {
1665 !matches!(self.state, DirectIoState::Prepared)
1666 }
1667
1668 pub(crate) fn retry_generation_conflict(&mut self) -> Result<(), DirectProxyError> {
1669 if self.state != DirectIoState::Prepared || self.generation_retries != 0 {
1670 return Err(ProxyErrorCode::SessionGenerationConflict.error(ProxyErrorPhase::Session));
1671 }
1672 self.generation_retries = 1;
1673 Ok(())
1674 }
1675
1676 pub(crate) fn mark_committed(&mut self) -> Result<(), DirectProxyError> {
1677 self.transition(DirectIoState::Prepared, DirectIoState::CommittedNoIo)
1678 }
1679
1680 pub(crate) fn mark_io_attempted(&mut self) -> Result<(), DirectProxyError> {
1681 self.transition(DirectIoState::CommittedNoIo, DirectIoState::IoAttempted)
1682 }
1683
1684 pub(crate) fn mark_response_complete(&mut self) -> Result<(), DirectProxyError> {
1685 self.transition(DirectIoState::IoAttempted, DirectIoState::ResponseComplete)
1686 }
1687
1688 fn transition(
1689 &mut self,
1690 expected: DirectIoState,
1691 next: DirectIoState,
1692 ) -> Result<(), DirectProxyError> {
1693 if self.state != expected {
1694 return Err(ProxyErrorCode::InvalidStateTransition.error(ProxyErrorPhase::Framework));
1695 }
1696 self.state = next;
1697 Ok(())
1698 }
1699
1700 fn require_committed_no_io(&self) -> Result<(), DirectProxyError> {
1701 if self.state == DirectIoState::CommittedNoIo {
1702 Ok(())
1703 } else {
1704 Err(ProxyErrorCode::InvalidStateTransition.error(ProxyErrorPhase::Framework))
1705 }
1706 }
1707}
1708
1709impl Default for DirectTransactionState {
1710 fn default() -> Self {
1711 Self::new()
1712 }
1713}
1714
1715#[derive(Clone)]
1716struct AppState {
1717 config: Arc<ProxyConfig>,
1718 pipeline: Arc<Pipeline>,
1719 client: Client,
1720 direct: Option<Arc<DirectRuntime>>,
1721 sessions: Arc<RwLock<HashMap<String, SessionEntry>>>,
1722 started_at: Instant,
1723}
1724
1725struct SessionEntry {
1726 session: Arc<LegacySession>,
1727 expires_at: Instant,
1728}
1729
1730struct LegacySession {
1733 session: Session,
1734 json_spelled_tokens: std::sync::Mutex<HashSet<String>>,
1738}
1739
1740impl LegacySession {
1741 fn new(session: Session) -> Self {
1742 Self {
1743 session,
1744 json_spelled_tokens: std::sync::Mutex::new(HashSet::new()),
1745 }
1746 }
1747
1748 fn json_spelled_tokens(&self) -> std::sync::MutexGuard<'_, HashSet<String>> {
1750 self.json_spelled_tokens
1751 .lock()
1752 .unwrap_or_else(std::sync::PoisonError::into_inner)
1753 }
1754}
1755
1756#[cfg(test)]
1757static SESSION_FOR_MISS_BARRIERS: std::sync::LazyLock<
1758 std::sync::Mutex<HashMap<String, Arc<tokio::sync::Barrier>>>,
1759> = std::sync::LazyLock::new(|| std::sync::Mutex::new(HashMap::new()));
1760
1761#[derive(Debug, Clone, Serialize)]
1762#[non_exhaustive]
1763pub struct HealthSnapshot {
1764 pub uptime_secs: u64,
1765 pub bind: String,
1766 pub adapters: Vec<AdapterSnapshot>,
1767}
1768
1769#[derive(Debug, Clone, Serialize)]
1770#[non_exhaustive]
1771pub struct AdapterSnapshot {
1772 pub name: &'static str,
1773 pub upstream: String,
1774}
1775
1776pub async fn serve(config: ProxyConfig, pipeline: Arc<Pipeline>) -> Result<(), ProxyError> {
1778 serve_with_ready(config, pipeline, || {}).await
1779}
1780
1781pub async fn serve_with_ready(
1783 config: ProxyConfig,
1784 pipeline: Arc<Pipeline>,
1785 ready: impl FnOnce(),
1786) -> Result<(), ProxyError> {
1787 let bind = config.bind;
1788 let app = build_app(config, pipeline)?;
1789 let listener = TcpListener::bind(bind)
1790 .await
1791 .map_err(|source| ProxyError::Server { source })?;
1792 ready();
1793 serve_app(listener, app).await
1794}
1795
1796pub async fn serve_with_listener(
1804 mut config: ProxyConfig,
1805 pipeline: Arc<Pipeline>,
1806 listener: TcpListener,
1807) -> Result<(), ProxyError> {
1808 let actual = listener
1809 .local_addr()
1810 .map_err(|source| ProxyError::Server { source })?;
1811 if config.bind.port() == 0 {
1812 config.bind.set_port(actual.port());
1813 }
1814 if config.bind != actual {
1815 return Err(ProxyError::Server {
1816 source: std::io::Error::new(
1817 std::io::ErrorKind::InvalidInput,
1818 "listener address does not match proxy configuration",
1819 ),
1820 });
1821 }
1822 let app = build_app(config, pipeline)?;
1823 serve_app(listener, app).await
1824}
1825
1826fn build_app(config: ProxyConfig, pipeline: Arc<Pipeline>) -> Result<Router, ProxyError> {
1827 let direct = DirectRuntime::from_config(&config)?.map(Arc::new);
1828 let state = AppState {
1829 config: Arc::new(config),
1830 pipeline,
1831 client: Client::new(),
1832 direct,
1833 sessions: Arc::new(RwLock::new(HashMap::new())),
1834 started_at: Instant::now(),
1835 };
1836 Ok(Router::new()
1837 .route("/_gaze_proxy/healthz", get(healthz))
1838 .fallback(proxy)
1839 .with_state(state))
1840}
1841
1842async fn serve_app(listener: TcpListener, app: Router) -> Result<(), ProxyError> {
1843 axum::serve(
1844 listener,
1845 app.into_make_service_with_connect_info::<SocketAddr>(),
1846 )
1847 .await
1848 .map_err(|source| ProxyError::Server { source })
1849}
1850
1851async fn healthz(State(state): State<AppState>) -> impl IntoResponse {
1852 axum::Json(health_snapshot(&state))
1853}
1854
1855fn health_snapshot(state: &AppState) -> HealthSnapshot {
1856 HealthSnapshot {
1857 uptime_secs: state.started_at.elapsed().as_secs(),
1858 bind: state.config.bind.to_string(),
1859 adapters: state
1860 .config
1861 .adapters
1862 .iter()
1863 .map(|adapter| AdapterSnapshot {
1864 name: adapter.name(),
1865 upstream: adapter.upstream_base().to_string(),
1866 })
1867 .collect(),
1868 }
1869}
1870
1871async fn proxy(
1872 State(state): State<AppState>,
1873 ConnectInfo(peer): ConnectInfo<SocketAddr>,
1874 request: Request<Body>,
1875) -> Response {
1876 let (parts, body) = request.into_parts();
1877 let direct_route_candidate =
1878 state.direct.is_some() && matches!(parts.uri.path(), "/v1/messages" | "/v1/messages/");
1879 let direct_ingress = if direct_route_candidate {
1880 match prepare_direct_ingress(&state, peer, &parts.method, &parts.uri, &parts.headers) {
1881 Ok(ingress) => Some(ingress),
1882 Err(error) => return direct_error_response(error),
1883 }
1884 } else {
1885 None
1886 };
1887 let is_direct = direct_ingress.is_some();
1888 let body_limit_bytes = if is_direct {
1889 state
1890 .direct
1891 .as_ref()
1892 .map_or(state.config.body_limit_bytes, |runtime| {
1893 state
1894 .config
1895 .body_limit_bytes
1896 .min(runtime.client_config.max_request_bytes as u64)
1897 })
1898 } else {
1899 state.config.body_limit_bytes
1900 };
1901 let body = match collect_inbound_body(body, &parts.headers, body_limit_bytes).await {
1902 Ok(body) => body,
1903 Err(_) if is_direct => {
1904 return direct_error_response(
1905 ProxyErrorCode::RequestBodyLimitExceeded.error(ProxyErrorPhase::RequestValidation),
1906 );
1907 }
1908 Err(error) => return proxy_error_response(error),
1909 };
1910 match proxy_inner(
1911 state,
1912 parts.method,
1913 parts.uri,
1914 parts.headers,
1915 body,
1916 direct_ingress,
1917 )
1918 .await
1919 {
1920 Ok(response) => response,
1921 Err(err) => proxy_error_response(err),
1922 }
1923}
1924
1925async fn collect_inbound_body(
1926 body: Body,
1927 headers: &HeaderMap,
1928 limit_bytes: u64,
1929) -> Result<Bytes, ProxyError> {
1930 if let Some(length) = headers
1931 .get(axum::http::header::CONTENT_LENGTH)
1932 .and_then(|value| std::str::from_utf8(value.as_bytes()).ok())
1933 .and_then(|value| value.parse::<u64>().ok())
1934 {
1935 if length > limit_bytes {
1936 return Err(ProxyError::BodyTooLarge { limit_bytes });
1937 }
1938 }
1939
1940 let mut stream = body.into_data_stream();
1941 let mut collected = Vec::new();
1942 let mut accumulated = 0_u64;
1943 while let Some(chunk) = stream.next().await {
1944 let chunk = chunk.map_err(|_| ProxyError::DaemonConfig {
1945 detail: "request_body_rejected".to_string(),
1946 })?;
1947 accumulated = accumulated
1948 .checked_add(chunk.len() as u64)
1949 .ok_or(ProxyError::BodyTooLarge { limit_bytes })?;
1950 if accumulated > limit_bytes {
1951 return Err(ProxyError::BodyTooLarge { limit_bytes });
1952 }
1953 collected.extend_from_slice(&chunk);
1954 }
1955 Ok(Bytes::from(collected))
1956}
1957
1958async fn proxy_inner(
1959 state: AppState,
1960 method: Method,
1961 uri: Uri,
1962 headers: HeaderMap,
1963 body: Bytes,
1964 direct_ingress: Option<PreparedDirectIngress>,
1965) -> Result<Response, ProxyError> {
1966 if body.len() as u64 > state.config.body_limit_bytes {
1967 return Err(ProxyError::BodyTooLarge {
1968 limit_bytes: state.config.body_limit_bytes,
1969 });
1970 }
1971 let path = uri.path();
1972 let adapter = state
1973 .config
1974 .adapters
1975 .iter()
1976 .find(|adapter| adapter.matches_path(&method, path))
1977 .cloned()
1978 .ok_or_else(|| ProxyError::AdapterNotFound {
1979 path: path.to_string(),
1980 method: method.clone(),
1981 })?;
1982 let contract = adapter.contract();
1983 if let ProtocolContract::Codec(codec) = contract.protocol() {
1984 return Ok(
1985 match direct_proxy_inner(
1986 &state,
1987 direct_ingress.ok_or_else(|| ProxyError::DaemonConfig {
1988 detail: "direct_ingress_missing".to_string(),
1989 })?,
1990 &body,
1991 codec,
1992 )
1993 .await
1994 {
1995 Ok(response) => response,
1996 Err(error) => direct_error_response(error),
1997 },
1998 );
1999 }
2000 let session = session_for(&state, &headers).await?;
2001 let mut json: Value =
2002 serde_json::from_slice(&body).map_err(|source| ProxyError::InvalidJson { source })?;
2003 let json_output = adapter.requests_json_output(&json);
2006 let locale_chain = state.config.locale_chain();
2010 let dictionaries = state.config.dictionaries();
2011 let redacted = redact_surfaces(
2012 &state.pipeline,
2013 &session.session,
2014 adapter.request_pii_surfaces(&mut json),
2015 locale_chain,
2016 dictionaries,
2017 )?;
2018 session
2019 .json_spelled_tokens()
2020 .extend(redacted.json_spelled_tokens);
2021 match contract.coverage_policy() {
2025 CoveragePolicy::LegacySurfaces => {
2028 residual_scan_request(
2029 &state.pipeline,
2030 &json,
2031 &redacted.field_paths,
2032 locale_chain,
2033 dictionaries,
2034 )?;
2035 }
2036 CoveragePolicy::CodecProved => return Err(ProxyError::UnprovenCoverage),
2039 }
2040
2041 let upstream_url = upstream_url(adapter.upstream_base(), &uri)?;
2042 let mut request = state
2043 .client
2044 .request(method.clone(), upstream_url.clone())
2045 .json(&json);
2046 request = forward_headers(request, &headers);
2047 let upstream = request
2048 .send()
2049 .await
2050 .map_err(|source| ProxyError::UpstreamUnreachable {
2051 url: upstream_url.clone(),
2052 source,
2053 })?;
2054
2055 let status = upstream.status();
2056 let upstream_headers = upstream.headers().clone();
2057 let content_type = upstream_headers
2058 .get(reqwest::header::CONTENT_TYPE)
2059 .and_then(|value| value.to_str().ok())
2060 .unwrap_or_default()
2061 .to_ascii_lowercase();
2062 let bytes = upstream
2063 .bytes()
2064 .await
2065 .map_err(|source| ProxyError::UpstreamUnreachable {
2066 url: upstream_url,
2067 source,
2068 })?;
2069
2070 let is_sse = content_type.contains("text/event-stream");
2071 let body = if is_sse {
2072 transform_sse(&adapter, &session, &bytes, json_output)?
2073 } else {
2074 let mut response_json: Value =
2075 serde_json::from_slice(&bytes).map_err(|source| ProxyError::InvalidJson { source })?;
2076 restore_surfaces(
2077 &session,
2078 adapter.response_pii_surfaces(&mut response_json),
2079 json_output,
2080 );
2081 serde_json::to_vec(&response_json).map_err(|source| ProxyError::InvalidJson { source })?
2082 };
2083
2084 let mut response = Response::builder().status(status);
2085 response = response.header(
2086 axum::http::header::CONTENT_TYPE,
2087 if is_sse {
2088 "text/event-stream"
2089 } else {
2090 "application/json"
2091 },
2092 );
2093 response
2094 .body(axum::body::Body::from(body))
2095 .map_err(|source| ProxyError::DaemonConfig {
2096 detail: source.to_string(),
2097 })
2098}
2099
2100fn prepare_direct_ingress(
2101 state: &AppState,
2102 peer: SocketAddr,
2103 method: &Method,
2104 uri: &Uri,
2105 headers: &HeaderMap,
2106) -> Result<PreparedDirectIngress, DirectProxyError> {
2107 let runtime = state
2108 .direct
2109 .as_ref()
2110 .ok_or_else(|| ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::Framework))?;
2111 if method != Method::POST || uri.path() != "/v1/messages" || uri.query().is_some() {
2112 return Err(ProxyErrorCode::RouteRejected.error(ProxyErrorPhase::Ingress));
2113 }
2114 let outbound_headers = validate_direct_request_headers(headers, runtime)?;
2115 let session = match runtime.session_policy {
2116 SessionPolicy::EphemeralSingleRequest => {
2117 if headers.get_all(SESSION_HEADER).iter().next().is_some() {
2118 return Err(ProxyErrorCode::UnexpectedSessionHeader
2119 .error(ProxyErrorPhase::RequestValidation));
2120 }
2121 DirectIngressSession::Ephemeral
2122 }
2123 SessionPolicy::OpaqueHeaderContinuity(_) => {
2124 let identifier = canonical_session_identifier(headers)?;
2125 let peer_scope = if peer.ip().is_loopback() {
2126 PeerScope::Loopback
2127 } else {
2128 PeerScope::NonLoopback
2129 };
2130 let context = PrincipalContext::new(runtime.listener_scope, peer_scope, None, None)
2131 .map_err(|_| ProxyErrorCode::PrincipalRejected.error(ProxyErrorPhase::Session))?;
2132 let principal = invoke_resolver(runtime.resolver.as_ref(), context)
2133 .map_err(|_| ProxyErrorCode::PrincipalRejected.error(ProxyErrorPhase::Session))?;
2134 let registry = runtime.registry.as_ref().ok_or_else(|| {
2135 ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::Session)
2136 })?;
2137 let lease = registry
2138 .acquire(principal, identifier, None)
2139 .map_err(map_registry_error)?;
2140 DirectIngressSession::Continuity(lease)
2141 }
2142 };
2143 Ok(PreparedDirectIngress {
2144 outbound_headers,
2145 session,
2146 })
2147}
2148
2149async fn direct_proxy_inner(
2150 state: &AppState,
2151 ingress: PreparedDirectIngress,
2152 body: &[u8],
2153 codec: &dyn BodyCodec,
2154) -> Result<Response, DirectProxyError> {
2155 let runtime = state
2156 .direct
2157 .as_ref()
2158 .ok_or_else(|| ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::Framework))?;
2159 if body.len() > runtime.client_config.max_request_bytes {
2160 return Err(
2161 ProxyErrorCode::RequestBodyLimitExceeded.error(ProxyErrorPhase::RequestValidation)
2162 );
2163 }
2164 let expected_response = parse_request_stream_format(body)?;
2165 let mut inspection = state
2166 .config
2167 .inspection
2168 .as_ref()
2169 .and_then(|producer| producer.begin_logical(runtime.inspection_endpoint_codes));
2170
2171 let mut continuity_lease = None;
2172 let mut committed = match ingress.session {
2173 DirectIngressSession::Ephemeral => {
2174 let session = Session::new(Scope::Ephemeral)
2175 .map_err(|_| ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::Session))?;
2176 prepare_and_commit_direct_request(
2177 &state.pipeline,
2178 state.config.locale_chain(),
2179 state.config.dictionaries(),
2180 &session,
2181 codec,
2182 body,
2183 &runtime.endpoint,
2184 &ingress.outbound_headers,
2185 expected_response,
2186 runtime.codec_limits,
2187 inspection
2188 .as_ref()
2189 .is_some_and(ProxyInspectionLogicalV1::provider_request_projection_selected),
2190 runtime.client_config.max_request_bytes,
2191 )?
2192 }
2193 DirectIngressSession::Continuity(lease) => {
2194 let request_guard = lease.lock_request().await;
2195 let committed = prepare_and_commit_direct_request(
2196 &state.pipeline,
2197 state.config.locale_chain(),
2198 state.config.dictionaries(),
2199 lease.session(),
2200 codec,
2201 body,
2202 &runtime.endpoint,
2203 &ingress.outbound_headers,
2204 expected_response,
2205 runtime.codec_limits,
2206 inspection
2207 .as_ref()
2208 .is_some_and(ProxyInspectionLogicalV1::provider_request_projection_selected),
2209 runtime.client_config.max_request_bytes,
2210 );
2211 drop(request_guard);
2212 let committed = committed?;
2213 continuity_lease = Some(lease);
2214 committed
2215 }
2216 };
2217
2218 if let Some(inspection) = inspection.as_mut() {
2219 let inspection_parsed_facts = committed.request.inspection_parsed_facts.take();
2220 inspection.emit_request_stages(
2221 body,
2222 &committed.request.body,
2223 inspection_parsed_facts,
2224 committed.request.signed_or_encrypted_surface,
2225 );
2226 }
2227
2228 if expected_response == WireFormat::Sse {
2229 let opened = runtime.client.execute_sse(committed).await?;
2230 return stage_and_replay_direct_sse(
2231 Arc::clone(&state.pipeline),
2232 state.config.locale_chain.clone(),
2233 state.config.dictionaries().clone(),
2234 opened,
2235 runtime.codec_limits,
2236 runtime.client_config.max_response_bytes,
2237 runtime.ping_interval,
2238 inspection,
2239 continuity_lease,
2240 );
2241 }
2242 let executed = runtime.client.execute(committed).await?;
2243 let response = restore_direct_response(
2244 &state.pipeline,
2245 state.config.locale_chain(),
2246 state.config.dictionaries(),
2247 codec,
2248 executed,
2249 runtime.codec_limits,
2250 inspection.as_mut(),
2251 );
2252 drop(continuity_lease);
2253 response
2254}
2255
2256#[allow(clippy::too_many_arguments)]
2257fn stage_and_replay_direct_sse(
2258 pipeline: Arc<Pipeline>,
2259 locale_chain: gaze::LocaleChain,
2260 dictionaries: DictionaryBundle,
2261 opened: CommittedDirectSse,
2262 limits: CodecLimits,
2263 max_response_bytes: usize,
2264 ping_interval: Duration,
2265 inspection: Option<ProxyInspectionLogicalV1>,
2266 continuity_lease: Option<SessionLease>,
2267) -> Result<Response, DirectProxyError> {
2268 let CommittedDirectSse {
2269 snapshot,
2270 response: upstream_response,
2271 head,
2272 lifecycle,
2273 remaining_total_timeout,
2274 } = opened;
2275 let status = head.status();
2276 let headers = head.headers().clone();
2277 let (sender, receiver) = tokio::sync::mpsc::channel::<Result<Bytes, std::io::Error>>(8);
2278 let mut response = Response::builder().status(status);
2279 for (name, value) in &headers {
2280 response = response.header(name, value);
2281 }
2282 let response = response
2283 .body(Body::from_stream(
2284 tokio_stream::wrappers::ReceiverStream::new(receiver),
2285 ))
2286 .map_err(|_| ProxyErrorCode::UpstreamProtocol.error(ProxyErrorPhase::Framework))?;
2287 tokio::spawn(async move {
2288 let _continuity_lease = continuity_lease;
2289 let mut inspection = inspection;
2290 let proof = async move {
2291 let body = collect_response_body(upstream_response, max_response_bytes).await?;
2292 validate_declared_body(WireFormat::Sse, &body, max_response_bytes)?;
2293 let mut residual = PipelineResponseResidualValidator::new(
2294 &pipeline,
2295 locale_chain.as_slice(),
2296 &dictionaries,
2297 )?;
2298 let mut context =
2299 ResponseTransformContext::new(&snapshot, &mut residual, WireFormat::Sse, limits);
2300 if inspection
2301 .as_ref()
2302 .is_some_and(ProxyInspectionLogicalV1::owner_restored_response_projection_selected)
2303 {
2304 context.request_inspection_projection();
2305 }
2306 let mut proved = crate::codecs::anthropic::AnthropicMessagesCodec
2307 .restore_response(&body, &mut context)
2308 .map_err(map_codec_error)?;
2309 if let Some(inspection) = inspection.as_mut() {
2310 let format = proved.format();
2311 let signed_or_encrypted_surface = proved.signed_or_encrypted_surface();
2312 let inspection_parsed_facts = proved.take_inspection_parsed_facts();
2313 inspection.emit_response_stages(
2314 &body,
2315 proved.bytes(),
2316 format,
2317 inspection_parsed_facts,
2318 signed_or_encrypted_surface,
2319 );
2320 }
2321 let mut lifecycle = lifecycle;
2322 lifecycle.mark_response_complete()?;
2323 Ok::<Vec<u8>, DirectProxyError>(proved.into_bytes())
2324 };
2325 let timed_proof = tokio::time::timeout(remaining_total_timeout, proof);
2326 tokio::pin!(timed_proof);
2327 let mut ping = tokio::time::interval(ping_interval);
2328 ping.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
2329 ping.tick().await;
2330 loop {
2331 tokio::select! {
2332 () = sender.closed() => {
2333 break;
2334 }
2335 result = &mut timed_proof => {
2336 let output = match result {
2337 Ok(Ok(bytes)) => Bytes::from(bytes),
2338 Ok(Err(_)) | Err(_) => Bytes::from_static(DIRECT_PROXY_ERROR_FRAME),
2339 };
2340 let _ = sender.send(Ok(output)).await;
2341 break;
2342 }
2343 _ = ping.tick() => {
2344 if sender
2345 .send(Ok(Bytes::from_static(DIRECT_PROXY_PING_FRAME)))
2346 .await
2347 .is_err()
2348 {
2349 break;
2350 }
2351 }
2352 }
2353 }
2354 });
2355 Ok(response)
2356}
2357
2358fn validate_direct_request_headers(
2359 headers: &HeaderMap,
2360 runtime: &DirectRuntime,
2361) -> Result<HeaderMap, DirectProxyError> {
2362 validate_header_space(headers, runtime.client_config)?;
2363 for singleton in [
2364 "content-type",
2365 "content-length",
2366 "transfer-encoding",
2367 "content-encoding",
2368 "connection",
2369 "x-api-key",
2370 "anthropic-version",
2371 "anthropic-beta",
2372 SESSION_HEADER,
2373 ] {
2374 reject_duplicate_request_header(headers, singleton)?;
2375 }
2376 if headers.get(axum::http::header::CONTENT_ENCODING).is_some() {
2377 return Err(
2378 ProxyErrorCode::UnsupportedContentEncoding.error(ProxyErrorPhase::RequestValidation)
2379 );
2380 }
2381 if let Some(connection) = one_request_header(headers, "connection", false)? {
2382 let connection = connection.to_str().map_err(|_| {
2383 ProxyErrorCode::InvalidFraming.error(ProxyErrorPhase::RequestValidation)
2384 })?;
2385 if !connection.eq_ignore_ascii_case("keep-alive")
2386 && !connection.eq_ignore_ascii_case("close")
2387 {
2388 return Err(ProxyErrorCode::InvalidFraming.error(ProxyErrorPhase::RequestValidation));
2389 }
2390 }
2391 if headers.get(axum::http::header::TRANSFER_ENCODING).is_some()
2392 && headers.get(axum::http::header::CONTENT_LENGTH).is_some()
2393 {
2394 return Err(ProxyErrorCode::InvalidFraming.error(ProxyErrorPhase::RequestValidation));
2395 }
2396 let content_type = one_request_header(headers, "content-type", true)?
2397 .ok_or_else(|| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2398 if parse_wire_content_type(content_type)? != WireFormat::Json {
2399 return Err(ProxyErrorCode::InvalidRequestFormat.error(ProxyErrorPhase::RequestValidation));
2400 }
2401 let api_key = one_request_header(headers, "x-api-key", true)?
2402 .ok_or_else(|| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2403 if !api_key
2404 .as_bytes()
2405 .iter()
2406 .all(|byte| (0x21..=0x7e).contains(byte) && *byte != b',')
2407 {
2408 return Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation));
2409 }
2410 let version = one_request_header(headers, "anthropic-version", true)?
2411 .ok_or_else(|| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2412 let version_text = std::str::from_utf8(version.as_bytes())
2413 .map_err(|_| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2414 if !runtime
2415 .allowed_versions
2416 .iter()
2417 .any(|allowed| allowed == version_text)
2418 {
2419 return Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation));
2420 }
2421 let beta = one_request_header(headers, "anthropic-beta", false)?;
2422 if let Some(beta) = beta {
2423 let beta_text = std::str::from_utf8(beta.as_bytes()).map_err(|_| {
2424 ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation)
2425 })?;
2426 if !runtime
2427 .allowed_betas
2428 .iter()
2429 .any(|allowed| allowed == beta_text)
2430 {
2431 return Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation));
2432 }
2433 }
2434
2435 let mut outbound = HeaderMap::new();
2436 outbound.insert(
2437 axum::http::header::CONTENT_TYPE,
2438 HeaderValue::from_static("application/json"),
2439 );
2440 outbound.insert(HeaderName::from_static("x-api-key"), api_key.clone());
2441 outbound.insert(
2442 HeaderName::from_static("anthropic-version"),
2443 version.clone(),
2444 );
2445 if let Some(beta) = beta {
2446 outbound.insert(HeaderName::from_static("anthropic-beta"), beta.clone());
2447 }
2448 Ok(outbound)
2449}
2450
2451fn one_request_header<'a>(
2452 headers: &'a HeaderMap,
2453 name: &'static str,
2454 required: bool,
2455) -> Result<Option<&'a HeaderValue>, DirectProxyError> {
2456 let mut values = headers.get_all(name).iter();
2457 let first = values.next();
2458 if values.next().is_some() || first.is_some_and(|value| value.is_empty()) {
2459 return Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation));
2460 }
2461 match (first, required) {
2462 (Some(value), _) => Ok(Some(value)),
2463 (None, false) => Ok(None),
2464 (None, true) => {
2465 Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))
2466 }
2467 }
2468}
2469
2470fn reject_duplicate_request_header(
2471 headers: &HeaderMap,
2472 name: &'static str,
2473) -> Result<(), DirectProxyError> {
2474 let mut values = headers.get_all(name).iter();
2475 let _ = values.next();
2476 if values.next().is_some() {
2477 Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))
2478 } else {
2479 Ok(())
2480 }
2481}
2482
2483fn canonical_session_identifier(headers: &HeaderMap) -> Result<&str, DirectProxyError> {
2484 let value = one_request_header(headers, SESSION_HEADER, false)?
2485 .ok_or_else(|| {
2486 ProxyErrorCode::SessionIdentityRequired.error(ProxyErrorPhase::RequestValidation)
2487 })?
2488 .to_str()
2489 .map_err(|_| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2490 let identifier = uuid::Uuid::parse_str(value)
2491 .map_err(|_| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2492 if identifier.get_version() != Some(uuid::Version::Random)
2493 || identifier.get_variant() != uuid::Variant::RFC4122
2494 || identifier.hyphenated().to_string() != value
2495 {
2496 return Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation));
2497 }
2498 Ok(value)
2499}
2500
2501fn restore_direct_response(
2502 pipeline: &Pipeline,
2503 locale_chain: &[LocaleTag],
2504 dictionaries: &DictionaryBundle,
2505 codec: &dyn BodyCodec,
2506 mut executed: CommittedDirectResponse,
2507 limits: CodecLimits,
2508 inspection: Option<&mut ProxyInspectionLogicalV1>,
2509) -> Result<Response, DirectProxyError> {
2510 let (head, body) = executed.response.into_parts();
2511 let mut residual =
2512 PipelineResponseResidualValidator::new(pipeline, locale_chain, dictionaries)?;
2513 let mut context =
2514 ResponseTransformContext::new(&executed.snapshot, &mut residual, head.format(), limits);
2515 if inspection
2516 .as_deref()
2517 .is_some_and(ProxyInspectionLogicalV1::owner_restored_response_projection_selected)
2518 {
2519 context.request_inspection_projection();
2520 }
2521 let mut proved = codec
2522 .restore_response(&body, &mut context)
2523 .map_err(map_codec_error)?;
2524 if let Some(inspection) = inspection {
2525 let format = proved.format();
2526 let signed_or_encrypted_surface = proved.signed_or_encrypted_surface();
2527 let inspection_parsed_facts = proved.take_inspection_parsed_facts();
2528 inspection.emit_response_stages(
2529 &body,
2530 proved.bytes(),
2531 format,
2532 inspection_parsed_facts,
2533 signed_or_encrypted_surface,
2534 );
2535 }
2536 executed.lifecycle.mark_response_complete()?;
2537 let mut response = Response::builder().status(head.status());
2538 for (name, value) in head.headers() {
2539 response = response.header(name, value);
2540 }
2541 response
2542 .body(Body::from(proved.into_bytes()))
2543 .map_err(|_| ProxyErrorCode::UpstreamProtocol.error(ProxyErrorPhase::Framework))
2544}
2545
2546fn map_registry_error(error: RegistryError) -> DirectProxyError {
2547 let code = match error {
2548 RegistryError::InvalidSessionIdentifier => ProxyErrorCode::HeaderRejected,
2549 RegistryError::SessionExpired => ProxyErrorCode::SessionExpired,
2550 RegistryError::SessionGenerationConflict => ProxyErrorCode::SessionGenerationConflict,
2551 RegistryError::ActiveCapacityReached => ProxyErrorCode::RegistryCapacity,
2552 RegistryError::EntropyUnavailable
2553 | RegistryError::SessionCreationFailed
2554 | RegistryError::EpochExhausted
2555 | RegistryError::InternalFailure => ProxyErrorCode::ProxyConfiguration,
2556 };
2557 code.error(ProxyErrorPhase::Session)
2558}
2559
2560fn map_codec_error(error: CodecError) -> DirectProxyError {
2561 let codec_phase = error.phase();
2562 let opaque_carrier = error.opaque_carrier();
2563 let request_direction = matches!(
2564 codec_phase,
2565 CodecPhase::RequestParse | CodecPhase::RequestProtection | CodecPhase::RequestProof
2566 );
2567 let code = match error.code() {
2568 CodecErrorCode::DuplicateObjectKey => ProxyErrorCode::DuplicateObjectKey,
2569 CodecErrorCode::InternalCoverageFailure => ProxyErrorCode::InternalCoverageFailure,
2570 CodecErrorCode::ControlWouldMutate => ProxyErrorCode::ControlWouldMutate,
2571 CodecErrorCode::OpaqueMediaUninspected => ProxyErrorCode::OpaqueMediaUninspected,
2572 CodecErrorCode::SignedMutationRequired => ProxyErrorCode::SignedMutationRequired,
2573 CodecErrorCode::SignedSurfaceMalformed => ProxyErrorCode::SignedSurfaceMalformed,
2574 CodecErrorCode::InvalidSseLifecycle | CodecErrorCode::IncompleteSseStream => {
2575 ProxyErrorCode::InvalidSseLifecycle
2576 }
2577 CodecErrorCode::InvalidContentBlockIndex | CodecErrorCode::InvalidContentBlockDelta => {
2578 ProxyErrorCode::InvalidContentBlockIndex
2579 }
2580 CodecErrorCode::DepthLimitExceeded
2581 | CodecErrorCode::NodeLimitExceeded
2582 | CodecErrorCode::StringLimitExceeded
2583 | CodecErrorCode::NumberLimitExceeded
2584 | CodecErrorCode::BodyLimitExceeded
2585 | CodecErrorCode::GrowthLimitExceeded
2586 | CodecErrorCode::SseLineLimitExceeded
2587 | CodecErrorCode::SseFrameLimitExceeded
2588 | CodecErrorCode::SseEventLimitExceeded
2589 | CodecErrorCode::SseIndexLimitExceeded
2590 | CodecErrorCode::SseAccumulatorLimitExceeded => {
2591 if request_direction {
2592 ProxyErrorCode::RequestBodyLimitExceeded
2593 } else {
2594 ProxyErrorCode::ResponseBodyLimitExceeded
2595 }
2596 }
2597 CodecErrorCode::ProtectionFailedClosed
2598 | CodecErrorCode::RestoreFailedClosed
2599 | CodecErrorCode::ProviderOriginPii
2600 | CodecErrorCode::ProviderOriginNumeric
2601 | CodecErrorCode::TransformedKeyCollision
2602 | CodecErrorCode::InvalidUtf8
2603 | CodecErrorCode::InvalidJson
2604 | CodecErrorCode::InvalidWireFormat
2605 | CodecErrorCode::UnsupportedRequestSurface
2606 | CodecErrorCode::UnsupportedResponseSurface
2607 | CodecErrorCode::InvalidSseField
2608 | CodecErrorCode::InvalidSseFrame
2609 | CodecErrorCode::InvalidSseEventPair
2610 | CodecErrorCode::UnsafeFutureEvent
2611 | CodecErrorCode::ProviderErrorEvent => ProxyErrorCode::InvalidToken,
2612 };
2613 let phase = match codec_phase {
2614 CodecPhase::RequestParse => ProxyErrorPhase::RequestValidation,
2615 CodecPhase::RequestProtection | CodecPhase::RequestProof => {
2616 ProxyErrorPhase::RequestTransform
2617 }
2618 CodecPhase::ResponseParse => ProxyErrorPhase::ResponseValidation,
2619 CodecPhase::ResponseRestore | CodecPhase::ResponseProof => {
2620 ProxyErrorPhase::ResponseTransform
2621 }
2622 CodecPhase::SseDecode | CodecPhase::SseLifecycle => ProxyErrorPhase::ResponseValidation,
2623 CodecPhase::SseReplay => ProxyErrorPhase::ResponseReplay,
2624 };
2625 let direct_error = code.error(phase);
2626 if let Some(opaque_carrier) = opaque_carrier {
2627 direct_error.with_opaque_carrier(opaque_carrier)
2628 } else {
2629 direct_error
2630 }
2631}
2632
2633fn direct_error_response(error: DirectProxyError) -> Response {
2634 let mut error_body = serde_json::json!({
2635 "type": "api_error",
2636 "message": "proxy_validation_failed",
2637 "code": format!("{:?}", error.code()),
2638 "phase": format!("{:?}", error.phase()),
2639 });
2640 if let (Some(name), Some(reason)) = (error.header_name(), error.header_rejection_reason()) {
2641 error_body["header"] = serde_json::json!({
2642 "name": name,
2643 "reason": reason.as_str(),
2644 });
2645 }
2646 if let Some(carrier) = error.opaque_carrier() {
2647 error_body["carrier"] = serde_json::json!({
2648 "scheme": carrier.scheme().as_str(),
2649 "span": {
2650 "byte_offset": carrier.byte_offset(),
2651 "byte_length": carrier.byte_length(),
2652 },
2653 "content_block_index": carrier.content_block_index(),
2654 });
2655 }
2656 let body = serde_json::json!({
2657 "type": "error",
2658 "error": error_body,
2659 });
2660 let mut response = Response::builder()
2661 .status(error.http_status())
2662 .header(axum::http::header::CONTENT_TYPE, "application/json");
2663 if let Some(seconds) = error.retry_after_seconds() {
2664 response = response.header(axum::http::header::RETRY_AFTER, seconds.to_string());
2665 }
2666 response
2667 .body(Body::from(body.to_string()))
2668 .unwrap_or_else(|_| Response::new(Body::empty()))
2669}
2670
2671async fn session_for(
2672 state: &AppState,
2673 headers: &HeaderMap,
2674) -> Result<Arc<LegacySession>, ProxyError> {
2675 let now = Instant::now();
2676 let id = headers
2677 .get(SESSION_HEADER)
2678 .and_then(|value| value.to_str().ok())
2679 .filter(|value| !value.trim().is_empty())
2680 .map(str::to_string)
2681 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
2682
2683 #[cfg(test)]
2684 {
2685 let barrier = SESSION_FOR_MISS_BARRIERS.lock().unwrap().get(&id).cloned();
2686 if let Some(barrier) = barrier {
2687 barrier.wait().await;
2688 }
2689 }
2690
2691 let mut sessions = state.sessions.write().await;
2692 sessions.retain(|_, entry| entry.expires_at > now);
2693 match sessions.entry(id) {
2694 std::collections::hash_map::Entry::Occupied(entry) => Ok(entry.get().session.clone()),
2695 std::collections::hash_map::Entry::Vacant(entry) => {
2696 let session = Arc::new(LegacySession::new(
2697 Session::new(Scope::Conversation(entry.key().clone()))
2698 .map_err(|source| ProxyError::Pipeline { source })?,
2699 ));
2700 entry.insert(SessionEntry {
2701 session: session.clone(),
2702 expires_at: now + state.config.session_ttl,
2703 });
2704 Ok(session)
2705 }
2706 }
2707}
2708
2709const NON_CARRIER_NUMERIC_KEYS: &[&str] = &[
2727 "budget_tokens",
2728 "candidateCount",
2729 "frequencyPenalty",
2730 "frequency_penalty",
2731 "logprobs",
2732 "maxOutputTokens",
2733 "max_completion_tokens",
2734 "max_tokens",
2735 "n",
2736 "presencePenalty",
2737 "presence_penalty",
2738 "seed",
2739 "temperature",
2740 "topK",
2741 "topP",
2742 "top_k",
2743 "top_logprobs",
2744 "top_p",
2745];
2746
2747const CARRIER_SUBTREE_KEYS: &[&str] = &[
2755 "functionDeclarations",
2756 "functions",
2757 "input_schema",
2758 "json_schema",
2759 "parameters",
2760 "properties",
2761 "response_format",
2762 "schema",
2763 "tools",
2764];
2765
2766fn redact_surfaces(
2779 pipeline: &Pipeline,
2780 session: &Session,
2781 surfaces: Vec<crate::adapter::PiiSurface<'_>>,
2782 locale_chain: &[LocaleTag],
2783 dictionaries: &DictionaryBundle,
2784) -> Result<RedactedSurfaces, ProxyError> {
2785 let mut redacted = RedactedSurfaces::default();
2786 for surface in surfaces {
2787 let clean = pipeline
2788 .pseudonymize_with_detect_context(
2789 session,
2790 RawDocument::Text(surface.text.clone()),
2791 locale_chain,
2792 dictionaries,
2793 )
2794 .map_err(|source| ProxyError::Pipeline { source })?;
2795 if let CleanDocument::Text(text) = clean {
2796 pipeline
2797 .admit_safety_nets(session, &text, locale_chain, dictionaries)
2798 .map_err(|source| ProxyError::Pipeline {
2799 source: source.into(),
2800 })?;
2801 collect_json_spelled_tokens(surface.text, &text, &mut redacted.json_spelled_tokens);
2802 *surface.text = text;
2803 }
2804 redacted.field_paths.insert(surface.field_path);
2805 }
2806 Ok(redacted)
2807}
2808
2809#[derive(Default)]
2810struct RedactedSurfaces {
2811 field_paths: HashSet<String>,
2813 json_spelled_tokens: HashSet<String>,
2815}
2816
2817fn collect_json_spelled_tokens(original: &str, clean: &str, tokens: &mut HashSet<String>) {
2823 if !matches!(clean.trim_start().as_bytes().first(), Some(b'{' | b'[')) {
2824 return;
2825 }
2826 let Ok(document) = serde_json::from_str::<Value>(clean) else {
2827 return;
2828 };
2829 let mut stack = vec![&document];
2830 while let Some(value) = stack.pop() {
2831 let mut record = |text: &str| {
2832 for token in token_shape::find_tokens(text) {
2833 if !original.contains(token) {
2834 tokens.insert(token.to_owned());
2835 }
2836 }
2837 };
2838 match value {
2839 Value::String(text) => record(text),
2840 Value::Array(items) => stack.extend(items),
2841 Value::Object(members) => {
2842 for (key, member) in members {
2843 record(key);
2844 stack.push(member);
2845 }
2846 }
2847 Value::Null | Value::Bool(_) | Value::Number(_) => {}
2848 }
2849 }
2850}
2851
2852fn residual_scan_request(
2868 pipeline: &Pipeline,
2869 body: &Value,
2870 surfaced: &HashSet<String>,
2871 locale_chain: &[LocaleTag],
2872 dictionaries: &DictionaryBundle,
2873) -> Result<(), ProxyError> {
2874 let session =
2875 Session::new(Scope::Ephemeral).map_err(|source| ProxyError::Pipeline { source })?;
2876 let scan = RequestResidualScan {
2877 pipeline,
2878 session,
2879 surfaced,
2880 locale_chain,
2881 dictionaries,
2882 };
2883 scan.walk("", "", body, false, true)
2884}
2885
2886struct RequestResidualScan<'a> {
2888 pipeline: &'a Pipeline,
2889 session: Session,
2890 surfaced: &'a HashSet<String>,
2891 locale_chain: &'a [LocaleTag],
2893 dictionaries: &'a DictionaryBundle,
2895}
2896
2897impl RequestResidualScan<'_> {
2898 fn probe(&self, field_path: &str, text: &str) -> Result<(), ProxyError> {
2899 let (clean, spans, report) = self
2900 .pipeline
2901 .clean_with_safety_net_detect_context(
2902 &self.session,
2903 RawDocument::Text(text.to_owned()),
2904 self.locale_chain,
2905 self.dictionaries,
2906 )
2907 .map_err(|source| ProxyError::Pipeline { source })?;
2908 let untouched = matches!(&clean, CleanDocument::Text(cleaned) if cleaned == text);
2914 if untouched && spans.is_empty() && report.suspects.is_empty() {
2922 return Ok(());
2923 }
2924 Err(ProxyError::UnsurfacedPii {
2925 field_path: field_path.to_string(),
2926 })
2927 }
2928
2929 fn walk(
2936 &self,
2937 path: &str,
2938 leaf_key: &str,
2939 value: &Value,
2940 in_carrier_subtree: bool,
2941 unambiguous: bool,
2942 ) -> Result<(), ProxyError> {
2943 match value {
2944 Value::String(text) => {
2945 if !(unambiguous && self.surfaced.contains(path)) {
2946 self.probe(path, text)?;
2947 }
2948 }
2949 Value::Number(number) => {
2950 if in_carrier_subtree || !NON_CARRIER_NUMERIC_KEYS.contains(&leaf_key) {
2951 self.probe(path, &number.to_string())?;
2952 }
2953 }
2954 Value::Array(items) => {
2955 for (index, item) in items.iter().enumerate() {
2956 self.walk(
2957 &format!("{path}[{index}]"),
2958 leaf_key,
2959 item,
2960 in_carrier_subtree,
2961 unambiguous,
2962 )?;
2963 }
2964 }
2965 Value::Object(map) => {
2966 for (key, child) in map {
2967 let child_path = if path.is_empty() {
2968 key.clone()
2969 } else {
2970 format!("{path}.{key}")
2971 };
2972 self.walk(
2973 &child_path,
2974 key,
2975 child,
2976 in_carrier_subtree || CARRIER_SUBTREE_KEYS.contains(&key.as_str()),
2977 unambiguous && !key.contains('.') && !key.contains('['),
2978 )?;
2979 }
2980 }
2981 Value::Null | Value::Bool(_) => {}
2982 }
2983 Ok(())
2984 }
2985}
2986
2987fn restore_surfaces(
2991 session: &LegacySession,
2992 surfaces: Vec<crate::adapter::PiiSurface<'_>>,
2993 json_output: bool,
2994) {
2995 let json_spelled_tokens = session.json_spelled_tokens();
2996 for surface in surfaces {
2997 let into_json_string = match surface.syntax {
2998 SurfaceSyntax::Text => false,
2999 SurfaceSyntax::Json => true,
3000 SurfaceSyntax::ModelOutput => json_output,
3001 };
3002 *surface.text = restore_text(
3003 &session.session,
3004 &json_spelled_tokens,
3005 surface.text,
3006 into_json_string,
3007 );
3008 }
3009}
3010
3011fn restore_text(
3022 session: &Session,
3023 json_spelled_tokens: &HashSet<String>,
3024 clean: &str,
3025 into_json_string: bool,
3026) -> String {
3027 let mut restored = String::new();
3028 let mut last = 0;
3029 for matched in token_shape::pattern().find_iter(clean) {
3030 restored.push_str(&clean[last..matched.start()]);
3031 let token = matched.as_str();
3032 match session.restore(token) {
3033 Some(raw)
3034 if into_json_string
3035 && !(json_spelled_tokens.contains(token) && is_json_string_content(&raw)) =>
3036 {
3037 push_json_string_contents(&mut restored, &raw);
3038 }
3039 Some(raw) => restored.push_str(&raw),
3040 None => restored.push_str(token),
3041 }
3042 last = matched.end();
3043 }
3044 restored.push_str(&clean[last..]);
3045 restored
3046}
3047
3048fn push_json_string_contents(out: &mut String, raw: &str) {
3050 let literal = Value::String(raw.to_owned()).to_string();
3051 out.push_str(&literal[1..literal.len() - 1]);
3052}
3053
3054fn is_json_string_content(text: &str) -> bool {
3056 serde_json::from_str::<String>(&format!("\"{text}\"")).is_ok()
3057}
3058
3059fn transform_sse(
3060 adapter: &Arc<dyn ProviderAdapter>,
3061 session: &LegacySession,
3062 bytes: &[u8],
3063 json_output: bool,
3064) -> Result<Vec<u8>, ProxyError> {
3065 let text = std::str::from_utf8(bytes).map_err(|err| ProxyError::SsePartialFrame {
3066 reason: err.to_string(),
3067 })?;
3068 let mut out = String::new();
3069 for frame in text.split("\n\n") {
3070 if frame.trim().is_empty() {
3071 continue;
3072 }
3073 let mut event_name = None;
3074 let mut data_lines = Vec::new();
3075 for line in frame.lines() {
3076 if let Some(rest) = line.strip_prefix("event:") {
3077 event_name = Some(rest.trim().to_string());
3078 } else if let Some(rest) = line.strip_prefix("data:") {
3079 data_lines.push(rest.trim_start());
3080 } else {
3081 out.push_str(line);
3082 out.push('\n');
3083 }
3084 }
3085 let data = data_lines.join("\n");
3086 if data == "[DONE]" {
3087 if let Some(name) = event_name {
3088 out.push_str("event: ");
3089 out.push_str(&name);
3090 out.push('\n');
3091 }
3092 out.push_str("data: [DONE]\n\n");
3093 continue;
3094 }
3095 let mut event = SseEvent {
3096 event: event_name.clone(),
3097 data: serde_json::from_str(&data)
3098 .map_err(|source| ProxyError::InvalidJson { source })?,
3099 };
3100 restore_surfaces(
3101 session,
3102 adapter.sse_event_pii_surfaces(&mut event),
3103 json_output,
3104 );
3105 if let Some(name) = event_name {
3106 out.push_str("event: ");
3107 out.push_str(&name);
3108 out.push('\n');
3109 }
3110 out.push_str("data: ");
3111 out.push_str(
3112 &serde_json::to_string(&event.data)
3113 .map_err(|source| ProxyError::InvalidJson { source })?,
3114 );
3115 out.push_str("\n\n");
3116 }
3117 Ok(out.into_bytes())
3118}
3119
3120fn upstream_url(base: &Url, uri: &Uri) -> Result<Url, ProxyError> {
3121 let path = uri.path_and_query().map(|pq| pq.as_str()).unwrap_or("/");
3122 base.join(path).map_err(|err| ProxyError::DaemonConfig {
3123 detail: err.to_string(),
3124 })
3125}
3126
3127fn forward_headers(
3128 mut request: reqwest::RequestBuilder,
3129 headers: &HeaderMap,
3130) -> reqwest::RequestBuilder {
3131 const FORWARDED: &[&str] = &[
3132 "authorization",
3133 "x-api-key",
3134 "anthropic-version",
3135 "openai-beta",
3136 "content-type",
3137 ];
3138 for name in FORWARDED {
3139 let Ok(header_name) = HeaderName::from_bytes(name.as_bytes()) else {
3140 continue;
3141 };
3142 if let Some(value) = headers.get(&header_name) {
3143 request = request.header(name.to_string(), value.clone());
3144 }
3145 }
3146 request
3147}
3148
3149fn proxy_error_response(err: ProxyError) -> Response {
3150 let status = match err {
3151 ProxyError::AdapterNotFound { .. } => StatusCode::NOT_FOUND,
3152 ProxyError::BodyTooLarge { .. } => StatusCode::PAYLOAD_TOO_LARGE,
3153 ProxyError::InvalidJson { .. } => StatusCode::BAD_REQUEST,
3154 ProxyError::UnsurfacedPii { .. } => StatusCode::UNPROCESSABLE_ENTITY,
3155 ProxyError::UpstreamUnreachable { .. } => StatusCode::BAD_GATEWAY,
3156 _ => StatusCode::INTERNAL_SERVER_ERROR,
3157 };
3158 let body = serde_json::json!({
3159 "error": proxy_error_name(&err),
3160 });
3161 (
3162 status,
3163 [(
3164 axum::http::header::CONTENT_TYPE,
3165 HeaderValue::from_static("application/json"),
3166 )],
3167 body.to_string(),
3168 )
3169 .into_response()
3170}
3171
3172fn proxy_error_name(err: &ProxyError) -> &'static str {
3173 match err {
3174 ProxyError::UpstreamUnreachable { .. } => "UpstreamUnreachable",
3175 ProxyError::AdapterNotFound { .. } => "AdapterNotFound",
3176 ProxyError::BodyTooLarge { .. } => "BodyTooLarge",
3177 ProxyError::InvalidJson { .. } => "InvalidJson",
3178 ProxyError::SsePartialFrame { .. } => "SsePartialFrame",
3179 ProxyError::UnsurfacedPii { .. } => "UnsurfacedPii",
3180 ProxyError::UnprovenCoverage => "UnprovenCoverage",
3181 ProxyError::Pipeline { .. } => "Pipeline",
3182 ProxyError::Server { .. } => "Server",
3183 ProxyError::DaemonAlreadyRunning { .. } => "DaemonAlreadyRunning",
3184 ProxyError::DaemonNotRunning => "DaemonNotRunning",
3185 ProxyError::DaemonPidfileStale { .. } => "DaemonPidfileStale",
3186 ProxyError::DaemonIo { .. } => "DaemonIo",
3187 ProxyError::DaemonConfig { .. } => "DaemonConfig",
3188 ProxyError::DaemonExitedEarly { .. } => "DaemonExitedEarly",
3189 }
3190}
3191
3192#[cfg(test)]
3193mod tests {
3194 use super::*;
3195 use std::process::Command;
3196 use std::sync::atomic::{AtomicUsize, Ordering};
3197 use std::sync::Mutex;
3198
3199 use axum::body::to_bytes;
3200 use axum::routing::post;
3201 use gaze::{Action, ClassRule, DefaultRule, PiiClass};
3202 use gaze_recognizers::RegexDetector;
3203
3204 use crate::codec::OutputProvenance;
3205
3206 const SYNTHETIC_EMAIL: &str = "alice@example.invalid";
3207 const SYNTHETIC_PHONE: &str = "+1-555-0199";
3208
3209 struct SyntheticPrincipalResolver;
3210
3211 impl PrincipalResolver for SyntheticPrincipalResolver {
3212 fn resolve(
3213 &self,
3214 _context: PrincipalContext<'_>,
3215 ) -> Result<crate::AuthenticatedPrincipal, crate::principal::PrincipalResolveError>
3216 {
3217 crate::AuthenticatedPrincipal::try_from_canonical_bytes(b"synthetic-principal".to_vec())
3218 }
3219 }
3220
3221 #[derive(Clone, Default)]
3222 struct Capture {
3223 hits: Arc<AtomicUsize>,
3224 headers: Arc<Mutex<Option<HeaderMap>>>,
3225 body: Arc<Mutex<Vec<u8>>>,
3226 }
3227
3228 async fn capture_ok(
3229 State(capture): State<Capture>,
3230 headers: HeaderMap,
3231 body: Bytes,
3232 ) -> Response {
3233 capture.hits.fetch_add(1, Ordering::SeqCst);
3234 *capture.headers.lock().unwrap() = Some(headers);
3235 *capture.body.lock().unwrap() = body.to_vec();
3236 Response::builder()
3237 .status(StatusCode::OK)
3238 .header("content-type", "application/json")
3239 .header("set-cookie", "synthetic=drop")
3240 .header("x-request-id", "drop-me")
3241 .body(Body::from(br#"{"ok":true}"#.as_slice()))
3242 .unwrap()
3243 }
3244
3245 async fn spawn_test_server(app: Router) -> (Url, tokio::task::JoinHandle<()>) {
3246 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3247 let address = listener.local_addr().unwrap();
3248 let handle = tokio::spawn(async move {
3249 axum::serve(listener, app).await.unwrap();
3250 });
3251 (
3252 Url::parse(&format!("http://{address}/v1/messages")).unwrap(),
3253 handle,
3254 )
3255 }
3256
3257 fn proved(bytes: &[u8]) -> ProvedRequestBody {
3258 let mut provenance = OutputProvenance::default();
3259 provenance.mark_final_buffer_verified();
3260 ProvedRequestBody::new(bytes.to_vec(), WireFormat::Json, provenance)
3261 }
3262
3263 fn direct_request(
3264 url: Url,
3265 expected: WireFormat,
3266 bytes: &[u8],
3267 ) -> Result<DirectRequest, DirectProxyError> {
3268 let mut headers = HeaderMap::new();
3269 headers.insert("content-type", HeaderValue::from_static("application/json"));
3270 headers.insert("x-api-key", HeaderValue::from_static("synthetic-key"));
3271 headers.insert("anthropic-version", HeaderValue::from_static("2023-06-01"));
3272 DirectRequest::from_proved(url, headers, proved(bytes), expected)
3273 }
3274
3275 fn committed_direct_request(
3276 url: Url,
3277 expected: WireFormat,
3278 bytes: &[u8],
3279 ) -> CommittedDirectRequest {
3280 let session = Session::new(Scope::Ephemeral).unwrap();
3281 let transaction = session.begin_transaction();
3282 PreparedDirectRequest::new(transaction, direct_request(url, expected, bytes).unwrap())
3283 .commit()
3284 .unwrap()
3285 }
3286
3287 fn email_pipeline() -> Pipeline {
3288 Pipeline::builder()
3289 .detector(RegexDetector::emails().unwrap())
3290 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
3291 .rule(DefaultRule::new(Action::Preserve))
3292 .build()
3293 .unwrap()
3294 }
3295
3296 #[test]
3297 fn net_admission_failure_discards_direct_staging_but_retains_legacy_live_mappings() {
3298 struct RejectNet;
3299 impl gaze::SafetyNet for RejectNet {
3300 fn id(&self) -> &str {
3301 "synthetic-reject"
3302 }
3303 fn supported_locales(&self) -> &[LocaleTag] {
3304 &[LocaleTag::Global]
3305 }
3306 fn check(
3307 &self,
3308 text: &str,
3309 _: gaze::SafetyNetContext<'_>,
3310 ) -> Result<Vec<gaze::LeakSuspect>, gaze::SafetyNetError> {
3311 if text.contains("reject-this") {
3312 Err(gaze::SafetyNetError::Runtime {
3313 message: "synthetic failure".into(),
3314 })
3315 } else {
3316 Ok(vec![])
3317 }
3318 }
3319 }
3320 let pipeline = Pipeline::builder()
3321 .detector(RegexDetector::emails().unwrap())
3322 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
3323 .rule(DefaultRule::new(Action::Preserve))
3324 .register_safety_net(RejectNet)
3325 .enable_prefix_cache()
3326 .build()
3327 .unwrap();
3328 let session = Session::new(Scope::Ephemeral).unwrap();
3329 let body = format!(
3330 r#"{{"model":"claude-test","max_tokens":32,"system":"{SYNTHETIC_EMAIL}","messages":[{{"role":"user","content":"reject-this"}}]}}"#
3331 );
3332 let request = direct_request(
3333 Url::parse("http://127.0.0.1:1/v1/messages").unwrap(),
3334 WireFormat::Json,
3335 body.as_bytes(),
3336 )
3337 .unwrap();
3338 assert!(prepare_and_commit_direct_request(
3339 &pipeline,
3340 &[LocaleTag::Global],
3341 &DictionaryBundle::default(),
3342 &session,
3343 &crate::codecs::anthropic::AnthropicMessagesCodec,
3344 body.as_bytes(),
3345 &request.url,
3346 &request.headers,
3347 WireFormat::Json,
3348 CodecLimits::default(),
3349 false,
3350 DEFAULT_MAX_REQUEST_BYTES,
3351 )
3352 .is_err());
3353 assert!(session.tokens().is_empty());
3354 assert_eq!(session.prefix_cache_entry_count(), 0);
3355
3356 let mut first = SYNTHETIC_EMAIL.to_owned();
3357 let mut second = "reject-this".to_owned();
3358 assert!(redact_surfaces(
3359 &pipeline,
3360 &session,
3361 vec![
3362 crate::adapter::PiiSurface {
3363 field_path: "$.first".into(),
3364 text: &mut first,
3365 syntax: SurfaceSyntax::Text,
3366 },
3367 crate::adapter::PiiSurface {
3368 field_path: "$.second".into(),
3369 text: &mut second,
3370 syntax: SurfaceSyntax::Text,
3371 },
3372 ],
3373 &[LocaleTag::Global],
3374 &DictionaryBundle::default()
3375 )
3376 .is_err());
3377 assert_eq!(session.tokens().len(), 1);
3378 assert_eq!(
3379 session.restore_strict_text(&first).unwrap(),
3380 SYNTHETIC_EMAIL
3381 );
3382 }
3383
3384 const DIRECT_REQUEST_BODY: &[u8] = br#"{"model":"claude-test","max_tokens":32,"messages":[{"role":"user","content":"synthetic request"}]}"#;
3385 const DIRECT_SSE_REQUEST_BODY: &[u8] = br#"{"model":"claude-test","max_tokens":32,"messages":[{"role":"user","content":"synthetic request"}],"stream":true}"#;
3386 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}}"#;
3387 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";
3388 const NON_STREAM_SESSION_ID: &str = "00000000-0000-4000-8000-000000000003";
3389
3390 #[derive(Clone)]
3391 struct ResponseControl {
3392 hits: Arc<AtomicUsize>,
3393 body_started: Arc<tokio::sync::Notify>,
3394 release: Arc<tokio::sync::Notify>,
3395 stream_error: bool,
3396 }
3397
3398 impl ResponseControl {
3399 fn new(stream_error: bool) -> Self {
3400 Self {
3401 hits: Arc::new(AtomicUsize::new(0)),
3402 body_started: Arc::new(tokio::sync::Notify::new()),
3403 release: Arc::new(tokio::sync::Notify::new()),
3404 stream_error,
3405 }
3406 }
3407 }
3408
3409 async fn delayed_json_response(
3410 State(control): State<ResponseControl>,
3411 _body: Bytes,
3412 ) -> Response {
3413 control.hits.fetch_add(1, Ordering::SeqCst);
3414 control.body_started.notify_one();
3415 control.release.notified().await;
3416 let body = if control.stream_error {
3417 b"{".as_slice()
3418 } else {
3419 DIRECT_JSON_RESPONSE
3420 };
3421 Response::builder()
3422 .status(StatusCode::OK)
3423 .header("content-type", "application/json")
3424 .body(Body::from(body))
3425 .unwrap()
3426 }
3427
3428 async fn controlled_sse_response(
3429 State(control): State<ResponseControl>,
3430 _body: Bytes,
3431 ) -> Response {
3432 control.hits.fetch_add(1, Ordering::SeqCst);
3433 let body_control = control.clone();
3434 let chunks = futures_util::stream::once(async move {
3435 body_control.body_started.notify_one();
3436 body_control.release.notified().await;
3437 if body_control.stream_error {
3438 Err(std::io::Error::other("synthetic stream failure"))
3439 } else {
3440 Ok(Bytes::from_static(DIRECT_SSE_RESPONSE))
3441 }
3442 });
3443 Response::builder()
3444 .status(StatusCode::OK)
3445 .header("content-type", "text/event-stream")
3446 .body(Body::from_stream(chunks))
3447 .unwrap()
3448 }
3449
3450 fn synthetic_principal() -> crate::AuthenticatedPrincipal {
3451 crate::AuthenticatedPrincipal::try_from_canonical_bytes(b"synthetic-principal".to_vec())
3452 .unwrap()
3453 }
3454
3455 fn direct_test_state(
3456 endpoint: Url,
3457 pipeline: Pipeline,
3458 session_policy: SessionPolicy,
3459 registry: Option<SessionRegistry>,
3460 ) -> AppState {
3461 let client_config = DirectClientConfig::default();
3462 let runtime = DirectRuntime {
3463 client: DirectClient::new(client_config).unwrap(),
3464 endpoint: endpoint.clone(),
3465 session_policy,
3466 registry,
3467 resolver: Arc::new(SyntheticPrincipalResolver),
3468 listener_scope: ListenerScope::Loopback,
3469 allowed_versions: Arc::from([DEFAULT_ANTHROPIC_VERSION.to_string()]),
3470 allowed_betas: Arc::from([]),
3471 codec_limits: CodecLimits::default(),
3472 client_config,
3473 ping_interval: Duration::from_secs(10),
3474 inspection_endpoint_codes: ProxyInspectionEndpointCodesV1::from_validated_origin(
3475 &endpoint,
3476 ),
3477 };
3478 AppState {
3479 config: Arc::new(ProxyConfig::new("127.0.0.1:0".parse().unwrap(), Vec::new())),
3480 pipeline: Arc::new(pipeline),
3481 client: Client::new(),
3482 direct: Some(Arc::new(runtime)),
3483 sessions: Arc::new(RwLock::new(HashMap::new())),
3484 started_at: Instant::now(),
3485 }
3486 }
3487
3488 #[tokio::test(flavor = "multi_thread", worker_threads = 8)]
3489 async fn concurrent_first_requests_share_one_session() {
3490 const TASK_COUNT: usize = 8;
3491 const SESSION_ID: &str = "synthetic-concurrent-session";
3492
3493 let state = Arc::new(AppState {
3494 config: Arc::new(ProxyConfig::new("127.0.0.1:0".parse().unwrap(), Vec::new())),
3495 pipeline: Arc::new(Pipeline::builder().build().unwrap()),
3496 client: Client::new(),
3497 direct: None,
3498 sessions: Arc::new(RwLock::new(HashMap::new())),
3499 started_at: Instant::now(),
3500 });
3501 SESSION_FOR_MISS_BARRIERS.lock().unwrap().insert(
3502 SESSION_ID.to_string(),
3503 Arc::new(tokio::sync::Barrier::new(TASK_COUNT)),
3504 );
3505
3506 let mut tasks = Vec::with_capacity(TASK_COUNT);
3507 for _ in 0..TASK_COUNT {
3508 let state = Arc::clone(&state);
3509 tasks.push(tokio::spawn(async move {
3510 let mut headers = HeaderMap::new();
3511 headers.insert(SESSION_HEADER, HeaderValue::from_static(SESSION_ID));
3512 session_for(&state, &headers).await.unwrap()
3513 }));
3514 }
3515
3516 let mut sessions = Vec::with_capacity(TASK_COUNT);
3517 for task in tasks {
3518 sessions.push(task.await.unwrap());
3519 }
3520 SESSION_FOR_MISS_BARRIERS.lock().unwrap().remove(SESSION_ID);
3521
3522 let first = &sessions[0];
3523 assert!(
3524 sessions.iter().all(|session| Arc::ptr_eq(first, session)),
3525 "concurrent first requests returned orphaned Sessions"
3526 );
3527 }
3528
3529 fn direct_test_ingress(
3530 endpoint: Url,
3531 session: DirectIngressSession,
3532 expected: WireFormat,
3533 body: &[u8],
3534 ) -> PreparedDirectIngress {
3535 let request = direct_request(endpoint, expected, body).unwrap();
3536 PreparedDirectIngress {
3537 outbound_headers: request.headers,
3538 session,
3539 }
3540 }
3541
3542 async fn wait_for_lease_count(registry: &SessionRegistry, expected: usize) {
3543 tokio::time::timeout(Duration::from_secs(1), async {
3544 loop {
3545 if registry.stats().total_leases() == expected {
3546 return;
3547 }
3548 tokio::task::yield_now().await;
3549 }
3550 })
3551 .await
3552 .unwrap();
3553 }
3554
3555 #[derive(Clone, Copy)]
3556 enum LimitFailureDirection {
3557 Request,
3558 Response,
3559 }
3560
3561 struct LimitFailureCodec(LimitFailureDirection);
3562
3563 impl BodyCodec for LimitFailureCodec {
3564 fn protect_request(
3565 &self,
3566 input: &[u8],
3567 _ctx: &mut RequestTransformContext<'_>,
3568 ) -> Result<ProvedRequestBody, CodecError> {
3569 match self.0 {
3570 LimitFailureDirection::Request => Err(CodecError::new(
3571 CodecErrorCode::DepthLimitExceeded,
3572 CodecPhase::RequestParse,
3573 )),
3574 LimitFailureDirection::Response => Ok(proved(input)),
3575 }
3576 }
3577
3578 fn restore_response(
3579 &self,
3580 _input: &[u8],
3581 _ctx: &mut ResponseTransformContext<'_>,
3582 ) -> Result<crate::codec::ProvedResponseBody, CodecError> {
3583 Err(CodecError::new(
3584 CodecErrorCode::DepthLimitExceeded,
3585 CodecPhase::ResponseParse,
3586 ))
3587 }
3588 }
3589
3590 async fn begin_delayed_non_stream(
3591 invalid_response: bool,
3592 ) -> (
3593 tokio::task::JoinHandle<Result<Response, DirectProxyError>>,
3594 ResponseControl,
3595 SessionRegistry,
3596 Instant,
3597 tokio::task::JoinHandle<()>,
3598 ) {
3599 let control = ResponseControl::new(invalid_response);
3600 let (endpoint, server) = spawn_test_server(
3601 Router::new()
3602 .route("/v1/messages", post(delayed_json_response))
3603 .with_state(control.clone()),
3604 )
3605 .await;
3606 let registry_config = crate::SessionRegistryConfig::default();
3607 let registry = SessionRegistry::with_secret(registry_config, [0x7c; 32]);
3608 let request_lease = registry
3609 .acquire(synthetic_principal(), NON_STREAM_SESSION_ID, None)
3610 .unwrap();
3611 let state = direct_test_state(
3612 endpoint.clone(),
3613 email_pipeline(),
3614 SessionPolicy::OpaqueHeaderContinuity(registry_config),
3615 Some(registry.clone()),
3616 );
3617 let ingress = direct_test_ingress(
3618 endpoint,
3619 DirectIngressSession::Continuity(request_lease),
3620 WireFormat::Json,
3621 DIRECT_REQUEST_BODY,
3622 );
3623 let request = tokio::spawn(async move {
3624 direct_proxy_inner(
3625 &state,
3626 ingress,
3627 DIRECT_REQUEST_BODY,
3628 &crate::codecs::anthropic::AnthropicMessagesCodec,
3629 )
3630 .await
3631 });
3632 control.body_started.notified().await;
3633 let expired_at = Instant::now() + Duration::from_secs(60 * 60);
3634 (request, control, registry, expired_at, server)
3635 }
3636
3637 async fn begin_controlled_sse(
3638 stream_error: bool,
3639 ) -> (
3640 Response,
3641 ResponseControl,
3642 SessionRegistry,
3643 Instant,
3644 tokio::task::JoinHandle<()>,
3645 ) {
3646 const SESSION_ID: &str = "00000000-0000-4000-8000-000000000002";
3647
3648 let control = ResponseControl::new(stream_error);
3649 let (endpoint, server) = spawn_test_server(
3650 Router::new()
3651 .route("/v1/messages", post(controlled_sse_response))
3652 .with_state(control.clone()),
3653 )
3654 .await;
3655 let registry_config = crate::SessionRegistryConfig::default();
3656 let registry = SessionRegistry::with_secret(registry_config, [0x6b; 32]);
3657 let request_lease = registry
3658 .acquire(synthetic_principal(), SESSION_ID, None)
3659 .unwrap();
3660 let state = direct_test_state(
3661 endpoint.clone(),
3662 email_pipeline(),
3663 SessionPolicy::OpaqueHeaderContinuity(registry_config),
3664 Some(registry.clone()),
3665 );
3666 let ingress = direct_test_ingress(
3667 endpoint,
3668 DirectIngressSession::Continuity(request_lease),
3669 WireFormat::Sse,
3670 DIRECT_SSE_REQUEST_BODY,
3671 );
3672
3673 let response = direct_proxy_inner(
3674 &state,
3675 ingress,
3676 DIRECT_SSE_REQUEST_BODY,
3677 &crate::codecs::anthropic::AnthropicMessagesCodec,
3678 )
3679 .await
3680 .unwrap();
3681 control.body_started.notified().await;
3682 let expired_at = Instant::now() + Duration::from_secs(60 * 60);
3683 (response, control, registry, expired_at, server)
3684 }
3685
3686 async fn assert_owned_listener_health(config: ProxyConfig, listener: TcpListener) {
3687 let actual = listener.local_addr().unwrap();
3688 let destination = SocketAddr::from(([127, 0, 0, 1], actual.port()));
3689 let server = serve_with_listener(config, Arc::new(email_pipeline()), listener);
3690 let health = async {
3691 let response = Client::new()
3692 .get(format!("http://{destination}/_gaze_proxy/healthz"))
3693 .send()
3694 .await
3695 .unwrap()
3696 .error_for_status()
3697 .unwrap()
3698 .json::<serde_json::Value>()
3699 .await
3700 .unwrap();
3701 assert_eq!(response["bind"], actual.to_string());
3702 assert_eq!(response["adapters"][0]["name"], "anthropic");
3703 };
3704 tokio::select! {
3705 result = server => panic!("owned listener server stopped: {result:?}"),
3706 result = tokio::time::timeout(Duration::from_secs(3), health) => result.unwrap(),
3707 }
3708 }
3710
3711 #[tokio::test]
3712 async fn reserved_port_rebind_fails_but_owned_listener_handoff_serves() {
3713 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3714 let bind = listener.local_addr().unwrap();
3715 let config = ProxyConfig::anthropic_direct(
3716 bind,
3717 AnthropicAdapter::new(Url::parse("http://127.0.0.1:1").unwrap()),
3718 );
3719 let error = serve(config.clone(), Arc::new(email_pipeline()))
3720 .await
3721 .unwrap_err();
3722 assert!(matches!(error, ProxyError::Server { source }
3723 if source.kind() == std::io::ErrorKind::AddrInUse));
3724 assert_owned_listener_health(config, listener).await;
3725 }
3726
3727 #[tokio::test]
3728 async fn owned_listener_resolves_ephemeral_port_in_health() {
3729 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3730 let config = ProxyConfig::anthropic_direct(
3731 "127.0.0.1:0".parse().unwrap(),
3732 AnthropicAdapter::new(Url::parse("http://127.0.0.1:1").unwrap()),
3733 );
3734 assert_owned_listener_health(config, listener).await;
3735 }
3736
3737 #[tokio::test]
3738 async fn owned_listener_rejects_wrong_ip_wildcard_and_port() {
3739 for (listen, configured) in [
3740 ("0.0.0.0:0", "127.0.0.1:0"),
3741 ("127.0.0.1:0", "0.0.0.0:0"),
3742 ("127.0.0.1:0", "127.0.0.2:0"),
3743 ("127.0.0.1:0", "127.0.0.1:1"),
3744 ] {
3745 let listener = TcpListener::bind(listen).await.unwrap();
3746 let config = ProxyConfig::anthropic_direct(
3747 configured.parse().unwrap(),
3748 AnthropicAdapter::new(Url::parse("http://127.0.0.1:1").unwrap()),
3749 );
3750 let error = serve_with_listener(config, Arc::new(email_pipeline()), listener)
3751 .await
3752 .unwrap_err();
3753 assert!(matches!(error, ProxyError::Server { source }
3754 if source.kind() == std::io::ErrorKind::InvalidInput));
3755 }
3756 }
3757
3758 #[tokio::test]
3759 async fn owned_non_loopback_listener_requires_explicit_resolver() {
3760 let listener = TcpListener::bind("0.0.0.0:0").await.unwrap();
3761 let config = ProxyConfig::anthropic_direct(
3762 "0.0.0.0:0".parse().unwrap(),
3763 AnthropicAdapter::new(Url::parse("http://127.0.0.1:1").unwrap()),
3764 );
3765 let error = serve_with_listener(config, Arc::new(email_pipeline()), listener)
3766 .await
3767 .unwrap_err();
3768 assert!(matches!(error, ProxyError::DaemonConfig { detail }
3769 if detail == "direct_principal_required"));
3770
3771 let listener = TcpListener::bind("0.0.0.0:0").await.unwrap();
3772 let config = ProxyConfig::anthropic_direct(
3773 "0.0.0.0:0".parse().unwrap(),
3774 AnthropicAdapter::builder(Url::parse("http://127.0.0.1:1").unwrap())
3775 .principal_resolver(Arc::new(SyntheticPrincipalResolver))
3776 .build()
3777 .unwrap(),
3778 );
3779 assert_owned_listener_health(config, listener).await;
3780 }
3781
3782 #[tokio::test]
3783 async fn serve_validates_before_rebind_and_owned_listener_validates_before_serving() {
3784 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3785 let config = ProxyConfig::anthropic_direct(
3786 listener.local_addr().unwrap(),
3787 AnthropicAdapter::new(Url::parse("http://synthetic:credential@127.0.0.1").unwrap()),
3788 );
3789 for error in [
3790 serve(config.clone(), Arc::new(email_pipeline()))
3791 .await
3792 .unwrap_err(),
3793 serve_with_listener(config, Arc::new(email_pipeline()), listener)
3794 .await
3795 .unwrap_err(),
3796 ] {
3797 assert!(matches!(error, ProxyError::DaemonConfig { detail }
3798 if detail == "direct_upstream_invalid"));
3799 }
3800 }
3801
3802 #[test]
3803 fn non_loopback_readiness_requires_an_explicit_principal_resolver() {
3804 let bind = "0.0.0.0:8787".parse().unwrap();
3805 let upstream = Url::parse("https://api.anthropic.com").unwrap();
3806 let missing = ProxyConfig::anthropic_direct(bind, AnthropicAdapter::new(upstream.clone()));
3807 assert!(DirectRuntime::from_config(&missing).is_err());
3808
3809 let configured = ProxyConfig::anthropic_direct(
3810 bind,
3811 AnthropicAdapter::builder(upstream)
3812 .principal_resolver(Arc::new(SyntheticPrincipalResolver))
3813 .build()
3814 .unwrap(),
3815 );
3816 assert!(DirectRuntime::from_config(&configured).unwrap().is_some());
3817 }
3818
3819 #[test]
3820 fn invalid_direct_upstream_fails_readiness_before_any_request() {
3821 let config = ProxyConfig::anthropic_direct(
3822 "127.0.0.1:8787".parse().unwrap(),
3823 AnthropicAdapter::new(Url::parse("https://synthetic:credential@127.0.0.1").unwrap()),
3824 );
3825 assert!(DirectRuntime::from_config(&config).is_err());
3826 }
3827
3828 #[tokio::test]
3829 async fn direct_send_requires_commit_is_single_use_and_captures_exact_proof() {
3830 let capture = Capture::default();
3831 let (url, server) = spawn_test_server(
3832 Router::new()
3833 .route("/v1/messages", post(capture_ok))
3834 .with_state(capture.clone()),
3835 )
3836 .await;
3837 let exact = br#"{"messages":[{"role":"user","content":"<Email_1>"}]}"#;
3838 let client = DirectClient::new(DirectClientConfig::default()).unwrap();
3839
3840 let session = Session::new(Scope::Ephemeral).unwrap();
3841 let prepared = PreparedDirectRequest::new(
3842 session.begin_transaction(),
3843 direct_request(url.clone(), WireFormat::Json, exact).unwrap(),
3844 );
3845 drop(prepared);
3846 assert!(session.tokens().is_empty());
3847 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
3848
3849 let mut lifecycle = DirectTransactionState::new();
3850 lifecycle.retry_generation_conflict().unwrap();
3851 assert_eq!(lifecycle.generation_retries(), 1);
3852 assert_eq!(
3853 lifecycle.retry_generation_conflict().unwrap_err().code(),
3854 ProxyErrorCode::SessionGenerationConflict
3855 );
3856 let response = client
3857 .execute(committed_direct_request(url, WireFormat::Json, exact))
3858 .await
3859 .unwrap();
3860 assert_eq!(response.lifecycle.state(), &DirectIoState::IoAttempted);
3861 assert_eq!(capture.hits.load(Ordering::SeqCst), 1);
3862 assert_eq!(&*capture.body.lock().unwrap(), exact);
3863 assert_eq!(response.response.body(), br#"{"ok":true}"#);
3864 assert_eq!(response.response.headers().len(), 1);
3865 assert!(response.response.headers().get("set-cookie").is_none());
3866 server.abort();
3867 }
3868
3869 #[test]
3870 fn signed_shadow_probe_reuses_only_the_previously_staged_mapping() {
3871 let pipeline = Pipeline::builder()
3872 .detector(RegexDetector::emails().unwrap())
3873 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
3874 .rule(DefaultRule::new(Action::Preserve))
3875 .enable_prefix_cache()
3876 .build()
3877 .unwrap();
3878 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
3879 let session = Session::new(Scope::Ephemeral).unwrap();
3880 let expected_token = format!("<{}:Email_1>", session.session_hex());
3881 let body = format!(
3882 r#"{{"model":"claude-test","max_tokens":32,"system":"{SYNTHETIC_EMAIL}","messages":[{{"role":"assistant","content":[{{"type":"thinking","thinking":"{expected_token}","signature":"opaque-signature"}}]}}]}}"#
3883 );
3884 let request = direct_request(
3885 Url::parse("https://api.anthropic.com/v1/messages").unwrap(),
3886 WireFormat::Json,
3887 body.as_bytes(),
3888 )
3889 .unwrap();
3890
3891 let committed = prepare_and_commit_direct_request_with_hook(
3892 &pipeline,
3893 &[LocaleTag::Global],
3894 &DictionaryBundle::default(),
3895 &session,
3896 &codec,
3897 body.as_bytes(),
3898 &request.url,
3899 &request.headers,
3900 WireFormat::Json,
3901 CodecLimits::default(),
3902 false,
3903 DEFAULT_MAX_REQUEST_BYTES,
3904 |attempt, live, staged| {
3905 assert_eq!(attempt, 0);
3906 assert_eq!(staged.prefix_cache_entry_count(), 0);
3907 assert_eq!(live.prefix_cache_entry_count(), 0);
3908 Ok(())
3909 },
3910 )
3911 .unwrap();
3912
3913 assert_eq!(session.tokens(), vec![expected_token.clone()]);
3914 assert_eq!(committed.snapshot.tokens(), vec![expected_token.clone()]);
3915 assert_eq!(committed.snapshot.prefix_cache_entry_count(), 0);
3916 assert_eq!(session.prefix_cache_entry_count(), 0);
3917 let protected = std::str::from_utf8(&committed.request.body).unwrap();
3918 assert!(!protected.contains(SYNTHETIC_EMAIL));
3919 assert_eq!(protected.matches(&expected_token).count(), 2);
3920 assert!(protected.contains(&format!(
3921 r#""thinking":"{expected_token}","signature":"opaque-signature""#
3922 )));
3923 }
3924
3925 #[test]
3926 fn final_request_reproof_is_idempotent_for_ordinary_and_format_preserving_tokens() {
3927 let endpoint = Url::parse("https://api.anthropic.com/v1/messages").unwrap();
3928 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
3929
3930 for action in [Action::Tokenize, Action::FormatPreserve] {
3931 let pipeline = Pipeline::builder()
3932 .detector(RegexDetector::emails().unwrap())
3933 .rule(ClassRule::new(PiiClass::Email, action))
3934 .rule(DefaultRule::new(Action::Preserve))
3935 .build()
3936 .unwrap();
3937 let session = Session::new(Scope::Ephemeral).unwrap();
3938 let original = format!(
3939 r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":"{SYNTHETIC_EMAIL}"}}]}}"#
3940 );
3941 let headers = direct_request(endpoint.clone(), WireFormat::Json, original.as_bytes())
3942 .unwrap()
3943 .headers;
3944 let first = prepare_and_commit_direct_request(
3945 &pipeline,
3946 &[LocaleTag::Global],
3947 &DictionaryBundle::default(),
3948 &session,
3949 &codec,
3950 original.as_bytes(),
3951 &endpoint,
3952 &headers,
3953 WireFormat::Json,
3954 CodecLimits::default(),
3955 false,
3956 DEFAULT_MAX_REQUEST_BYTES,
3957 )
3958 .unwrap_or_else(|error| panic!("{action:?} first proof failed: {error:?}"));
3959 let proved_once = first.request.body.clone();
3960 assert_eq!(session.tokens().len(), 1);
3961
3962 let second = prepare_and_commit_direct_request(
3963 &pipeline,
3964 &[LocaleTag::Global],
3965 &DictionaryBundle::default(),
3966 &session,
3967 &codec,
3968 &proved_once,
3969 &endpoint,
3970 &headers,
3971 WireFormat::Json,
3972 CodecLimits::default(),
3973 false,
3974 DEFAULT_MAX_REQUEST_BYTES,
3975 )
3976 .unwrap();
3977 assert_eq!(second.request.body, proved_once);
3978 assert_eq!(session.tokens().len(), 1);
3979 }
3980 }
3981
3982 #[tokio::test]
3983 async fn late_signed_failure_discards_all_staged_state_and_does_zero_io() {
3984 let capture = Capture::default();
3985 let (url, server) = spawn_test_server(
3986 Router::new()
3987 .route("/v1/messages", post(capture_ok))
3988 .with_state(capture.clone()),
3989 )
3990 .await;
3991 let pipeline = Pipeline::builder()
3992 .detector(RegexDetector::emails().unwrap())
3993 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
3994 .rule(DefaultRule::new(Action::Preserve))
3995 .enable_prefix_cache()
3996 .build()
3997 .unwrap();
3998 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
3999 let session = Session::new(Scope::Ephemeral).unwrap();
4000 let body = format!(
4001 r#"{{"model":"claude-test","max_tokens":32,"system":"{SYNTHETIC_EMAIL}","messages":[{{"role":"assistant","content":[{{"type":"thinking","thinking":"{SYNTHETIC_EMAIL}","signature":"opaque-signature"}}]}}]}}"#
4002 );
4003 let request = direct_request(url, WireFormat::Json, body.as_bytes()).unwrap();
4004 let error = prepare_and_commit_direct_request(
4005 &pipeline,
4006 &[LocaleTag::Global],
4007 &DictionaryBundle::default(),
4008 &session,
4009 &codec,
4010 body.as_bytes(),
4011 &request.url,
4012 &request.headers,
4013 WireFormat::Json,
4014 CodecLimits::default(),
4015 false,
4016 DEFAULT_MAX_REQUEST_BYTES,
4017 )
4018 .unwrap_err();
4019
4020 assert_eq!(error.code(), ProxyErrorCode::SignedMutationRequired);
4021 assert!(session.tokens().is_empty());
4022 assert_eq!(session.prefix_cache_entry_count(), 0);
4023 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4024
4025 let valid = format!(
4026 r#"{{"model":"claude-test","max_tokens":32,"system":"{SYNTHETIC_EMAIL} reports","messages":[{{"role":"user","content":"synthetic"}}]}}"#
4027 );
4028 prepare_and_commit_direct_request(
4029 &pipeline,
4030 &[LocaleTag::Global],
4031 &DictionaryBundle::default(),
4032 &session,
4033 &codec,
4034 valid.as_bytes(),
4035 &request.url,
4036 &request.headers,
4037 WireFormat::Json,
4038 CodecLimits::default(),
4039 false,
4040 DEFAULT_MAX_REQUEST_BYTES,
4041 )
4042 .unwrap();
4043 assert_eq!(session.tokens().len(), 1);
4044 server.abort();
4045 }
4046
4047 #[tokio::test]
4048 async fn logical_domain_failure_discards_state_before_commit_and_does_zero_io() {
4049 let capture = Capture::default();
4050 let (url, server) = spawn_test_server(
4051 Router::new()
4052 .route("/v1/messages", post(capture_ok))
4053 .with_state(capture.clone()),
4054 )
4055 .await;
4056 let pipeline = Pipeline::builder()
4057 .detector(RegexDetector::emails().unwrap())
4058 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
4059 .rule(DefaultRule::new(Action::Preserve))
4060 .enable_prefix_cache()
4061 .build()
4062 .unwrap();
4063 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4064 let session = Session::new(Scope::Ephemeral).unwrap();
4065 let split = 6;
4066 let (left, right) = SYNTHETIC_EMAIL.split_at(split);
4067
4068 let mut prefix_transaction = session.begin_transaction();
4069 pipeline
4070 .pseudonymize_transaction_with_detect_context(
4071 &mut prefix_transaction,
4072 RawDocument::Text(left.to_owned()),
4073 &[LocaleTag::Global],
4074 &DictionaryBundle::default(),
4075 )
4076 .unwrap();
4077 prefix_transaction.commit().unwrap();
4078 let baseline_entries = session.snapshot_entries();
4079 let baseline_tokens = session.tokens();
4080 let baseline_cache_entries = session.prefix_cache_entry_count();
4081 assert_eq!(baseline_cache_entries, 0);
4082
4083 let body = format!(
4084 r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":[{{"type":"text","text":"{left}"}},{{"type":"text","text":"{right}"}}]}}]}}"#
4085 );
4086 let request = direct_request(url, WireFormat::Json, body.as_bytes()).unwrap();
4087 let mut before_commit_calls = 0usize;
4088 let error = match prepare_and_commit_direct_request_with_hook(
4089 &pipeline,
4090 &[LocaleTag::Global],
4091 &DictionaryBundle::default(),
4092 &session,
4093 &codec,
4094 body.as_bytes(),
4095 &request.url,
4096 &request.headers,
4097 WireFormat::Json,
4098 CodecLimits::default(),
4099 false,
4100 DEFAULT_MAX_REQUEST_BYTES,
4101 |_, _, _| {
4102 before_commit_calls += 1;
4103 Ok(())
4104 },
4105 ) {
4106 Ok(_) => panic!("expected closed request proof failure"),
4107 Err(error) => error,
4108 };
4109
4110 assert_eq!(error.code(), ProxyErrorCode::ControlWouldMutate);
4111 assert_eq!(error.phase(), ProxyErrorPhase::RequestTransform);
4112 assert_eq!(error.http_status(), StatusCode::UNPROCESSABLE_ENTITY);
4113 assert_eq!(before_commit_calls, 0);
4114 assert_eq!(session.snapshot_entries(), baseline_entries);
4115 assert_eq!(session.tokens(), baseline_tokens);
4116 assert_eq!(session.prefix_cache_entry_count(), baseline_cache_entries);
4117 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4118 server.abort();
4119 }
4120
4121 #[test]
4122 fn suppression_only_request_sites_do_not_populate_the_prefix_cache() {
4123 let pipeline = Pipeline::builder()
4124 .detector(RegexDetector::emails().unwrap())
4125 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
4126 .rule(DefaultRule::new(Action::Preserve))
4127 .enable_prefix_cache()
4128 .build()
4129 .unwrap();
4130 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4131 let session = Session::new(Scope::Ephemeral).unwrap();
4132 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"}}}}]}"#;
4133 let request = direct_request(
4134 Url::parse("https://api.anthropic.com/v1/messages").unwrap(),
4135 WireFormat::Json,
4136 body,
4137 )
4138 .unwrap();
4139
4140 let committed = prepare_and_commit_direct_request_with_hook(
4141 &pipeline,
4142 &[LocaleTag::Global],
4143 &DictionaryBundle::default(),
4144 &session,
4145 &codec,
4146 body,
4147 &request.url,
4148 &request.headers,
4149 WireFormat::Json,
4150 CodecLimits::default(),
4151 false,
4152 DEFAULT_MAX_REQUEST_BYTES,
4153 |attempt, live, staged| {
4154 assert_eq!(attempt, 0);
4155 assert_eq!(live.prefix_cache_entry_count(), 0);
4156 assert_eq!(staged.prefix_cache_entry_count(), 0);
4157 Ok(())
4158 },
4159 )
4160 .unwrap();
4161
4162 assert!(session.tokens().is_empty());
4163 assert_eq!(committed.snapshot.prefix_cache_entry_count(), 0);
4164 assert_eq!(session.prefix_cache_entry_count(), 0);
4165 }
4166
4167 #[test]
4168 fn every_logical_domain_probe_family_leaves_prefix_storage_empty() {
4169 let pipeline = Pipeline::builder()
4170 .detector(RegexDetector::emails().unwrap())
4171 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
4172 .rule(DefaultRule::new(Action::Preserve))
4173 .enable_prefix_cache()
4174 .build()
4175 .unwrap();
4176 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4177 let session = Session::new(Scope::Ephemeral).unwrap();
4178 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"}]}"#;
4179 let request = direct_request(
4180 Url::parse("https://api.anthropic.com/v1/messages").unwrap(),
4181 WireFormat::Json,
4182 body,
4183 )
4184 .unwrap();
4185
4186 let committed = prepare_and_commit_direct_request_with_hook(
4187 &pipeline,
4188 &[LocaleTag::Global],
4189 &DictionaryBundle::default(),
4190 &session,
4191 &codec,
4192 body,
4193 &request.url,
4194 &request.headers,
4195 WireFormat::Json,
4196 CodecLimits::default(),
4197 false,
4198 DEFAULT_MAX_REQUEST_BYTES,
4199 |attempt, live, staged| {
4200 assert_eq!(attempt, 0);
4201 assert_eq!(live.prefix_cache_entry_count(), 0);
4202 assert_eq!(staged.prefix_cache_entry_count(), 0);
4203 Ok(())
4204 },
4205 )
4206 .unwrap();
4207
4208 assert_eq!(committed.snapshot.prefix_cache_entry_count(), 0);
4209 assert_eq!(session.prefix_cache_entry_count(), 0);
4210 }
4211
4212 struct RejectOnSecondAttemptCodec {
4213 protect_calls: AtomicUsize,
4214 }
4215
4216 impl BodyCodec for RejectOnSecondAttemptCodec {
4217 fn protect_request(
4218 &self,
4219 input: &[u8],
4220 ctx: &mut RequestTransformContext<'_>,
4221 ) -> Result<ProvedRequestBody, CodecError> {
4222 if self.protect_calls.fetch_add(1, Ordering::SeqCst) == 0 {
4223 crate::codecs::anthropic::AnthropicMessagesCodec.protect_request(input, ctx)
4224 } else {
4225 Err(CodecError::new(
4226 CodecErrorCode::ControlWouldMutate,
4227 CodecPhase::RequestProof,
4228 ))
4229 }
4230 }
4231
4232 fn restore_response(
4233 &self,
4234 input: &[u8],
4235 ctx: &mut ResponseTransformContext<'_>,
4236 ) -> Result<crate::codec::ProvedResponseBody, CodecError> {
4237 crate::codecs::anthropic::AnthropicMessagesCodec.restore_response(input, ctx)
4238 }
4239 }
4240
4241 #[tokio::test]
4242 async fn generation_conflict_reproofs_once_and_sends_only_the_winning_body() {
4243 let capture = Capture::default();
4244 let (url, server) = spawn_test_server(
4245 Router::new()
4246 .route("/v1/messages", post(capture_ok))
4247 .with_state(capture.clone()),
4248 )
4249 .await;
4250 let pipeline = email_pipeline();
4251 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4252 let session = Session::new(Scope::Ephemeral).unwrap();
4253 let body = format!(
4254 r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":"{SYNTHETIC_EMAIL}"}}]}}"#
4255 );
4256 let request = direct_request(url.clone(), WireFormat::Json, body.as_bytes()).unwrap();
4257 let mut hook_calls = 0;
4258 let committed = prepare_and_commit_direct_request_with_hook(
4259 &pipeline,
4260 &[LocaleTag::Global],
4261 &DictionaryBundle::default(),
4262 &session,
4263 &codec,
4264 body.as_bytes(),
4265 &request.url,
4266 &request.headers,
4267 WireFormat::Json,
4268 CodecLimits::default(),
4269 false,
4270 DEFAULT_MAX_REQUEST_BYTES,
4271 |attempt, live, _staged| {
4272 hook_calls += 1;
4273 if attempt == 0 {
4274 live.tokenize(&PiiClass::Name, "Dr. Schmidt").unwrap();
4275 }
4276 Ok(())
4277 },
4278 )
4279 .unwrap();
4280 assert_eq!(hook_calls, 2);
4281 assert_eq!(session.tokens().len(), 2);
4282
4283 DirectClient::new(DirectClientConfig::default())
4284 .unwrap()
4285 .execute(committed)
4286 .await
4287 .unwrap();
4288 assert_eq!(capture.hits.load(Ordering::SeqCst), 1);
4289 let captured = String::from_utf8(capture.body.lock().unwrap().clone()).unwrap();
4290 assert!(!captured.contains(SYNTHETIC_EMAIL));
4291 assert_eq!(captured.matches(":Email_1>").count(), 1);
4292 server.abort();
4293 }
4294
4295 #[tokio::test]
4296 async fn second_attempt_request_proof_failure_rejects_without_upstream_io() {
4297 let capture = Capture::default();
4298 let (url, server) = spawn_test_server(
4299 Router::new()
4300 .route("/v1/messages", post(capture_ok))
4301 .with_state(capture.clone()),
4302 )
4303 .await;
4304 let pipeline = Pipeline::builder()
4305 .rule(DefaultRule::new(Action::Preserve))
4306 .build()
4307 .unwrap();
4308 let codec = RejectOnSecondAttemptCodec {
4309 protect_calls: AtomicUsize::new(0),
4310 };
4311 let session = Session::new(Scope::Ephemeral).unwrap();
4312 let body = br#"{"model":"claude-test","max_tokens":32,"messages":[{"role":"user","content":"synthetic"}]}"#;
4313 let request = direct_request(url, WireFormat::Json, body).unwrap();
4314 let mut hook_calls = 0usize;
4315
4316 let error = prepare_and_commit_direct_request_with_hook(
4317 &pipeline,
4318 &[LocaleTag::Global],
4319 &DictionaryBundle::default(),
4320 &session,
4321 &codec,
4322 body,
4323 &request.url,
4324 &request.headers,
4325 WireFormat::Json,
4326 CodecLimits::default(),
4327 false,
4328 DEFAULT_MAX_REQUEST_BYTES,
4329 |attempt, live, _staged| {
4330 hook_calls += 1;
4331 assert_eq!(attempt, 0);
4332 live.tokenize(&PiiClass::Name, "Dr. Schmidt").unwrap();
4333 Ok(())
4334 },
4335 )
4336 .unwrap_err();
4337
4338 assert_eq!(codec.protect_calls.load(Ordering::SeqCst), 2);
4339 assert_eq!(hook_calls, 1);
4340 assert_eq!(error.code(), ProxyErrorCode::ControlWouldMutate);
4341 assert_eq!(error.phase(), ProxyErrorPhase::RequestTransform);
4342 assert_eq!(error.http_status(), StatusCode::UNPROCESSABLE_ENTITY);
4343 assert_eq!(session.tokens().len(), 1);
4344 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4345 server.abort();
4346 }
4347
4348 #[tokio::test]
4349 async fn second_generation_conflict_fails_closed_without_upstream_io() {
4350 let capture = Capture::default();
4351 let (url, server) = spawn_test_server(
4352 Router::new()
4353 .route("/v1/messages", post(capture_ok))
4354 .with_state(capture.clone()),
4355 )
4356 .await;
4357 let pipeline = Pipeline::builder()
4358 .detector(RegexDetector::emails().unwrap())
4359 .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
4360 .rule(DefaultRule::new(Action::Preserve))
4361 .enable_prefix_cache()
4362 .build()
4363 .unwrap();
4364 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4365 let session = Session::new(Scope::Ephemeral).unwrap();
4366 let body = format!(
4367 r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":"{SYNTHETIC_EMAIL}"}}]}}"#
4368 );
4369 let request = direct_request(url, WireFormat::Json, body.as_bytes()).unwrap();
4370 let error = prepare_and_commit_direct_request_with_hook(
4371 &pipeline,
4372 &[LocaleTag::Global],
4373 &DictionaryBundle::default(),
4374 &session,
4375 &codec,
4376 body.as_bytes(),
4377 &request.url,
4378 &request.headers,
4379 WireFormat::Json,
4380 CodecLimits::default(),
4381 false,
4382 DEFAULT_MAX_REQUEST_BYTES,
4383 |attempt, live, staged| {
4384 assert_eq!(staged.prefix_cache_entry_count(), 0);
4385 assert_eq!(live.prefix_cache_entry_count(), 0);
4386 live.tokenize(&PiiClass::Name, &format!("Dr. Schmidt {attempt}"))
4387 .unwrap();
4388 Ok(())
4389 },
4390 )
4391 .unwrap_err();
4392 assert_eq!(error.code(), ProxyErrorCode::SessionGenerationConflict);
4393 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4394 assert_eq!(session.prefix_cache_entry_count(), 0);
4395 assert!(session
4396 .tokens()
4397 .iter()
4398 .all(|token| !token.contains(":Email_")));
4399 server.abort();
4400 }
4401
4402 #[tokio::test]
4403 async fn cancellation_publishes_only_after_commit_and_never_rolls_it_back() {
4404 let session = Session::new(Scope::Ephemeral).unwrap();
4405 let mut staged = session.begin_transaction();
4406 staged.tokenize(&PiiClass::Email, SYNTHETIC_EMAIL).unwrap();
4407 let prepared = PreparedDirectRequest::new(
4408 staged,
4409 direct_request(
4410 Url::parse("https://api.anthropic.com/v1/messages").unwrap(),
4411 WireFormat::Json,
4412 br#"{}"#,
4413 )
4414 .unwrap(),
4415 );
4416 drop(prepared);
4417 assert!(session.tokens().is_empty());
4418
4419 let (url, server) = spawn_test_server(Router::new().route(
4420 "/v1/messages",
4421 post(|| async {
4422 tokio::time::sleep(Duration::from_secs(1)).await;
4423 ([("content-type", "application/json")], "{}")
4424 }),
4425 ))
4426 .await;
4427 let mut staged = session.begin_transaction();
4428 staged.tokenize(&PiiClass::Email, SYNTHETIC_EMAIL).unwrap();
4429 let committed = PreparedDirectRequest::new(
4430 staged,
4431 direct_request(url, WireFormat::Json, br#"{}"#).unwrap(),
4432 )
4433 .commit()
4434 .unwrap();
4435 assert_eq!(session.tokens().len(), 1);
4436 let client = DirectClient::new(DirectClientConfig::default()).unwrap();
4437 assert!(
4438 tokio::time::timeout(Duration::from_millis(20), client.execute(committed))
4439 .await
4440 .is_err()
4441 );
4442 assert_eq!(session.tokens().len(), 1);
4443 server.abort();
4444 }
4445
4446 #[tokio::test]
4447 async fn response_restore_uses_the_exact_committed_snapshot_after_live_mutation() {
4448 let pipeline = email_pipeline();
4449 let session = Session::new(Scope::Ephemeral).unwrap();
4450 let mut transaction = session.begin_transaction();
4451 let committed_token = transaction
4452 .tokenize(&PiiClass::Email, SYNTHETIC_EMAIL)
4453 .unwrap();
4454 let snapshot = transaction.commit().unwrap();
4455 session.tokenize(&PiiClass::Name, "Dr. Schmidt").unwrap();
4456 assert_eq!(session.tokens().len(), 2);
4457
4458 let body = format!(
4459 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}}}}"#
4460 );
4461 let mut headers = HeaderMap::new();
4462 headers.insert(
4463 axum::http::header::CONTENT_TYPE,
4464 HeaderValue::from_static("application/json"),
4465 );
4466 let executed = CommittedDirectResponse {
4467 snapshot,
4468 response: DirectResponse {
4469 head: DirectResponseHead {
4470 status: StatusCode::OK,
4471 format: WireFormat::Json,
4472 headers,
4473 },
4474 body: body.into_bytes(),
4475 },
4476 lifecycle: DirectTransactionState {
4477 state: DirectIoState::IoAttempted,
4478 generation_retries: 0,
4479 },
4480 };
4481 let response = restore_direct_response(
4482 &pipeline,
4483 &[LocaleTag::Global],
4484 &DictionaryBundle::default(),
4485 &crate::codecs::anthropic::AnthropicMessagesCodec,
4486 executed,
4487 CodecLimits::default(),
4488 None,
4489 )
4490 .unwrap();
4491 let restored = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
4492 assert!(std::str::from_utf8(&restored)
4493 .unwrap()
4494 .contains(SYNTHETIC_EMAIL));
4495 assert!(!std::str::from_utf8(&restored)
4496 .unwrap()
4497 .contains("Dr. Schmidt"));
4498 }
4499
4500 #[tokio::test]
4501 async fn transformed_request_overlimit_rejects_before_commit_and_zero_io() {
4502 const SESSION_ID: &str = "00000000-0000-4000-8000-000000000001";
4503
4504 let capture = Capture::default();
4505 let (url, server) = spawn_test_server(
4506 Router::new()
4507 .route("/v1/messages", post(capture_ok))
4508 .with_state(capture.clone()),
4509 )
4510 .await;
4511 let phone_class = PiiClass::Custom("phone".to_string());
4512 let pipeline = Pipeline::builder()
4513 .detector(RegexDetector::new(r"\+1-555-0199", phone_class.clone()).unwrap())
4514 .rule(ClassRule::new(phone_class, Action::Tokenize))
4515 .rule(DefaultRule::new(Action::Preserve))
4516 .enable_prefix_cache()
4517 .build()
4518 .unwrap();
4519 let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4520 let body = format!(
4521 r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":"{SYNTHETIC_PHONE}"}}]}}"#
4522 );
4523 let request_limit = body.len() + 1;
4524 let probe_session = Session::new(Scope::Ephemeral).unwrap();
4525 let probe_request = direct_request(url.clone(), WireFormat::Json, body.as_bytes()).unwrap();
4526 let transformed = prepare_and_commit_direct_request(
4527 &pipeline,
4528 &[LocaleTag::Global],
4529 &DictionaryBundle::default(),
4530 &probe_session,
4531 &codec,
4532 body.as_bytes(),
4533 &probe_request.url,
4534 &probe_request.headers,
4535 WireFormat::Json,
4536 CodecLimits::default(),
4537 false,
4538 DEFAULT_MAX_REQUEST_BYTES,
4539 )
4540 .unwrap();
4541 assert!(body.len() < request_limit);
4542 assert!(transformed.request.body.len() > request_limit);
4543
4544 let registry_config = crate::SessionRegistryConfig::default();
4545 let registry = SessionRegistry::with_secret(registry_config, [0x5a; 32]);
4546 let observed_lease = registry
4547 .acquire(
4548 crate::AuthenticatedPrincipal::try_from_canonical_bytes(
4549 b"synthetic-principal".to_vec(),
4550 )
4551 .unwrap(),
4552 SESSION_ID,
4553 None,
4554 )
4555 .unwrap();
4556 let request_lease = registry
4557 .acquire(
4558 crate::AuthenticatedPrincipal::try_from_canonical_bytes(
4559 b"synthetic-principal".to_vec(),
4560 )
4561 .unwrap(),
4562 SESSION_ID,
4563 None,
4564 )
4565 .unwrap();
4566 let client_config = DirectClientConfig::default()
4567 .try_with_max_request_bytes(request_limit)
4568 .unwrap();
4569 let runtime = DirectRuntime {
4570 client: DirectClient::new(client_config).unwrap(),
4571 endpoint: url.clone(),
4572 session_policy: SessionPolicy::OpaqueHeaderContinuity(registry_config),
4573 registry: Some(registry),
4574 resolver: Arc::new(SyntheticPrincipalResolver),
4575 listener_scope: ListenerScope::Loopback,
4576 allowed_versions: Arc::from([DEFAULT_ANTHROPIC_VERSION.to_string()]),
4577 allowed_betas: Arc::from([]),
4578 codec_limits: CodecLimits::default(),
4579 client_config,
4580 ping_interval: Duration::from_secs(10),
4581 inspection_endpoint_codes: ProxyInspectionEndpointCodesV1::from_validated_origin(&url),
4582 };
4583 let state = AppState {
4584 config: Arc::new(ProxyConfig::new("127.0.0.1:0".parse().unwrap(), Vec::new())),
4585 pipeline: Arc::new(pipeline),
4586 client: Client::new(),
4587 direct: Some(Arc::new(runtime)),
4588 sessions: Arc::new(RwLock::new(HashMap::new())),
4589 started_at: Instant::now(),
4590 };
4591 let ingress = PreparedDirectIngress {
4592 outbound_headers: probe_request.headers.clone(),
4593 session: DirectIngressSession::Continuity(request_lease),
4594 };
4595
4596 let error = direct_proxy_inner(&state, ingress, body.as_bytes(), &codec)
4597 .await
4598 .unwrap_err();
4599
4600 assert_eq!(error.code(), ProxyErrorCode::RequestBodyLimitExceeded);
4601 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4602 assert!(observed_lease.session().tokens().is_empty());
4603 assert!(observed_lease.session().snapshot_entries().is_empty());
4604
4605 let extended_body = format!(
4606 r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":"{SYNTHETIC_PHONE} reports"}}]}}"#
4607 );
4608 let recovered = prepare_and_commit_direct_request(
4609 &state.pipeline,
4610 state.config.locale_chain(),
4611 state.config.dictionaries(),
4612 observed_lease.session(),
4613 &codec,
4614 extended_body.as_bytes(),
4615 &url,
4616 &probe_request.headers,
4617 WireFormat::Json,
4618 CodecLimits::default(),
4619 false,
4620 DEFAULT_MAX_REQUEST_BYTES,
4621 )
4622 .unwrap();
4623 let first_token = format!(
4624 "<{}:Custom:phone_1>",
4625 observed_lease.session().session_hex()
4626 );
4627 assert_eq!(observed_lease.session().tokens(), vec![first_token.clone()]);
4628 assert!(std::str::from_utf8(&recovered.request.body)
4629 .unwrap()
4630 .contains(&first_token));
4631 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4632 server.abort();
4633 }
4634
4635 #[tokio::test]
4636 async fn overlimit_proved_request_and_duplicate_credentials_do_zero_io() {
4637 let capture = Capture::default();
4638 let (url, server) = spawn_test_server(
4639 Router::new()
4640 .route("/v1/messages", post(capture_ok))
4641 .with_state(capture.clone()),
4642 )
4643 .await;
4644 let config = DirectClientConfig::default()
4645 .try_with_max_request_bytes(4)
4646 .unwrap();
4647 let client = DirectClient::new(config).unwrap();
4648 let error = client
4649 .execute(committed_direct_request(
4650 url.clone(),
4651 WireFormat::Json,
4652 br#"{"over":true}"#,
4653 ))
4654 .await
4655 .unwrap_err();
4656 assert_eq!(error.code(), ProxyErrorCode::RequestBodyLimitExceeded);
4657 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4658
4659 for singleton in ["x-api-key", "anthropic-version", "anthropic-beta"] {
4660 let mut duplicate = HeaderMap::new();
4661 duplicate.insert("content-type", HeaderValue::from_static("application/json"));
4662 duplicate.append(singleton, HeaderValue::from_static("synthetic-value-1"));
4663 duplicate.append(singleton, HeaderValue::from_static("synthetic-value-2"));
4664 let error = DirectRequest::from_proved(
4665 url.clone(),
4666 duplicate,
4667 proved(br#"{}"#),
4668 WireFormat::Json,
4669 )
4670 .err()
4671 .unwrap();
4672 assert_eq!(error.code(), ProxyErrorCode::HeaderRejected);
4673 }
4674
4675 let mut unverified_headers = HeaderMap::new();
4676 unverified_headers.insert("content-type", HeaderValue::from_static("application/json"));
4677 let unverified = ProvedRequestBody::new(
4678 br#"{}"#.to_vec(),
4679 WireFormat::Json,
4680 OutputProvenance::default(),
4681 );
4682 let error =
4683 DirectRequest::from_proved(url, unverified_headers, unverified, WireFormat::Json)
4684 .err()
4685 .unwrap();
4686 assert_eq!(error.code(), ProxyErrorCode::InvalidProvenance);
4687 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4688 server.abort();
4689 }
4690
4691 #[tokio::test]
4692 async fn header_rejections_name_only_the_header_and_disambiguate_the_cause() {
4693 const SECRET_ONE: &str = "synthetic-credential-one";
4694 const SECRET_TWO: &str = "synthetic-credential-two";
4695
4696 let mut not_allowlisted = HeaderMap::new();
4697 not_allowlisted.insert("x-synthetic-extra", HeaderValue::from_static(SECRET_ONE));
4698
4699 let mut empty = HeaderMap::new();
4700 empty.insert("anthropic-beta", HeaderValue::from_static(""));
4701
4702 let mut duplicate = HeaderMap::new();
4703 duplicate.append("x-api-key", HeaderValue::from_static(SECRET_ONE));
4704 duplicate.append("x-api-key", HeaderValue::from_static(SECRET_TWO));
4705
4706 let cases = [
4707 (
4708 not_allowlisted,
4709 "x-synthetic-extra",
4710 HeaderRejectionReason::NameNotAllowlisted,
4711 ),
4712 (empty, "anthropic-beta", HeaderRejectionReason::EmptyValue),
4713 (
4714 duplicate,
4715 "x-api-key",
4716 HeaderRejectionReason::DuplicateSingleton,
4717 ),
4718 ];
4719
4720 for (headers, expected_name, expected_reason) in cases {
4721 let error = validate_outbound_headers(&headers).unwrap_err();
4722 assert_eq!(error.code(), ProxyErrorCode::HeaderRejected);
4723 assert_eq!(error.phase(), ProxyErrorPhase::RequestValidation);
4724 assert_eq!(error.header_name().as_deref(), Some(expected_name));
4725 assert_eq!(error.header_rejection_reason(), Some(expected_reason));
4726
4727 let debug = format!("{error:?}");
4728 assert!(debug.contains(expected_name));
4729 assert!(!debug.contains(SECRET_ONE));
4730 assert!(!debug.contains(SECRET_TWO));
4731
4732 let response = direct_error_response(error);
4733 let bytes = to_bytes(response.into_body(), 4 * 1024).await.unwrap();
4734 let body: Value = serde_json::from_slice(&bytes).unwrap();
4735 assert_eq!(body["error"]["phase"], "RequestValidation");
4736 assert_eq!(body["error"]["header"]["name"], expected_name);
4737 assert_eq!(body["error"]["header"]["reason"], expected_reason.as_str());
4738 let rendered = String::from_utf8(bytes.to_vec()).unwrap();
4739 assert!(!rendered.contains(SECRET_ONE));
4740 assert!(!rendered.contains(SECRET_TWO));
4741 }
4742
4743 let maximum_name = format!("x-{}", "a".repeat(126));
4744 let mut maximum_length = HeaderMap::new();
4745 maximum_length.insert(
4746 HeaderName::from_bytes(maximum_name.as_bytes()).unwrap(),
4747 HeaderValue::from_static(SECRET_ONE),
4748 );
4749 let error = validate_outbound_headers(&maximum_length).unwrap_err();
4750 assert_eq!(error.header_name().as_deref(), Some(maximum_name.as_str()));
4751 assert_eq!(
4752 error.header_rejection_reason(),
4753 Some(HeaderRejectionReason::NameNotAllowlisted)
4754 );
4755 }
4756
4757 #[test]
4758 fn actionable_header_reporting_preserves_the_rejected_request_set() {
4759 let mut valid = HeaderMap::new();
4760 valid.insert("content-type", HeaderValue::from_static("application/json"));
4761 valid.insert("x-api-key", HeaderValue::from_static("synthetic-key"));
4762 valid.insert("anthropic-version", HeaderValue::from_static("2023-06-01"));
4763
4764 let mut not_allowlisted = valid.clone();
4765 not_allowlisted.insert(
4766 "x-synthetic-extra",
4767 HeaderValue::from_static("synthetic-value"),
4768 );
4769
4770 let mut empty = valid.clone();
4771 empty.insert("anthropic-version", HeaderValue::from_static(""));
4772
4773 let mut duplicate = valid.clone();
4774 duplicate.append("x-api-key", HeaderValue::from_static("synthetic-key-two"));
4775
4776 let rejected = [valid, not_allowlisted, empty, duplicate]
4777 .iter()
4778 .map(|headers| validate_outbound_headers(headers).is_err())
4779 .collect::<Vec<_>>();
4780 assert_eq!(rejected, vec![false, true, true, true]);
4781 }
4782
4783 #[tokio::test]
4784 async fn redirects_retries_and_response_overflow_remain_disabled() {
4785 let target_hits = Arc::new(AtomicUsize::new(0));
4786 let target_hits_for_handler = Arc::clone(&target_hits);
4787 let (target_url, target) = spawn_test_server(Router::new().route(
4788 "/v1/messages",
4789 post(move || {
4790 let hits = Arc::clone(&target_hits_for_handler);
4791 async move {
4792 hits.fetch_add(1, Ordering::SeqCst);
4793 StatusCode::OK
4794 }
4795 }),
4796 ))
4797 .await;
4798 let origin_hits = Arc::new(AtomicUsize::new(0));
4799 let origin_hits_for_handler = Arc::clone(&origin_hits);
4800 let location = target_url.to_string();
4801 let (origin_url, origin) = spawn_test_server(Router::new().route(
4802 "/v1/messages",
4803 post(move || {
4804 let hits = Arc::clone(&origin_hits_for_handler);
4805 let location = location.clone();
4806 async move {
4807 hits.fetch_add(1, Ordering::SeqCst);
4808 Response::builder()
4809 .status(StatusCode::TEMPORARY_REDIRECT)
4810 .header("location", location)
4811 .body(Body::from("synthetic redirect body"))
4812 .unwrap()
4813 }
4814 }),
4815 ))
4816 .await;
4817 let client = DirectClient::new(DirectClientConfig::default()).unwrap();
4818 let error = client
4819 .execute(committed_direct_request(
4820 origin_url,
4821 WireFormat::Json,
4822 br#"{}"#,
4823 ))
4824 .await
4825 .unwrap_err();
4826 assert_eq!(error.code(), ProxyErrorCode::UpstreamRedirect);
4827 assert_eq!(origin_hits.load(Ordering::SeqCst), 1);
4828 assert_eq!(target_hits.load(Ordering::SeqCst), 0);
4829 origin.abort();
4830 target.abort();
4831
4832 let attempts = Arc::new(AtomicUsize::new(0));
4833 let attempts_for_handler = Arc::clone(&attempts);
4834 let (url, server) = spawn_test_server(Router::new().route(
4835 "/v1/messages",
4836 post(move || {
4837 let attempts = Arc::clone(&attempts_for_handler);
4838 async move {
4839 attempts.fetch_add(1, Ordering::SeqCst);
4840 StatusCode::SERVICE_UNAVAILABLE
4841 }
4842 }),
4843 ))
4844 .await;
4845 let error = client
4846 .execute(committed_direct_request(url, WireFormat::Json, br#"{}"#))
4847 .await
4848 .unwrap_err();
4849 assert_eq!(error.code(), ProxyErrorCode::UpstreamUnavailable);
4850 assert_eq!(attempts.load(Ordering::SeqCst), 1);
4851 server.abort();
4852
4853 let chunks = futures_util::stream::iter([
4854 Ok::<_, std::convert::Infallible>(Bytes::from_static(b"12345")),
4855 Ok(Bytes::from_static(b"67890")),
4856 ]);
4857 let body = Arc::new(Mutex::new(Some(Body::from_stream(chunks))));
4858 let body_for_handler = Arc::clone(&body);
4859 let (url, server) = spawn_test_server(Router::new().route(
4860 "/v1/messages",
4861 post(move || {
4862 let body = body_for_handler.lock().unwrap().take().unwrap();
4863 async move {
4864 Response::builder()
4865 .status(StatusCode::OK)
4866 .header("content-type", "application/json")
4867 .body(body)
4868 .unwrap()
4869 }
4870 }),
4871 ))
4872 .await;
4873 let config = DirectClientConfig::default()
4874 .try_with_max_response_bytes(8)
4875 .unwrap();
4876 let client = DirectClient::new(config).unwrap();
4877 let error = client
4878 .execute(committed_direct_request(url, WireFormat::Json, br#"{}"#))
4879 .await
4880 .unwrap_err();
4881 assert_eq!(error.code(), ProxyErrorCode::ResponseBodyLimitExceeded);
4882 server.abort();
4883 }
4884
4885 #[tokio::test]
4886 async fn timeout_and_body_stream_error_retain_io_attempted() {
4887 let (slow_url, slow_server) = spawn_test_server(Router::new().route(
4888 "/v1/messages",
4889 post(|| async {
4890 tokio::time::sleep(Duration::from_millis(100)).await;
4891 ([("content-type", "application/json")], "{}")
4892 }),
4893 ))
4894 .await;
4895 let config = DirectClientConfig::default()
4896 .try_with_timeouts(
4897 Duration::from_millis(20),
4898 Duration::from_millis(80),
4899 Duration::from_millis(30),
4900 )
4901 .unwrap();
4902 let client = DirectClient::new(config).unwrap();
4903 let error = client
4904 .execute(committed_direct_request(
4905 slow_url,
4906 WireFormat::Json,
4907 br#"{}"#,
4908 ))
4909 .await
4910 .unwrap_err();
4911 assert_eq!(error.code(), ProxyErrorCode::TotalTimeout);
4912 slow_server.abort();
4913
4914 let chunks =
4915 futures_util::stream::once(async { Ok::<_, std::io::Error>(Bytes::from_static(b"{")) })
4916 .chain(futures_util::stream::once(async {
4917 tokio::time::sleep(Duration::from_millis(20)).await;
4918 Err(std::io::Error::other("synthetic stream failure"))
4919 }));
4920 let body = Body::from_stream(chunks);
4921 let body = Arc::new(Mutex::new(Some(body)));
4922 let body_for_handler = Arc::clone(&body);
4923 let (url, server) = spawn_test_server(Router::new().route(
4924 "/v1/messages",
4925 post(move || {
4926 let body = body_for_handler.lock().unwrap().take().unwrap();
4927 async move {
4928 Response::builder()
4929 .status(StatusCode::OK)
4930 .header("content-type", "application/json")
4931 .body(body)
4932 .unwrap()
4933 }
4934 }),
4935 ))
4936 .await;
4937 let client = DirectClient::new(DirectClientConfig::default()).unwrap();
4938 let error = client
4939 .execute(committed_direct_request(url, WireFormat::Json, br#"{}"#))
4940 .await
4941 .unwrap_err();
4942 assert_eq!(error.phase(), ProxyErrorPhase::UpstreamBody);
4943 server.abort();
4944 }
4945
4946 #[tokio::test]
4947 async fn codec_limit_direction_pins_code_phase_status_body_and_io() {
4948 let capture = Capture::default();
4949 let (endpoint, server) = spawn_test_server(
4950 Router::new()
4951 .route("/v1/messages", post(capture_ok))
4952 .with_state(capture.clone()),
4953 )
4954 .await;
4955 let state = direct_test_state(
4956 endpoint.clone(),
4957 email_pipeline(),
4958 SessionPolicy::EphemeralSingleRequest,
4959 None,
4960 );
4961
4962 let request_error = direct_proxy_inner(
4963 &state,
4964 direct_test_ingress(
4965 endpoint.clone(),
4966 DirectIngressSession::Ephemeral,
4967 WireFormat::Json,
4968 DIRECT_REQUEST_BODY,
4969 ),
4970 DIRECT_REQUEST_BODY,
4971 &LimitFailureCodec(LimitFailureDirection::Request),
4972 )
4973 .await
4974 .unwrap_err();
4975 assert_eq!(
4976 request_error.code(),
4977 ProxyErrorCode::RequestBodyLimitExceeded
4978 );
4979 assert_eq!(request_error.phase(), ProxyErrorPhase::RequestValidation);
4980 assert_eq!(request_error.http_status(), StatusCode::PAYLOAD_TOO_LARGE);
4981 assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4982 let request_response = direct_error_response(request_error);
4983 assert_eq!(request_response.status(), StatusCode::PAYLOAD_TOO_LARGE);
4984 let request_body = to_bytes(request_response.into_body(), 4 * 1024)
4985 .await
4986 .unwrap();
4987 assert_eq!(
4988 serde_json::from_slice::<serde_json::Value>(&request_body).unwrap(),
4989 serde_json::json!({
4990 "type": "error",
4991 "error": {
4992 "type": "api_error",
4993 "message": "proxy_validation_failed",
4994 "code": "RequestBodyLimitExceeded",
4995 "phase": "RequestValidation"
4996 }
4997 })
4998 );
4999
5000 let response_error = direct_proxy_inner(
5001 &state,
5002 direct_test_ingress(
5003 endpoint,
5004 DirectIngressSession::Ephemeral,
5005 WireFormat::Json,
5006 DIRECT_REQUEST_BODY,
5007 ),
5008 DIRECT_REQUEST_BODY,
5009 &LimitFailureCodec(LimitFailureDirection::Response),
5010 )
5011 .await
5012 .unwrap_err();
5013 assert_eq!(
5014 response_error.code(),
5015 ProxyErrorCode::ResponseBodyLimitExceeded
5016 );
5017 assert_eq!(response_error.phase(), ProxyErrorPhase::ResponseValidation);
5018 assert_eq!(response_error.http_status(), StatusCode::BAD_GATEWAY);
5019 assert_eq!(capture.hits.load(Ordering::SeqCst), 1);
5020 let response = direct_error_response(response_error);
5021 assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
5022 let response_body = to_bytes(response.into_body(), 4 * 1024).await.unwrap();
5023 assert_eq!(
5024 serde_json::from_slice::<serde_json::Value>(&response_body).unwrap(),
5025 serde_json::json!({
5026 "type": "error",
5027 "error": {
5028 "type": "api_error",
5029 "message": "proxy_validation_failed",
5030 "code": "ResponseBodyLimitExceeded",
5031 "phase": "ResponseValidation"
5032 }
5033 })
5034 );
5035 server.abort();
5036 }
5037
5038 #[test]
5039 fn future_non_conflict_session_commit_failure_is_never_a_conflict() {
5040 let generation_conflict =
5041 map_session_commit_failure(SessionCommitFailureKind::GenerationConflict);
5042 assert_eq!(
5043 generation_conflict.code(),
5044 ProxyErrorCode::SessionGenerationConflict
5045 );
5046 assert_eq!(generation_conflict.phase(), ProxyErrorPhase::Session);
5047 assert_eq!(generation_conflict.http_status(), StatusCode::CONFLICT);
5048
5049 let non_conflict = map_session_commit_failure(SessionCommitFailureKind::NonConflict);
5050 assert_eq!(non_conflict.code(), ProxyErrorCode::SessionCommitFailure);
5051 assert_eq!(non_conflict.phase(), ProxyErrorPhase::Session);
5052 assert_eq!(
5053 non_conflict.http_status(),
5054 StatusCode::INTERNAL_SERVER_ERROR
5055 );
5056 assert_ne!(non_conflict.code(), generation_conflict.code());
5057 }
5058
5059 #[tokio::test]
5060 async fn continuity_lease_spans_non_stream_response_without_holding_request_mutex() {
5061 let (request, control, registry, expired_at, server) =
5062 begin_delayed_non_stream(false).await;
5063
5064 assert_eq!(control.hits.load(Ordering::SeqCst), 1);
5065 assert_eq!(registry.stats().total_leases(), 1);
5066 let concurrent_lease = registry
5067 .acquire(synthetic_principal(), NON_STREAM_SESSION_ID, None)
5068 .unwrap();
5069 let concurrent_request_guard = concurrent_lease
5070 .try_lock_request()
5071 .expect("request mutex must be free before upstream I/O");
5072 drop(concurrent_request_guard);
5073 drop(concurrent_lease);
5074 assert_eq!(registry.stats().total_leases(), 1);
5075
5076 registry.cleanup_at(expired_at);
5077 assert_eq!(registry.stats().active_entries(), 1);
5078 assert_eq!(registry.stats().tombstones(), 0);
5079
5080 control.release.notify_one();
5081 let response = request.await.unwrap().unwrap();
5082 assert_eq!(response.status(), StatusCode::OK);
5083 wait_for_lease_count(®istry, 0).await;
5084 registry.cleanup_at(expired_at);
5085 assert_eq!(registry.stats().active_entries(), 0);
5086 assert_eq!(registry.stats().tombstones(), 1);
5087 server.abort();
5088 }
5089
5090 #[tokio::test]
5091 async fn continuity_lease_releases_on_non_stream_error() {
5092 let (request, control, registry, expired_at, server) = begin_delayed_non_stream(true).await;
5093 assert_eq!(registry.stats().total_leases(), 1);
5094 registry.cleanup_at(expired_at);
5095 assert_eq!(registry.stats().active_entries(), 1);
5096
5097 control.release.notify_one();
5098 let error = request.await.unwrap().unwrap_err();
5099 assert_eq!(error.phase(), ProxyErrorPhase::ResponseValidation);
5100 assert_eq!(error.http_status(), StatusCode::BAD_GATEWAY);
5101 wait_for_lease_count(®istry, 0).await;
5102 registry.cleanup_at(expired_at);
5103 assert_eq!(registry.stats().active_entries(), 0);
5104 assert_eq!(registry.stats().tombstones(), 1);
5105 server.abort();
5106 }
5107
5108 #[tokio::test]
5109 async fn continuity_lease_releases_on_non_stream_cancellation() {
5110 let (request, _control, registry, expired_at, server) =
5111 begin_delayed_non_stream(false).await;
5112 assert_eq!(registry.stats().total_leases(), 1);
5113 registry.cleanup_at(expired_at);
5114 assert_eq!(registry.stats().active_entries(), 1);
5115
5116 request.abort();
5117 assert!(request.await.unwrap_err().is_cancelled());
5118 wait_for_lease_count(®istry, 0).await;
5119 registry.cleanup_at(expired_at);
5120 assert_eq!(registry.stats().active_entries(), 0);
5121 assert_eq!(registry.stats().tombstones(), 1);
5122 server.abort();
5123 }
5124
5125 #[tokio::test]
5126 async fn continuity_lease_spans_sse_message_stop_then_becomes_expirable() {
5127 let (response, control, registry, expired_at, server) = begin_controlled_sse(false).await;
5128 assert_eq!(control.hits.load(Ordering::SeqCst), 1);
5129 assert_eq!(registry.stats().total_leases(), 1);
5130 registry.cleanup_at(expired_at);
5131 assert_eq!(registry.stats().active_entries(), 1);
5132 assert_eq!(registry.stats().tombstones(), 0);
5133
5134 control.release.notify_one();
5135 let body = to_bytes(response.into_body(), DEFAULT_MAX_RESPONSE_BYTES)
5136 .await
5137 .unwrap();
5138 assert_eq!(body.as_ref(), DIRECT_SSE_RESPONSE);
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_once_on_sse_error() {
5148 let (response, control, registry, expired_at, server) = begin_controlled_sse(true).await;
5149 assert_eq!(registry.stats().total_leases(), 1);
5150 registry.cleanup_at(expired_at);
5151 assert_eq!(registry.stats().active_entries(), 1);
5152
5153 control.release.notify_one();
5154 let body = to_bytes(response.into_body(), DEFAULT_MAX_RESPONSE_BYTES)
5155 .await
5156 .unwrap();
5157 assert_eq!(body.as_ref(), DIRECT_PROXY_ERROR_FRAME);
5158 wait_for_lease_count(®istry, 0).await;
5159 registry.cleanup_at(expired_at);
5160 assert_eq!(registry.stats().active_entries(), 0);
5161 assert_eq!(registry.stats().tombstones(), 1);
5162 server.abort();
5163 }
5164
5165 #[tokio::test]
5166 async fn continuity_lease_releases_once_on_sse_downstream_cancellation() {
5167 let (response, _control, registry, expired_at, server) = begin_controlled_sse(false).await;
5168 assert_eq!(registry.stats().total_leases(), 1);
5169 registry.cleanup_at(expired_at);
5170 assert_eq!(registry.stats().active_entries(), 1);
5171
5172 drop(response);
5173 wait_for_lease_count(®istry, 0).await;
5174 registry.cleanup_at(expired_at);
5175 assert_eq!(registry.stats().active_entries(), 0);
5176 assert_eq!(registry.stats().tombstones(), 1);
5177 server.abort();
5178 }
5179
5180 #[test]
5181 fn proxy_environment_is_ignored_in_isolated_helper() {
5182 let output = Command::new(std::env::current_exe().unwrap())
5183 .args([
5184 "--exact",
5185 "server::tests::proxy_environment_helper",
5186 "--nocapture",
5187 ])
5188 .env("GAZE_PROXY_ENV_HELPER", "1")
5189 .output()
5190 .unwrap();
5191 assert!(
5192 output.status.success(),
5193 "isolated proxy helper failed: {}",
5194 String::from_utf8_lossy(&output.stderr)
5195 );
5196 }
5197
5198 #[test]
5199 fn proxy_environment_helper() {
5200 if std::env::var_os("GAZE_PROXY_ENV_HELPER").is_none() {
5201 return;
5202 }
5203 tokio::runtime::Runtime::new().unwrap().block_on(async {
5204 let capture = Capture::default();
5205 let (url, server) = spawn_test_server(
5206 Router::new()
5207 .route("/v1/messages", post(capture_ok))
5208 .with_state(capture.clone()),
5209 )
5210 .await;
5211 let trap = TcpListener::bind("127.0.0.1:0").await.unwrap();
5212 let proxy = format!("http://{}", trap.local_addr().unwrap());
5213 for name in [
5214 "HTTP_PROXY",
5215 "HTTPS_PROXY",
5216 "ALL_PROXY",
5217 "http_proxy",
5218 "https_proxy",
5219 ] {
5220 std::env::set_var(name, &proxy);
5221 }
5222 std::env::remove_var("NO_PROXY");
5223 std::env::remove_var("no_proxy");
5224
5225 let client = DirectClient::new(DirectClientConfig::default()).unwrap();
5226 client
5227 .execute(committed_direct_request(url, WireFormat::Json, br#"{}"#))
5228 .await
5229 .unwrap();
5230 assert_eq!(capture.hits.load(Ordering::SeqCst), 1);
5231 assert!(
5232 tokio::time::timeout(Duration::from_millis(100), trap.accept())
5233 .await
5234 .is_err()
5235 );
5236 server.abort();
5237 });
5238 }
5239
5240 #[test]
5241 fn direct_frames_match_anthropic_codec_bytes() {
5242 assert_eq!(
5243 DIRECT_PROXY_PING_FRAME,
5244 crate::codecs::anthropic::ANTHROPIC_PROXY_PING_FRAME
5245 );
5246 assert_eq!(
5247 DIRECT_PROXY_ERROR_FRAME,
5248 crate::codecs::anthropic::ANTHROPIC_PROXY_ERROR_FRAME
5249 );
5250 }
5251
5252 #[test]
5253 fn restore_text_json_escapes_raw_values_only_into_json_syntax() {
5254 let session = Session::new(Scope::Conversation("test".to_string())).unwrap();
5255 let raw = "q\"b\\s\u{1}\u{1f}\u{7f}\u{2028}\u{e9}\r\n";
5256 let mut transaction = session.begin_transaction();
5257 let token = transaction
5258 .tokenize(&PiiClass::Custom("secret".to_string()), raw)
5259 .unwrap();
5260 transaction.commit().unwrap();
5261
5262 let none = HashSet::new();
5263 let restored = restore_text(&session, &none, &format!("{{\"v\":\"x{token}y\"}}"), true);
5264 let parsed: Value = serde_json::from_str(&restored).unwrap();
5265 assert_eq!(parsed["v"].as_str(), Some(format!("x{raw}y").as_str()));
5266 assert_eq!(
5267 restore_text(&session, &none, &format!("x{token}y"), false),
5268 format!("x{raw}y")
5269 );
5270 let json_spelled = HashSet::from([token.clone()]);
5272 let restored = restore_text(
5273 &session,
5274 &json_spelled,
5275 &format!("{{\"v\":\"{token}\"}}"),
5276 true,
5277 );
5278 assert_eq!(
5279 serde_json::from_str::<Value>(&restored).unwrap()["v"].as_str(),
5280 Some(raw)
5281 );
5282 }
5283
5284 #[test]
5285 fn json_spelled_tokens_are_the_ones_minted_inside_json_string_literals() {
5286 let mut tokens = HashSet::new();
5287 collect_json_spelled_tokens(
5288 r#"{"a":"Robert \"Bob\" Smith","b":"<Email_9>"}"#,
5289 r#"{"a":"<Custom:name_1>","<Custom:key_1>":["<Custom:name_2>"],"b":"<Email_9>"}"#,
5290 &mut tokens,
5291 );
5292 assert_eq!(
5293 tokens,
5294 HashSet::from([
5295 "<Custom:name_1>".to_string(),
5296 "<Custom:key_1>".to_string(),
5297 "<Custom:name_2>".to_string(),
5298 ])
5299 );
5300
5301 let mut none = HashSet::new();
5302 collect_json_spelled_tokens("prose", "prose <Custom:name_1>", &mut none);
5303 collect_json_spelled_tokens("{broken", "{broken <Custom:name_1>", &mut none);
5304 assert!(none.is_empty());
5305 }
5306
5307 #[test]
5308 fn restore_text_leaves_unknown_tokens() {
5309 let session = Session::new(Scope::Conversation("test".to_string())).unwrap();
5310 let none = HashSet::new();
5311 assert_eq!(
5312 restore_text(&session, &none, "hello <Email_1>", false),
5313 "hello <Email_1>"
5314 );
5315 assert_eq!(
5316 restore_text(&session, &none, "{\"to\":\"<Email_1>\"}", true),
5317 "{\"to\":\"<Email_1>\"}"
5318 );
5319 }
5320
5321 #[test]
5322 fn upstream_url_preserves_path_and_query() {
5323 let base = Url::parse("https://api.openai.com").unwrap();
5324 let uri: Uri = "/v1/chat/completions?x=1".parse().unwrap();
5325 assert_eq!(
5326 upstream_url(&base, &uri).unwrap().as_str(),
5327 "https://api.openai.com/v1/chat/completions?x=1"
5328 );
5329 }
5330
5331 #[test]
5332 fn health_lists_adapters() {
5333 let config = ProxyConfig::new(
5334 "127.0.0.1:0".parse().unwrap(),
5335 vec![Arc::new(crate::adapters::OpenAiAdapter::new(
5336 Url::parse("https://api.openai.com").unwrap(),
5337 ))],
5338 );
5339 let state = AppState {
5340 config: Arc::new(config),
5341 pipeline: Arc::new(Pipeline::builder().build().unwrap()),
5342 client: Client::new(),
5343 sessions: Arc::new(RwLock::new(HashMap::new())),
5344 direct: None,
5345 started_at: Instant::now(),
5346 };
5347 assert_eq!(health_snapshot(&state).adapters[0].name, "openai");
5348 }
5349
5350 fn span(clean: std::ops::Range<usize>, raw: std::ops::Range<usize>) -> EmittedTokenSpan {
5351 EmittedTokenSpan::new(clean, raw, gaze::PiiClass::Email)
5352 }
5353
5354 #[test]
5355 fn manifest_accounting_accepts_an_untouched_pair_and_a_described_substitution() {
5356 assert!(manifest_accounts_for_every_change(
5357 "no pii here",
5358 "no pii here",
5359 &[]
5360 ));
5361 assert!(manifest_accounts_for_every_change(
5363 "lead a@b tail",
5364 "lead <T> tail",
5365 &[span(5..8, 5..8)],
5366 ));
5367 assert!(manifest_accounts_for_every_change(
5370 "a@b and c@d",
5371 "<TOKEN1> and <TOKEN2>",
5372 &[span(0..8, 0..3), span(13..21, 8..11)],
5373 ));
5374 }
5375
5376 #[test]
5377 fn manifest_accounting_refuses_a_deletion_no_entry_describes() {
5378 assert!(!manifest_accounts_for_every_change(
5381 "lead deleted@example.invalid tail",
5382 "lead tail",
5383 &[],
5384 ));
5385 assert!(!manifest_accounts_for_every_change(
5387 "a@b and deleted@example.invalid",
5388 "<T> and ",
5389 &[span(0..3, 0..3)],
5390 ));
5391 }
5392
5393 #[test]
5394 fn manifest_accounting_refuses_entries_that_cannot_describe_the_pair() {
5395 assert!(!manifest_accounts_for_every_change(
5397 "short",
5398 "short",
5399 &[span(0..5, 0..99)]
5400 ));
5401 assert!(!manifest_accounts_for_every_change(
5403 "a@b c@d",
5404 "<T><T>",
5405 &[span(0..3, 0..4), span(3..6, 2..7)],
5406 ));
5407 assert!(!manifest_accounts_for_every_change(
5409 "a@b tail",
5410 "<T> other",
5411 &[span(0..3, 0..3)]
5412 ));
5413 assert!(!manifest_accounts_for_every_change(
5415 "\u{00e4}x",
5416 "\u{00e4}x",
5417 &[span(0..1, 0..1)]
5418 ));
5419 }
5420
5421 #[test]
5422 fn manifest_accounting_refuses_a_deletion_that_falls_between_two_entries() {
5423 assert!(!manifest_accounts_for_every_change(
5436 "a@b GONE c@d",
5437 "<T1> <T2>",
5438 &[span(0..4, 0..3), span(5..9, 9..12)],
5439 ));
5440 }
5441}