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