gproxy_protocol/adapt/guardian/
resources.rs1use crate::{
6 capability::{ResourceAccess, ResourceReference},
7 codec::{self, CodecLimits},
8 transform::{
9 Converted, TransformError, TransformErrorKind,
10 guardian::{self, GuardianPreparedRequest, GuardianRequestContext, GuardianTarget},
11 },
12 wire::openai::guardian::GuardianRequestBody,
13};
14use base64::{Engine as _, engine::general_purpose::STANDARD};
15use std::collections::HashMap;
16
17#[derive(Debug, Clone, Copy)]
18pub struct GuardianResourceLimits {
19 pub codec: CodecLimits,
20 pub max_media: usize,
21 pub max_resource_bytes: u64,
22 pub max_total_resource_bytes: u64,
23}
24
25fn limit(field: &str) -> TransformError {
26 TransformError::new(
27 TransformErrorKind::Limit,
28 field,
29 "Guardian media exceeds configured cap",
30 )
31}
32
33pub async fn prepare_with_resources<A: ResourceAccess>(
37 access: &A,
38 scope: &A::Scope,
39 input: GuardianRequestBody,
40 target: GuardianTarget,
41 context: GuardianRequestContext,
42 limits: GuardianResourceLimits,
43) -> Result<Converted<GuardianPreparedRequest>, TransformError> {
44 let (mut prepared, mut media) =
45 guardian::prepare_unattached(input, target, context, limits.codec)?;
46 if media.len() > limits.max_media {
47 return Err(limit("guardian.media.count"));
48 }
49 let mut total = 0u64;
50 for part in &media {
51 if let Some((_, data)) = part
52 .url()
53 .strip_prefix("data:")
54 .and_then(|v| v.split_once(";base64,"))
55 {
56 let bytes = STANDARD
57 .decode(data)
58 .map_err(|e| TransformError::shape("guardian.media", e.to_string()))?;
59 let len = bytes.len() as u64;
60 if len > limits.max_resource_bytes {
61 return Err(limit("guardian.media.bytes"));
62 }
63 total = total
64 .checked_add(len)
65 .ok_or_else(|| limit("guardian.media.total"))?;
66 }
67 }
68 if total > limits.max_total_resource_bytes {
69 return Err(limit("guardian.media.total"));
70 }
71 let mut resolved: HashMap<String, (String, u64)> = HashMap::new();
74 for part in &mut media {
75 if !part.needs_resolution(target) {
76 continue;
77 }
78 let original = part.url().to_owned();
79 if let Some((uri, len)) = resolved.get(&original) {
80 total = total
81 .checked_add(*len)
82 .ok_or_else(|| limit("guardian.media.total"))?;
83 if total > limits.max_total_resource_bytes {
84 return Err(limit("guardian.media.total"));
85 }
86 part.set_url(uri.clone());
87 continue;
88 }
89 let cap = limits
90 .max_resource_bytes
91 .min(limits.max_total_resource_bytes - total)
92 .min(access.limits().read_bytes)
93 .min(limits.codec.max_body_bytes);
94 if cap == 0 {
95 return Err(limit("guardian.media.read"));
96 }
97 let read = access
98 .read(scope, &ResourceReference::Url(original.clone()))
99 .await
100 .map_err(TransformError::from)?;
101 if read.metadata.length.is_some_and(|n| n > cap) {
102 return Err(limit("guardian.media.length"));
103 }
104 let mime = read
105 .metadata
106 .mime
107 .ok_or_else(|| TransformError::missing_metadata("guardian.media.mime"))?;
108 let bytes = codec::read_http_body(
109 read.body,
110 CodecLimits {
111 max_body_bytes: cap,
112 ..limits.codec
113 },
114 )
115 .await
116 .map_err(|e| {
117 TransformError::new(
118 if e.kind() == codec::CodecErrorKind::Limit {
119 TransformErrorKind::Limit
120 } else {
121 TransformErrorKind::Host
122 },
123 "guardian.media.read",
124 e.to_string(),
125 )
126 })?;
127 let len = bytes.len() as u64;
128 if read.metadata.length.is_some_and(|n| n != len) {
129 return Err(TransformError::invalid_result(
130 "guardian.media.length",
131 "host metadata differs from actual bytes",
132 ));
133 }
134
135 let uri = format!("data:{mime};base64,{}", STANDARD.encode(&bytes));
136 total += len;
137 resolved.insert(original, (uri.clone(), len));
138 part.set_url(uri);
139 }
140 guardian::finish_attachments(&mut prepared.value, &media, limits.codec)?;
141 if !resolved.is_empty() {
142 prepared.report.changed(
143 "guardian.media",
144 "scoped ResourceAccess reads materialized native inline attachments",
145 );
146 }
147 Ok(prepared)
148}