gproxy_protocol/adapt/generate/resources/
mod.rs1mod chat;
5mod claude;
6mod gemini;
7mod responses;
8use crate::{
9 capability::{ResourceAccess, ResourceReference},
10 codec::{self, CodecLimits},
11 transform::{TransformError, TransformErrorKind},
12};
13use base64::{Engine as _, engine::general_purpose::STANDARD};
14use std::time::SystemTime;
15
16pub struct GenerationResources<'a, R: ResourceAccess> {
17 pub access: &'a R,
18 pub scope: &'a R::Scope,
19 pub limits: CodecLimits,
20 pub max_references: usize,
21 pub now: SystemTime,
22}
23
24pub(super) struct Budget<'a, 'b, R: ResourceAccess> {
25 ctx: &'a GenerationResources<'b, R>,
26 remaining: u64,
27 references: usize,
28}
29
30struct Media {
31 bytes: bytes::Bytes,
32 mime: String,
33 filename: Option<String>,
34}
35
36impl Media {
37 fn data_uri(&self) -> String {
38 format!("data:{};base64,{}", self.mime, STANDARD.encode(&self.bytes))
39 }
40}
41
42impl<R: ResourceAccess> GenerationResources<'_, R> {
43 fn budget(&self) -> Budget<'_, '_, R> {
44 Budget {
45 ctx: self,
46 remaining: self
47 .limits
48 .max_body_bytes
49 .min(self.access.limits().read_bytes),
50 references: 0,
51 }
52 }
53}
54
55impl<R: ResourceAccess> Budget<'_, '_, R> {
56 async fn read(
57 &mut self,
58 reference: ResourceReference,
59 image: bool,
60 ) -> Result<Media, TransformError> {
61 if self.references >= self.ctx.max_references || self.remaining == 0 {
62 return Err(limit());
63 }
64 self.references += 1;
65 let result = self.ctx.access.read(self.ctx.scope, &reference).await?;
66 if result
67 .metadata
68 .length
69 .is_some_and(|length| length > self.remaining)
70 {
71 return Err(limit());
72 }
73 if result
74 .metadata
75 .expires_at
76 .is_some_and(|expiry| expiry <= self.ctx.now)
77 {
78 return Err(TransformError::invalid_result(
79 "resource.expiry",
80 "resource is expired",
81 ));
82 }
83 let mime = result
84 .metadata
85 .mime
86 .ok_or_else(|| TransformError::missing_metadata("resource.mime"))?;
87 if mime.trim().is_empty() || mime.contains([';', ',', '\r', '\n']) {
88 return Err(TransformError::invalid_result(
89 "resource.mime",
90 "exact MIME type required",
91 ));
92 }
93 let mut limits = self.ctx.limits;
94 limits.max_body_bytes = limits.max_body_bytes.min(self.remaining);
95 limits.max_buffer_bytes = limits.max_buffer_bytes.min(self.remaining);
96 let bytes = codec::read_http_body(result.body, limits)
97 .await
98 .map_err(|e| {
99 TransformError::new(
100 if e.kind() == codec::CodecErrorKind::Limit {
101 TransformErrorKind::Limit
102 } else {
103 TransformErrorKind::Host
104 },
105 "resource.body",
106 e.to_string(),
107 )
108 })?;
109 if result
110 .metadata
111 .length
112 .is_some_and(|n| n != bytes.len() as u64)
113 {
114 return Err(TransformError::invalid_result(
115 "resource.length",
116 "declared length differs from actual bytes",
117 ));
118 }
119 self.remaining = self
120 .remaining
121 .checked_sub(bytes.len() as u64)
122 .ok_or_else(limit)?;
123 if (image || matches!(mime.as_str(), "image/png" | "image/jpeg" | "image/webp"))
124 && crate::transform::images::inspect_image(&bytes)?.mime() != mime
125 {
126 return Err(TransformError::invalid_result(
127 "resource.mime",
128 "image bytes differ from MIME",
129 ));
130 }
131 Ok(Media {
132 bytes,
133 mime,
134 filename: result.metadata.filename,
135 })
136 }
137}
138
139fn limit() -> TransformError {
140 TransformError::new(
141 TransformErrorKind::Limit,
142 "generation.resources",
143 "resource count or aggregate byte budget exceeded",
144 )
145}