1use crate::{Error, ErrorKind, Result};
4use std::{collections::BTreeSet, fmt, future::Future, time::Duration};
5use tokio::time::Instant;
6use tokio_util::sync::CancellationToken;
7use url::Url;
8
9pub(crate) const MAX_URL_BYTES: usize = 16 * 1024;
10const MAX_PATH_SEGMENT_BYTES: usize = 4096;
11
12#[derive(Clone)]
14pub struct Endpoint {
15 url: Url,
16 origin: String,
17 loopback_http: bool,
18}
19
20impl Endpoint {
21 pub fn parse(value: &str) -> Result<Self> {
22 Self::parse_inner(value, false)
23 }
24 pub fn loopback(value: &str) -> Result<Self> {
26 Self::parse_inner(value, true)
27 }
28 fn parse_inner(value: &str, local: bool) -> Result<Self> {
29 if value.len() > MAX_URL_BYTES {
30 return Err(Error::new(ErrorKind::Configuration));
31 }
32 let url = Url::parse(value).map_err(|_| Error::new(ErrorKind::Configuration))?;
33 if !valid_url(&url, local) || url.query().is_some() {
34 return Err(Error::new(ErrorKind::Configuration));
35 }
36 let origin = url.origin().ascii_serialization();
37 Ok(Self {
38 url,
39 origin,
40 loopback_http: local,
41 })
42 }
43 pub fn url(&self) -> &Url {
45 &self.url
46 }
47 pub fn origin(&self) -> &str {
48 &self.origin
49 }
50 pub fn admit(&self, url: &Url) -> Result<()> {
52 if !valid_url(url, self.loopback_http) || url.origin().ascii_serialization() != self.origin
53 {
54 return Err(Error::new(ErrorKind::Destination));
55 }
56 Ok(())
57 }
58 pub(crate) fn resolve(
59 &self,
60 path: &str,
61 segments: &[String],
62 query: &[(String, String)],
63 ) -> Result<Url> {
64 if path.len() > MAX_URL_BYTES
65 || query.len() > 256
66 || query
67 .iter()
68 .try_fold(0usize, |n, (k, v)| {
69 n.checked_add(k.len())?.checked_add(v.len())
70 })
71 .is_none_or(|n| n > MAX_URL_BYTES)
72 {
73 return Err(Error::new(ErrorKind::Limit));
74 }
75 if path.starts_with("//") || path.contains(['?', '#', '\\']) || path.contains("://") {
76 return Err(Error::new(ErrorKind::Validation));
77 }
78 for part in path.split('/') {
79 let lowered = part.to_ascii_lowercase();
82 if ["%2f", "%5c", "%25"]
83 .iter()
84 .any(|value| lowered.contains(value))
85 {
86 return Err(Error::new(ErrorKind::Validation));
87 }
88 let normalized = lowered.replace("%2e", ".");
89 if matches!(normalized.as_str(), "." | "..") {
90 return Err(Error::new(ErrorKind::Validation));
91 }
92 }
93 let mut url = self.url.clone();
94 let mut full = format!(
95 "{}/{}",
96 url.path().trim_end_matches('/'),
97 path.trim_start_matches('/')
98 );
99 for segment in segments {
100 full.push('/');
101 full.push_str(&encode_segment(segment)?);
102 if full.len() > MAX_URL_BYTES {
103 return Err(Error::new(ErrorKind::Limit));
104 }
105 }
106 url.set_path(&full);
107 if !query.is_empty() {
108 url.query_pairs_mut()
109 .extend_pairs(query.iter().map(|(k, v)| (k, v)));
110 }
111 self.admit(&url)?;
112 if url.as_str().len() > MAX_URL_BYTES {
113 return Err(Error::new(ErrorKind::Limit));
114 }
115 Ok(url)
116 }
117}
118
119pub(crate) fn encode_segment(value: &str) -> Result<String> {
122 if value.is_empty() || value.len() > MAX_PATH_SEGMENT_BYTES {
123 return Err(Error::new(ErrorKind::Validation));
124 }
125 let mut decoded = value.as_bytes().to_vec();
126 for depth in 0..=8 {
127 if decoded.iter().any(|b| b.is_ascii_control() || *b == b'\\')
128 || decoded
129 .split(|b| *b == b'/')
130 .any(|p| p == b"." || p == b"..")
131 {
132 return Err(Error::new(ErrorKind::Validation));
133 }
134 let mut next = Vec::with_capacity(decoded.len());
135 let mut i = 0;
136 while i < decoded.len() {
137 if decoded[i] == b'%'
138 && i + 2 < decoded.len()
139 && let (Some(a), Some(b)) = (hex(decoded[i + 1]), hex(decoded[i + 2]))
140 {
141 next.push(a * 16 + b);
142 i += 3;
143 } else {
144 next.push(decoded[i]);
145 i += 1;
146 }
147 }
148 if next == decoded {
149 break;
150 }
151 if depth == 8 {
152 return Err(Error::new(ErrorKind::Validation));
153 }
154 decoded = next;
155 }
156 const HEX: &[u8; 16] = b"0123456789ABCDEF";
157 let mut encoded = String::with_capacity(value.len());
158 for byte in value.bytes() {
159 if byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.' | b'~') {
160 encoded.push(byte as char);
161 } else {
162 encoded.push('%');
163 encoded.push(HEX[(byte >> 4) as usize] as char);
164 encoded.push(HEX[(byte & 15) as usize] as char);
165 }
166 }
167 Ok(encoded)
168}
169fn hex(value: u8) -> Option<u8> {
170 match value {
171 b'0'..=b'9' => Some(value - b'0'),
172 b'a'..=b'f' => Some(value - b'a' + 10),
173 b'A'..=b'F' => Some(value - b'A' + 10),
174 _ => None,
175 }
176}
177fn valid_url(url: &Url, local: bool) -> bool {
178 if url.host().is_none()
179 || !url.username().is_empty()
180 || url.password().is_some()
181 || url.fragment().is_some()
182 {
183 return false;
184 }
185 if url.scheme() == "https" {
186 return true;
187 }
188 if !local || url.scheme() != "http" {
189 return false;
190 }
191 match url.host() {
192 Some(url::Host::Domain(name)) => name == "localhost",
193 Some(url::Host::Ipv4(ip)) => ip.is_loopback(),
194 Some(url::Host::Ipv6(ip)) => ip.is_loopback(),
195 None => false,
196 }
197}
198impl fmt::Debug for Endpoint {
199 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
200 f.debug_struct("Endpoint")
201 .field("origin", &self.origin)
202 .finish_non_exhaustive()
203 }
204}
205
206#[derive(Debug, Clone, Copy, PartialEq, Eq)]
208#[non_exhaustive]
209pub enum DestinationPurpose {
210 Api,
211 OAuthToken,
212 Download,
213 Realtime,
214}
215
216pub trait DestinationPolicy: Send + Sync {
220 fn check(&self, destination: &Url, purpose: DestinationPurpose) -> Result<()>;
221}
222
223#[derive(Clone, Debug)]
225pub struct OriginPolicy {
226 origins: BTreeSet<String>,
227}
228impl OriginPolicy {
229 pub fn new(endpoints: impl IntoIterator<Item = Endpoint>) -> Self {
230 Self {
231 origins: endpoints.into_iter().map(|e| e.origin).collect(),
232 }
233 }
234}
235impl DestinationPolicy for OriginPolicy {
236 fn check(&self, url: &Url, purpose: DestinationPurpose) -> Result<()> {
237 let mut admitted = url.clone();
238 if purpose == DestinationPurpose::Realtime {
240 match url.scheme() {
241 "wss" => {
242 admitted
243 .set_scheme("https")
244 .map_err(|_| Error::new(ErrorKind::Destination))?;
245 }
246 "ws" => {
247 admitted
248 .set_scheme("http")
249 .map_err(|_| Error::new(ErrorKind::Destination))?;
250 }
251 _ => return Err(Error::new(ErrorKind::Destination)),
252 }
253 }
254 if self
255 .origins
256 .contains(&admitted.origin().ascii_serialization())
257 {
258 Ok(())
259 } else {
260 Err(Error::new(ErrorKind::Destination))
261 }
262 }
263}
264
265#[derive(Clone, Debug)]
270#[non_exhaustive]
271pub struct Limits {
272 pub operation_timeout: Duration,
273 pub max_request_bytes: usize,
274 pub max_response_bytes: usize,
275 pub max_error_bytes: usize,
276 pub max_concurrency: usize,
277 pub max_retries: u32,
278 pub retry_base_delay: Duration,
279 pub max_retry_delay: Duration,
280 pub stream_handshake_timeout: Duration,
281 pub stream_idle_timeout: Duration,
282 pub stream_session_timeout: Option<Duration>,
283 pub max_event_bytes: usize,
284}
285impl Default for Limits {
286 fn default() -> Self {
287 Self {
288 operation_timeout: Duration::from_secs(60),
289 max_request_bytes: 8 * 1024 * 1024,
290 max_response_bytes: 8 * 1024 * 1024,
291 max_error_bytes: 64 * 1024,
292 max_concurrency: 32,
293 max_retries: 2,
294 retry_base_delay: Duration::from_millis(250),
295 max_retry_delay: Duration::from_secs(5),
296 stream_handshake_timeout: Duration::from_secs(30),
297 stream_idle_timeout: Duration::from_secs(30),
298 stream_session_timeout: None,
299 max_event_bytes: 1024 * 1024,
300 }
301 }
302}
303impl Limits {
304 pub fn validate(&self) -> Result<()> {
305 let durations = [
306 self.operation_timeout,
307 self.retry_base_delay,
308 self.max_retry_delay,
309 self.stream_handshake_timeout,
310 self.stream_idle_timeout,
311 ];
312 let byte_limits = [
313 self.max_request_bytes,
314 self.max_response_bytes,
315 self.max_error_bytes,
316 self.max_event_bytes,
317 ];
318 if durations
319 .iter()
320 .any(|v| v.is_zero() || *v > Duration::from_secs(86400))
321 || self
322 .stream_session_timeout
323 .is_some_and(|v| v.is_zero() || v > Duration::from_secs(86400))
324 || byte_limits
325 .iter()
326 .any(|v| *v == 0 || *v > 512 * 1024 * 1024)
327 || self.max_error_bytes > 1024 * 1024
328 || self.max_event_bytes > 16 * 1024 * 1024
329 || !(1..=4096).contains(&self.max_concurrency)
330 || self.max_retries > 5
331 || self.retry_base_delay > self.max_retry_delay
332 {
333 return Err(Error::new(ErrorKind::Configuration));
334 }
335 Ok(())
336 }
337}
338
339#[derive(Clone, Debug, Default)]
341pub struct Session {
342 cancellation: CancellationToken,
343}
344impl Session {
345 pub fn new() -> Self {
346 Self::default()
347 }
348 pub fn cancel(&self) {
349 self.cancellation.cancel();
350 }
351 pub fn is_cancelled(&self) -> bool {
352 self.cancellation.is_cancelled()
353 }
354 pub async fn cancelled(&self) {
355 self.cancellation.cancelled().await;
356 }
357}
358
359#[derive(Clone, Debug)]
361pub struct RequestOptions {
362 deadline: Instant,
363 cancellation: CancellationToken,
364 retry_limit: Option<u32>,
365 response_defaults: bool,
366}
367impl RequestOptions {
368 pub fn with_timeout(timeout: Duration) -> Result<Self> {
369 if timeout.is_zero() || timeout > Duration::from_secs(86400) {
370 return Err(Error::new(ErrorKind::Configuration));
371 }
372 let deadline = Instant::now()
373 .checked_add(timeout)
374 .ok_or_else(|| Error::new(ErrorKind::Configuration))?;
375 Ok(Self {
376 deadline,
377 cancellation: CancellationToken::new(),
378 retry_limit: None,
379 response_defaults: true,
380 })
381 }
382 pub fn deadline(&self) -> Instant {
383 self.deadline
384 }
385 pub fn with_retries(mut self, value: u32) -> Result<Self> {
388 if value > 5 {
389 return Err(Error::new(ErrorKind::Configuration));
390 }
391 self.retry_limit = Some(value);
392 Ok(self)
393 }
394 pub fn retry_limit(&self) -> Option<u32> {
395 self.retry_limit
396 }
397 pub fn without_response_defaults(mut self) -> Self {
407 self.response_defaults = false;
408 self
409 }
410 pub fn response_defaults(&self) -> bool {
412 self.response_defaults
413 }
414 #[cfg(any(
415 feature = "scan",
416 feature = "management",
417 feature = "gateway",
418 feature = "dlp",
419 feature = "iam",
420 feature = "model-security",
421 feature = "red-team",
422 feature = "agentguard"
423 ))]
424 pub(crate) fn response_mode(&self) -> crate::validation::Mode {
425 if self.response_defaults {
426 crate::validation::Mode::Response
427 } else {
428 crate::validation::Mode::ResponseWithoutDefaults
429 }
430 }
431 pub fn cancel(&self) {
432 self.cancellation.cancel();
433 }
434 pub fn is_cancelled(&self) -> bool {
435 self.cancellation.is_cancelled()
436 }
437 pub fn cancellation_token(&self) -> CancellationToken {
438 self.cancellation.clone()
439 }
440 pub async fn run<T>(
442 &self,
443 session: &Session,
444 future: impl Future<Output = Result<T>>,
445 ) -> Result<T> {
446 if self.is_cancelled() || session.is_cancelled() {
447 return Err(Error::new(ErrorKind::Cancelled));
448 }
449 if Instant::now() >= self.deadline {
450 return Err(Error::new(ErrorKind::Deadline));
451 }
452 tokio::select! { biased;
453 _ = self.cancellation.cancelled() => Err(Error::new(ErrorKind::Cancelled)),
454 _ = session.cancelled() => Err(Error::new(ErrorKind::Cancelled)),
455 _ = tokio::time::sleep_until(self.deadline) => Err(Error::new(ErrorKind::Deadline)),
456 value = future => {
457 if self.is_cancelled() || session.is_cancelled() { Err(Error::new(ErrorKind::Cancelled)) }
458 else if Instant::now() >= self.deadline { Err(Error::new(ErrorKind::Deadline)) }
459 else { value }
460 },
461 }
462 }
463}