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