gproxy_protocol/adapt/generate/image_resources/
publish.rs1use super::*;
2use crate::{
3 HttpBody,
4 capability::{
5 PublicationKind, PublicationStatus, ResourceAccess, ResourceMetadata, ResourceReference,
6 },
7 transform::{identity::IdNamespace, images::decode_image},
8 wire::DeclaredFields,
9};
10
11impl<H> GeminiImagePublications<H> {
12 pub async fn publish<R: ResourceAccess<PublishedHandle = H>>(
15 &mut self,
16 input: &g::GenerateContentResponseBody,
17 namespace: IdNamespace,
18 expires_at: std::time::SystemTime,
19 resources: &super::super::GenerationResources<'_, R>,
20 ) -> Result<g::GenerateContentResponseBody, TransformError> {
21 self.publish_event(input, namespace, 0, expires_at, resources)
22 .await
23 }
24 pub(crate) async fn publish_event<R: ResourceAccess<PublishedHandle = H>>(
25 &mut self,
26 input: &g::GenerateContentResponseBody,
27 namespace: IdNamespace,
28 event_index: u64,
29 expires_at: std::time::SystemTime,
30 resources: &super::super::GenerationResources<'_, R>,
31 ) -> Result<g::GenerateContentResponseBody, TransformError> {
32 if self.failed {
33 return Err(conflict("image publication previously failed permanently"));
34 }
35 if expires_at <= resources.now {
36 return Err(invalid("future image publication expiry required"));
37 }
38 let input = input.clone().into_declared();
39 crate::codec::encode_json(&input, resources.limits).map_err(codec_error)?;
40 if let Some(original) = &self.original {
41 if original != &input {
42 return Err(conflict(
43 "client image response changed during publication recovery",
44 ));
45 }
46 } else {
47 self.original = Some(input.clone());
48 }
49 let mut validated = Vec::new();
50 let mut total = 0u64;
51 for (candidate, value) in input.candidates.iter().flatten().enumerate() {
52 for (index, part) in value
53 .content
54 .iter()
55 .flat_map(|c| c.parts.iter().flatten())
56 .enumerate()
57 {
58 if part.file_data.is_some() {
59 return Err(invalid(
60 "URI publication requires actual inline image bytes",
61 ));
62 }
63 let Some(blob) = &part.inline_data else {
64 continue;
65 };
66 if part.thought_signature.is_some() {
67 return Err(invalid(
68 "cannot replace a signed inline Part with a published fileData Part",
69 ));
70 }
71 let image = decode_image(
72 &blob.data,
73 Some(&blob.mime_type),
74 resources.limits.max_part_bytes,
75 )?;
76 total = total
77 .checked_add(image.bytes.len() as u64)
78 .ok_or_else(limit)?;
79 if image.bytes.len() as u64 > resources.access.limits().write_bytes
80 || total > resources.limits.max_body_bytes
81 || validated.len() >= resources.max_references
82 {
83 return Err(limit());
84 }
85 validated.push(((candidate, index), blob.clone(), image));
86 }
87 }
88 let mut output = input;
89 for (key, blob, image) in validated {
90 let operation = format!(
91 "generation-image:{}:{}:{}:{}",
92 namespace.hex(),
93 event_index,
94 key.0,
95 key.1
96 );
97 let metadata = ResourceMetadata {
98 mime: Some(image.metadata.mime().into()),
99 length: Some(image.bytes.len() as u64),
100 filename: None,
101 expires_at: Some(expires_at),
102 };
103 let recovering = self.publications.contains_key(&key);
104 if let Some(known) = self.publications.get(&key) {
105 if known.source != blob
106 || known.operation != operation
107 || known.expiry != expires_at
108 {
109 return Err(conflict("publication identity or source image changed"));
110 }
111 } else {
112 self.publications.insert(
113 key,
114 Publication {
115 source: blob,
116 operation: operation.clone(),
117 expiry: expires_at,
118 receipt: None,
119 },
120 );
121 let receipt = resources
124 .access
125 .publish(
126 resources.scope,
127 &operation,
128 PublicationKind::Url,
129 metadata.clone(),
130 HttpBody::Bytes(image.bytes),
131 )
132 .await?;
133 let index = self.receipts.len();
134 self.receipts.push(receipt);
135 self.publications
136 .get_mut(&key)
137 .expect("reserved publication")
138 .receipt = Some(index);
139 }
140 if recovering {
141 let receipt = match resources
142 .access
143 .publication_status(resources.scope, &operation)
144 .await?
145 {
146 PublicationStatus::Published(receipt) => receipt,
147 PublicationStatus::Missing => {
148 return Err(missing(
149 "unacknowledged publication has no durable result; body cannot be replayed",
150 ));
151 }
152 PublicationStatus::Pending => {
153 return Err(conflict("image publication is still pending"));
154 }
155 PublicationStatus::Expired => return Err(missing("image publication expired")),
156 };
157 if let Some(index) = self.publications[&key].receipt {
158 let old = &self.receipts[index];
159 if old.reference != receipt.reference || old.metadata != receipt.metadata {
160 self.receipts.push(receipt);
161 self.failed = true;
162 return Err(conflict(
163 "acknowledged publication changed in supplied resource scope",
164 ));
165 }
166 } else {
167 let index = self.receipts.len();
168 self.receipts.push(receipt);
169 self.publications
170 .get_mut(&key)
171 .expect("reserved publication")
172 .receipt = Some(index);
173 }
174 }
175 let receipt = &self.receipts[self.publications[&key]
176 .receipt
177 .expect("acknowledged publication")];
178 let ResourceReference::Url(url) = &receipt.reference else {
179 self.failed = true;
180 return Err(invalid("image publisher returned an ID instead of URL"));
181 };
182 if url.trim().is_empty()
183 || receipt.metadata.mime != metadata.mime
184 || receipt.metadata.length != metadata.length
185 || receipt.metadata.expires_at != metadata.expires_at
186 {
187 self.failed = true;
188 return Err(invalid(
189 "image publication metadata contradicts requested bytes or expiry",
190 ));
191 }
192 let part = output.candidates.as_mut().expect("candidate")[key.0]
193 .content
194 .as_mut()
195 .and_then(|content| content.parts.as_mut())
196 .expect("parts")
197 .get_mut(key.1)
198 .expect("part");
199 part.inline_data = None;
200 part.file_data = Some(
201 g::FileData::builder(url.clone())
202 .mime_type(image.metadata.mime())
203 .build(),
204 );
205 }
206 crate::codec::encode_json(&output, resources.limits).map_err(codec_error)?;
207 Ok(output)
208 }
209}