1use runifold_model::{
2 Artifact, ArtifactError, ArtifactFuture, ArtifactPage, ArtifactRef, ArtifactScope,
3 ArtifactStore, ArtifactWrite, DEFAULT_MAX_ARTIFACT_BYTES, MAX_ARTIFACT_PAGE_SIZE,
4};
5use sha2::{Digest, Sha256};
6use std::time::{SystemTime, UNIX_EPOCH};
7
8use super::PostgresConversationStore;
9
10impl ArtifactStore for PostgresConversationStore {
11 fn put(&self, write: ArtifactWrite) -> ArtifactFuture<'_, Result<ArtifactRef, ArtifactError>> {
12 Box::pin(async move {
13 if write.bytes().len() > DEFAULT_MAX_ARTIFACT_BYTES {
14 return Err(ArtifactError::InvalidInput(format!(
15 "artifact is {} bytes and exceeds the {}-byte limit",
16 write.bytes().len(),
17 DEFAULT_MAX_ARTIFACT_BYTES
18 )));
19 }
20 let digest = sha256(write.bytes());
21 let artifact_id = artifact_identity(write.media_type(), write.bytes());
22 let size_bytes = i64::try_from(write.bytes().len()).map_err(|_| {
23 ArtifactError::InvalidInput("artifact size exceeds PostgreSQL BIGINT".into())
24 })?;
25 let created_at_ms = i64_value(unix_time_ms()?, "creation time")?;
26 let expires_at_ms = write
27 .expires_at_unix_ms()
28 .map(|value| i64_value(value, "expiration time"))
29 .transpose()?;
30 let artifacts = format!("{}_artifacts", self.table);
31 let idempotency = format!("{}_artifact_idempotency", self.table);
32 let mut client = self.transaction_client.lock().await;
33 let transaction = client
34 .transaction()
35 .await
36 .map_err(|error| storage(&error))?;
37 if let Some(reference) =
38 load_replay_reference(&transaction, &artifacts, &idempotency, &write, &artifact_id)
39 .await?
40 {
41 return Ok(reference);
42 }
43 transaction
44 .execute(
45 &format!(
46 "INSERT INTO {artifacts}
47 (scope, artifact_id, media_type, size_bytes, sha256, name, bytes,
48 created_at_ms, expires_at_ms)
49 VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
50 ON CONFLICT (scope, artifact_id) DO NOTHING"
51 ),
52 &[
53 &write.scope().as_str(),
54 &artifact_id,
55 &write.media_type(),
56 &size_bytes,
57 &digest,
58 &write.name(),
59 &write.bytes(),
60 &created_at_ms,
61 &expires_at_ms,
62 ],
63 )
64 .await
65 .map_err(|error| storage(&error))?;
66 let reference =
67 load_reference(&transaction, &artifacts, write.scope(), &artifact_id).await?;
68 if !write.matches_immutable_reference(&reference) {
69 return Err(ArtifactError::MetadataConflict(artifact_id));
70 }
71 transaction
72 .execute(
73 &format!(
74 "INSERT INTO {idempotency} (scope, idempotency_key, artifact_id)
75 VALUES ($1, $2, $3)
76 ON CONFLICT (scope, idempotency_key) DO NOTHING"
77 ),
78 &[
79 &write.scope().as_str(),
80 &write.idempotency_key(),
81 &artifact_id,
82 ],
83 )
84 .await
85 .map_err(|error| storage(&error))?;
86 let selected = transaction
87 .query_one(
88 &format!(
89 "SELECT artifact_id FROM {idempotency}
90 WHERE scope = $1 AND idempotency_key = $2"
91 ),
92 &[&write.scope().as_str(), &write.idempotency_key()],
93 )
94 .await
95 .map_err(|error| storage(&error))?
96 .get::<_, String>("artifact_id");
97 if selected != artifact_id {
98 return Err(ArtifactError::IdempotencyConflict(
99 write.idempotency_key().into(),
100 ));
101 }
102 transaction
103 .commit()
104 .await
105 .map_err(|error| storage(&error))?;
106 Ok(reference)
107 })
108 }
109
110 fn get(&self, reference: &ArtifactRef) -> ArtifactFuture<'_, Result<Artifact, ArtifactError>> {
111 let reference = reference.clone();
112 Box::pin(async move {
113 let sql = format!(
114 "SELECT media_type, size_bytes, sha256, name, bytes,
115 created_at_ms, expires_at_ms
116 FROM {}_artifacts WHERE scope = $1 AND artifact_id = $2",
117 self.table
118 );
119 let row = self
120 .client
121 .query_opt(&sql, &[&reference.scope.as_str(), &reference.artifact_id])
122 .await
123 .map_err(|error| storage(&error))?
124 .ok_or_else(|| ArtifactError::NotFound(reference.artifact_id.clone()))?;
125 let size = decode_size(row.get("size_bytes"), &reference.artifact_id)?;
126 let artifact = Artifact {
127 reference: ArtifactRef {
128 scope: reference.scope.clone(),
129 artifact_id: reference.artifact_id.clone(),
130 media_type: row.get("media_type"),
131 size_bytes: size,
132 sha256: row.get("sha256"),
133 name: row.get("name"),
134 created_at_unix_ms: decode_time(
135 row.get("created_at_ms"),
136 &reference.artifact_id,
137 )?,
138 expires_at_unix_ms: row
139 .get::<_, Option<i64>>("expires_at_ms")
140 .map(|value| decode_time(value, &reference.artifact_id))
141 .transpose()?,
142 },
143 bytes: row.get("bytes"),
144 };
145 verify(&artifact)?;
146 if artifact.reference != reference {
147 return Err(ArtifactError::Integrity(reference.artifact_id));
148 }
149 ensure_not_expired(&artifact.reference, unix_time_ms()?)?;
150 Ok(artifact)
151 })
152 }
153
154 fn list(
155 &self,
156 scope: &ArtifactScope,
157 after: Option<&str>,
158 limit: u32,
159 ) -> ArtifactFuture<'_, Result<ArtifactPage, ArtifactError>> {
160 let scope = scope.clone();
161 let after = after.unwrap_or_default().to_owned();
162 Box::pin(async move {
163 validate_limit(limit)?;
164 let sql = format!(
165 "SELECT artifact_id, media_type, size_bytes, sha256, name,
166 created_at_ms, expires_at_ms
167 FROM {}_artifacts
168 WHERE scope = $1 AND artifact_id > $2
169 ORDER BY artifact_id LIMIT $3",
170 self.table
171 );
172 let fetch = i64::from(limit) + 1;
173 let rows = self
174 .client
175 .query(&sql, &[&scope.as_str(), &after, &fetch])
176 .await
177 .map_err(|error| storage(&error))?;
178 let mut items = rows
179 .iter()
180 .map(|row| decode_reference(row, &scope))
181 .collect::<Result<Vec<_>, _>>()?;
182 let next_cursor = if items.len() > limit as usize {
183 items.pop();
184 items.last().map(|item| item.artifact_id.clone())
185 } else {
186 None
187 };
188 Ok(ArtifactPage { items, next_cursor })
189 })
190 }
191
192 fn delete(
193 &self,
194 scope: &ArtifactScope,
195 artifact_id: &str,
196 ) -> ArtifactFuture<'_, Result<bool, ArtifactError>> {
197 let scope = scope.clone();
198 let artifact_id = artifact_id.to_owned();
199 Box::pin(async move {
200 let artifacts = format!("{}_artifacts", self.table);
201 let idempotency = format!("{}_artifact_idempotency", self.table);
202 let mut client = self.transaction_client.lock().await;
203 let transaction = client
204 .transaction()
205 .await
206 .map_err(|error| storage(&error))?;
207 transaction
208 .execute(
209 &format!("DELETE FROM {idempotency} WHERE scope = $1 AND artifact_id = $2"),
210 &[&scope.as_str(), &artifact_id],
211 )
212 .await
213 .map_err(|error| storage(&error))?;
214 let removed = transaction
215 .execute(
216 &format!("DELETE FROM {artifacts} WHERE scope = $1 AND artifact_id = $2"),
217 &[&scope.as_str(), &artifact_id],
218 )
219 .await
220 .map_err(|error| storage(&error))?
221 > 0;
222 transaction
223 .commit()
224 .await
225 .map_err(|error| storage(&error))?;
226 Ok(removed)
227 })
228 }
229
230 fn purge_expired(
231 &self,
232 scope: &ArtifactScope,
233 now_unix_ms: u64,
234 limit: u32,
235 ) -> ArtifactFuture<'_, Result<u32, ArtifactError>> {
236 let scope = scope.clone();
237 Box::pin(async move {
238 validate_limit(limit)?;
239 let now = i64_value(now_unix_ms, "purge time")?;
240 let artifacts = format!("{}_artifacts", self.table);
241 let sql = format!(
242 "DELETE FROM {artifacts} WHERE (scope, artifact_id) IN (
243 SELECT scope, artifact_id FROM {artifacts}
244 WHERE scope = $1 AND expires_at_ms IS NOT NULL AND expires_at_ms <= $2
245 ORDER BY expires_at_ms, artifact_id LIMIT $3
246 )"
247 );
248 let removed = self
249 .client
250 .execute(&sql, &[&scope.as_str(), &now, &i64::from(limit)])
251 .await
252 .map_err(|error| storage(&error))?;
253 u32::try_from(removed)
254 .map_err(|_| ArtifactError::Storage("purge count exceeds u32".into()))
255 })
256 }
257}
258
259async fn load_replay_reference(
260 transaction: &tokio_postgres::Transaction<'_>,
261 artifacts: &str,
262 idempotency: &str,
263 write: &ArtifactWrite,
264 artifact_id: &str,
265) -> Result<Option<ArtifactRef>, ArtifactError> {
266 let existing = transaction
267 .query_opt(
268 &format!(
269 "SELECT artifact_id FROM {idempotency}
270 WHERE scope = $1 AND idempotency_key = $2"
271 ),
272 &[&write.scope().as_str(), &write.idempotency_key()],
273 )
274 .await
275 .map_err(|error| storage(&error))?;
276 let Some(existing) = existing else {
277 return Ok(None);
278 };
279 let existing = existing.get::<_, String>("artifact_id");
280 if existing != artifact_id {
281 return Err(ArtifactError::IdempotencyConflict(
282 write.idempotency_key().into(),
283 ));
284 }
285 let reference = load_reference(transaction, artifacts, write.scope(), &existing).await?;
286 if !write.matches_immutable_reference(&reference) {
287 return Err(ArtifactError::IdempotencyConflict(
288 write.idempotency_key().into(),
289 ));
290 }
291 Ok(Some(reference))
292}
293
294async fn load_reference(
295 transaction: &tokio_postgres::Transaction<'_>,
296 table: &str,
297 scope: &ArtifactScope,
298 artifact_id: &str,
299) -> Result<ArtifactRef, ArtifactError> {
300 let row = transaction
301 .query_one(
302 &format!(
303 "SELECT artifact_id, media_type, size_bytes, sha256, name,
304 created_at_ms, expires_at_ms
305 FROM {table} WHERE scope = $1 AND artifact_id = $2"
306 ),
307 &[&scope.as_str(), &artifact_id],
308 )
309 .await
310 .map_err(|error| storage(&error))?;
311 decode_reference(&row, scope)
312}
313
314fn decode_reference(
315 row: &tokio_postgres::Row,
316 scope: &ArtifactScope,
317) -> Result<ArtifactRef, ArtifactError> {
318 let artifact_id: String = row.get("artifact_id");
319 Ok(ArtifactRef {
320 scope: scope.clone(),
321 artifact_id: artifact_id.clone(),
322 media_type: row.get("media_type"),
323 size_bytes: decode_size(row.get("size_bytes"), &artifact_id)?,
324 sha256: row.get("sha256"),
325 name: row.get("name"),
326 created_at_unix_ms: decode_time(row.get("created_at_ms"), &artifact_id)?,
327 expires_at_unix_ms: row
328 .get::<_, Option<i64>>("expires_at_ms")
329 .map(|value| decode_time(value, &artifact_id))
330 .transpose()?,
331 })
332}
333
334fn verify(artifact: &Artifact) -> Result<(), ArtifactError> {
335 if artifact.reference.size_bytes != artifact.bytes.len() as u64
336 || artifact.reference.sha256 != sha256(&artifact.bytes)
337 || artifact.reference.artifact_id
338 != artifact_identity(&artifact.reference.media_type, &artifact.bytes)
339 {
340 return Err(ArtifactError::Integrity(
341 artifact.reference.artifact_id.clone(),
342 ));
343 }
344 Ok(())
345}
346
347fn decode_size(value: i64, artifact_id: &str) -> Result<u64, ArtifactError> {
348 u64::try_from(value).map_err(|_| ArtifactError::Integrity(artifact_id.into()))
349}
350
351fn decode_time(value: i64, artifact_id: &str) -> Result<u64, ArtifactError> {
352 u64::try_from(value).map_err(|_| ArtifactError::Integrity(artifact_id.into()))
353}
354
355fn i64_value(value: u64, label: &str) -> Result<i64, ArtifactError> {
356 i64::try_from(value)
357 .map_err(|_| ArtifactError::InvalidInput(format!("artifact {label} exceeds i64")))
358}
359
360fn unix_time_ms() -> Result<u64, ArtifactError> {
361 let elapsed = SystemTime::now()
362 .duration_since(UNIX_EPOCH)
363 .map_err(|error| ArtifactError::Storage(error.to_string()))?;
364 u64::try_from(elapsed.as_millis())
365 .map_err(|_| ArtifactError::Storage("system time exceeds u64 milliseconds".into()))
366}
367
368fn ensure_not_expired(reference: &ArtifactRef, now: u64) -> Result<(), ArtifactError> {
369 if reference
370 .expires_at_unix_ms
371 .is_some_and(|expires| expires <= now)
372 {
373 return Err(ArtifactError::Expired(reference.artifact_id.clone()));
374 }
375 Ok(())
376}
377
378fn validate_limit(limit: u32) -> Result<(), ArtifactError> {
379 if limit == 0 || limit > MAX_ARTIFACT_PAGE_SIZE {
380 return Err(ArtifactError::InvalidInput(format!(
381 "artifact page limit must be between 1 and {MAX_ARTIFACT_PAGE_SIZE}"
382 )));
383 }
384 Ok(())
385}
386
387fn sha256(bytes: &[u8]) -> String {
388 hex_digest(Sha256::digest(bytes))
389}
390
391fn artifact_identity(media_type: &str, bytes: &[u8]) -> String {
392 let mut digest = Sha256::new();
393 digest.update(media_type.as_bytes());
394 digest.update([0]);
395 digest.update(bytes);
396 format!("sha256:{}", hex_digest(digest.finalize()))
397}
398
399fn hex_digest(digest: impl AsRef<[u8]>) -> String {
400 const HEX: &[u8; 16] = b"0123456789abcdef";
401 let bytes = digest.as_ref();
402 let mut output = String::with_capacity(bytes.len().saturating_mul(2));
403 for byte in bytes {
404 output.push(char::from(HEX[usize::from(byte >> 4)]));
405 output.push(char::from(HEX[usize::from(byte & 0x0f)]));
406 }
407 output
408}
409
410fn storage(error: &tokio_postgres::Error) -> ArtifactError {
411 ArtifactError::Storage(error.to_string())
412}