1use std::collections::HashSet;
2use std::path::Path;
3
4use anyhow::{Context, Result, bail};
5use sqlx::{Row, SqliteConnection};
6
7use super::wire::{AttachmentAddPayload, BlobUploadContract, ChangeWire};
8use crate::attachments::decode::validate_image;
9use crate::attachments::lifecycle::{
10 ByteCount, LifecyclePolicy, PruneSummary, SystemClock, acquire_lease, ensure_local_capacity,
11 prune, release_lease, release_reservation, reserve_upload,
12};
13use crate::attachments::storage::{
14 blob_inventory_row, object_path, sha256_hex, store_validated_blob,
15};
16use crate::change_log::op_type;
17use crate::db::Database;
18
19#[derive(Debug, Clone)]
20pub(super) struct MissingLocalBlob {
21 pub sha256: String,
22 pub byte_size: i64,
23 pub media_type: String,
24 pub width: Option<i64>,
25 pub height: Option<i64>,
26}
27
28#[derive(Debug)]
29pub(super) struct MissingBlobPage {
30 pub blobs: Vec<MissingLocalBlob>,
31 pub total: ByteCount,
32}
33
34#[derive(Debug)]
35pub(super) struct PreparedBlobUpload {
36 pub bytes: Vec<u8>,
37 pub lease_id: String,
38}
39
40#[derive(Debug)]
41pub struct ServerBlobDownload {
42 pub bytes: Vec<u8>,
43}
44
45pub(super) fn attachment_blob_contract(change: &ChangeWire) -> Result<Option<BlobUploadContract>> {
46 if change.op_type != op_type::ATTACHMENT_ADD {
47 return Ok(None);
48 }
49 let payload = AttachmentAddPayload::from_change(change)?;
50 Ok(Some(BlobUploadContract {
51 workspace_id: payload.workspace_id,
52 sha256: payload.sha256,
53 byte_size: payload.byte_size,
54 media_type: payload.media_type,
55 width: payload
56 .width
57 .context("validated attachment payload missing width")?,
58 height: payload
59 .height
60 .context("validated attachment payload missing height")?,
61 }))
62}
63
64pub(super) fn unique_blob_contracts(changes: &[ChangeWire]) -> Result<Vec<BlobUploadContract>> {
65 let mut seen = HashSet::new();
66 let mut contracts = Vec::new();
67 for change in changes {
68 if let Some(contract) = attachment_blob_contract(change)?
69 && seen.insert((contract.workspace_id.clone(), contract.sha256.clone()))
70 {
71 contracts.push(contract);
72 }
73 }
74 Ok(contracts)
75}
76
77pub(super) fn contract_by_hash(
78 contracts: &[BlobUploadContract],
79 sha256: &str,
80) -> Result<BlobUploadContract> {
81 contracts
82 .iter()
83 .find(|contract| contract.sha256 == sha256)
84 .cloned()
85 .context("payload missing attachment blob contract")
86}
87
88impl Database {
89 pub(super) async fn prepare_blob_upload(
90 &self,
91 blob_dir: &Path,
92 contract: &BlobUploadContract,
93 ) -> Result<PreparedBlobUpload> {
94 let (row, lease_id) = {
95 let mut conn = self.acquire_writer().await?;
96 let row = blob_inventory_row(&mut conn, &contract.sha256)
97 .await?
98 .filter(|row| row.available)
99 .context("error attachment-blob-local-missing")?;
100 let lease_id =
101 acquire_lease(&mut conn, &contract.sha256, "transfer", &SystemClock).await?;
102 (row, lease_id)
103 };
104 let result = async {
105 let bytes = tokio::fs::read(object_path(blob_dir, &contract.sha256)?)
106 .await
107 .context("error attachment-blob-local-missing")?;
108 if sha256_hex(&bytes) != contract.sha256
109 || row.byte_size
110 != i64::try_from(bytes.len()).context("attachment bytes exceed i64")?
111 {
112 bail!("error attachment-blob-local-invalid");
113 }
114 let validated = validate_image(bytes.clone(), Some(row.media_type.clone())).await?;
115 if (validated.facts.width, validated.facts.height) != (contract.width, contract.height)
116 || validated.facts.media_type != contract.media_type
117 || contract.byte_size != row.byte_size
118 {
119 bail!("error attachment-blob-local-invalid");
120 }
121 Ok(PreparedBlobUpload {
122 bytes,
123 lease_id: lease_id.clone(),
124 })
125 }
126 .await;
127 if result.is_err() {
128 let _ = self.finish_blob_upload(&lease_id).await;
129 }
130 result
131 }
132
133 pub(super) async fn finish_blob_upload(&self, lease_id: &str) -> Result<()> {
134 let mut conn = self.acquire_writer().await?;
135 release_lease(&mut conn, lease_id).await
136 }
137
138 pub(super) async fn missing_local_blob_page(
139 &self,
140 blob_dir: &Path,
141 limit: usize,
142 ) -> Result<MissingBlobPage> {
143 let mut conn = self.acquire_reader().await?;
144 missing_local_blob_page(&mut conn, blob_dir, limit).await
145 }
146
147 pub async fn missing_sync_attachment_counts(&self, blob_dir: &Path) -> Result<ByteCount> {
148 let mut conn = self.acquire_reader().await?;
149 Ok(missing_local_blob_page(&mut conn, blob_dir, 0).await?.total)
150 }
151
152 pub(super) async fn store_downloaded_blob(
153 &self,
154 blob_dir: &Path,
155 policy: LifecyclePolicy,
156 blob: &MissingLocalBlob,
157 bytes: Vec<u8>,
158 ) -> Result<()> {
159 let expected = usize::try_from(blob.byte_size).context("attachment bytes exceed usize")?;
160 if bytes.len() != expected || sha256_hex(&bytes) != blob.sha256 {
161 bail!("error attachment-blob-remote-invalid");
162 }
163 let validated = validate_image(bytes, Some(blob.media_type.clone()))
164 .await
165 .context("error attachment-blob-remote-invalid")?;
166 if (blob.width, blob.height) != (Some(validated.facts.width), Some(validated.facts.height))
167 {
168 bail!("error attachment-blob-remote-invalid");
169 }
170 let mut conn = self.acquire_writer().await?;
171 let reservation = ensure_local_capacity(
172 &mut conn,
173 blob_dir,
174 &blob.sha256,
175 blob.byte_size,
176 policy,
177 &SystemClock,
178 )
179 .await?;
180 let result = store_validated_blob(&mut conn, blob_dir, validated).await;
181 if let Some(reservation) = reservation {
182 release_reservation(&mut conn, &reservation).await?;
183 }
184 result.map(|_| ())
185 }
186
187 pub async fn prepare_server_blob_uploads(
188 &self,
189 blob_dir: &Path,
190 policy: LifecyclePolicy,
191 blobs: &[BlobUploadContract],
192 ) -> Result<Vec<String>> {
193 super::wire::validate_blob_contracts(blobs)?;
194 let mut conn = self.acquire_writer().await?;
195 let mut missing = Vec::new();
196 let mut missing_hashes = HashSet::new();
197 for blob in blobs {
198 if blob_available(&mut conn, blob_dir, &blob.sha256).await? {
199 reserve_upload(
200 &mut conn,
201 &blob.workspace_id,
202 &blob.sha256,
203 blob.byte_size,
204 policy.quota_bytes,
205 &SystemClock,
206 )
207 .await?;
208 } else if missing_hashes.insert(blob.sha256.as_str()) {
209 missing.push(blob.sha256.clone());
210 }
211 }
212 Ok(missing)
213 }
214
215 pub async fn store_server_blob(
216 &self,
217 blob_dir: &Path,
218 policy: LifecyclePolicy,
219 contract: &BlobUploadContract,
220 bytes: Vec<u8>,
221 ) -> Result<()> {
222 super::wire::validate_blob_contracts(std::slice::from_ref(contract))?;
223 if sha256_hex(&bytes) != contract.sha256
224 || usize::try_from(contract.byte_size).ok() != Some(bytes.len())
225 {
226 bail!("error blob-hash-or-size-mismatch");
227 }
228 let validated = validate_image(bytes, Some(contract.media_type.clone()))
229 .await
230 .context("error blob-validation-failed")?;
231 if (validated.facts.width, validated.facts.height) != (contract.width, contract.height) {
232 bail!("error blob-validation-failed");
233 }
234 let mut conn = self.acquire_writer().await?;
235 let reservation = reserve_upload(
236 &mut conn,
237 &contract.workspace_id,
238 &contract.sha256,
239 contract.byte_size,
240 policy.quota_bytes,
241 &SystemClock,
242 )
243 .await?;
244 let result = store_validated_blob(&mut conn, blob_dir, validated).await;
245 if result.is_err()
246 && let Some(reservation) = reservation
247 {
248 release_reservation(&mut conn, &reservation).await?;
249 }
250 result.map(|_| ())
251 }
252
253 pub async fn read_server_blob(
254 &self,
255 blob_dir: &Path,
256 sha256: &str,
257 ) -> Result<Option<ServerBlobDownload>> {
258 super::wire::validate_blob_hashes(&[sha256.to_string()])?;
259 let lease = {
260 let mut conn = self.acquire_writer().await?;
261 if !blob_available(&mut conn, blob_dir, sha256).await? {
262 return Ok(None);
263 }
264 acquire_lease(&mut conn, sha256, "transfer", &SystemClock).await?
265 };
266 let result = tokio::fs::read(object_path(blob_dir, sha256)?).await;
267 let mut conn = self.acquire_writer().await?;
268 release_lease(&mut conn, &lease).await?;
269 Ok(Some(ServerBlobDownload {
270 bytes: result.context("error blob-read-failed")?,
271 }))
272 }
273
274 pub async fn maintain_server_blobs(
275 &self,
276 blob_dir: &Path,
277 policy: LifecyclePolicy,
278 ) -> Result<PruneSummary> {
279 let mut conn = self.acquire_writer().await?;
280 prune(&mut conn, blob_dir, policy, true, &SystemClock).await
281 }
282}
283
284async fn missing_local_blob_page(
285 conn: &mut SqliteConnection,
286 blob_dir: &Path,
287 limit: usize,
288) -> Result<MissingBlobPage> {
289 let rows = sqlx::query(
290 "SELECT ta.sha256, MAX(ta.byte_size) AS byte_size,
291 MAX(ta.media_type) AS media_type, MAX(ta.width) AS width,
292 MAX(ta.height) AS height, MAX(COALESCE(bi.available, 0)) AS available
293 FROM task_attachments ta
294 JOIN tasks t ON t.workspace_id = ta.workspace_id AND t.id = ta.task_id
295 LEFT JOIN blob_inventory bi ON bi.sha256 = ta.sha256
296 WHERE ta.deleted = 0 AND t.deleted = 0
297 GROUP BY ta.sha256
298 ORDER BY ta.sha256",
299 )
300 .fetch_all(&mut *conn)
301 .await?;
302 let mut blobs = Vec::new();
303 let mut total = ByteCount::default();
304 for row in rows {
305 let sha256: String = row.get("sha256");
306 let available: i64 = row.get("available");
307 if available != 0 && object_path(blob_dir, &sha256)?.exists() {
308 continue;
309 }
310 let blob = MissingLocalBlob {
311 sha256,
312 byte_size: row.get("byte_size"),
313 media_type: row.get("media_type"),
314 width: row.get("width"),
315 height: row.get("height"),
316 };
317 total.count += 1;
318 total.bytes += u64::try_from(blob.byte_size)?;
319 if blobs.len() < limit {
320 blobs.push(blob);
321 }
322 }
323 Ok(MissingBlobPage { blobs, total })
324}
325
326pub(super) async fn blob_available(
327 conn: &mut SqliteConnection,
328 blob_dir: &Path,
329 sha256: &str,
330) -> Result<bool> {
331 let Some(row) = blob_inventory_row(conn, sha256).await? else {
332 return Ok(false);
333 };
334 Ok(row.available && object_path(blob_dir, sha256)?.exists())
335}
336
337#[cfg(test)]
338mod tests {
339 use serde_json::json;
340
341 use super::*;
342
343 fn attachment_change(op_type: &str) -> ChangeWire {
344 ChangeWire {
345 change_id: "AAAAAAAAAAAAAAA0".to_string(),
346 client_id: "client".to_string(),
347 local_seq: 1,
348 entity_type: "task".to_string(),
349 entity_id: "BBBBBBBBBBBBBBBB".to_string(),
350 field: Some("attachments".to_string()),
351 op_type: op_type.to_string(),
352 payload: json!({
353 "workspace_id": "0000000000000000",
354 "workspace_key": "default",
355 "attachment_id": "7KQ9A1X4MV2P8D6R",
356 "sha256": "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789",
357 "byte_size": 12,
358 "media_type": "image/png",
359 "filename": null,
360 "alt_text": null,
361 "width": 320,
362 "height": 240,
363 "created_at": "2026-06-01T00:00:00Z"
364 }),
365 base_version: None,
366 created_at: "2026-06-01T00:00:00Z".to_string(),
367 server_seq: None,
368 }
369 }
370
371 #[test]
372 fn attachment_blob_contract_uses_typed_add_payload() {
373 let add = attachment_change(op_type::ATTACHMENT_ADD);
374 let contract = attachment_blob_contract(&add).unwrap().unwrap();
375 assert_eq!(contract.workspace_id, "0000000000000000");
376 assert_eq!(
377 contract.sha256,
378 "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789"
379 );
380 assert_eq!(contract.byte_size, 12);
381 assert_eq!(contract.media_type, "image/png");
382 assert_eq!((contract.width, contract.height), (320, 240));
383
384 let mut malformed = add;
385 malformed.payload["sha256"] = json!("invalid");
386 assert!(
387 attachment_blob_contract(&malformed)
388 .unwrap_err()
389 .to_string()
390 .contains("invalid-sha256")
391 );
392
393 assert!(
394 attachment_blob_contract(&attachment_change(op_type::ATTACHMENT_DELETE))
395 .unwrap()
396 .is_none()
397 );
398 assert!(
399 attachment_blob_contract(&attachment_change(op_type::NOTE_ADD))
400 .unwrap()
401 .is_none()
402 );
403 }
404
405 async fn seed_missing_attachments(database: &Database, count: usize) {
406 let mut conn = database.acquire_writer().await.unwrap();
407 for index in 0..count {
408 let task_id = format!("{index:016X}");
409 let attachment_id = format!("{:016X}", index + 100);
410 let sha256 = format!("{index:064x}");
411 sqlx::query(
412 "INSERT INTO tasks(
413 workspace_id, id, title, description, project_id, status, priority,
414 created_at, updated_at, queue_activity_at
415 ) VALUES ('0000000000000000', ?, 'task', '', 'project', 'inbox', 'none',
416 '2026-01-01T00:00:00Z', '2026-01-01T00:00:00Z',
417 '2026-01-01T00:00:00Z')",
418 )
419 .bind(&task_id)
420 .execute(&mut *conn)
421 .await
422 .unwrap();
423 sqlx::query(
424 "INSERT INTO task_attachments(
425 workspace_id, attachment_id, task_id, sha256, byte_size, media_type,
426 width, height, created_at
427 ) VALUES ('0000000000000000', ?, ?, ?, 4, 'image/png', 1, 1,
428 '2026-01-01T00:00:00Z')",
429 )
430 .bind(attachment_id)
431 .bind(task_id)
432 .bind(sha256)
433 .execute(&mut *conn)
434 .await
435 .unwrap();
436 }
437 }
438
439 #[tokio::test]
440 async fn missing_blob_page_is_bounded_and_reports_more() {
441 let temp = tempfile::tempdir().unwrap();
442 let database = Database::open(&temp.path().join("replica.sqlite"))
443 .await
444 .unwrap();
445 seed_missing_attachments(&database, 3).await;
446 {
447 let mut conn = database.acquire_writer().await.unwrap();
448 crate::attachments::storage::upsert_inventory_available(
449 &mut conn,
450 &format!("{:064x}", 0),
451 4,
452 "image/png",
453 )
454 .await
455 .unwrap();
456 }
457
458 let page = database
459 .missing_local_blob_page(&temp.path().join("blobs"), 2)
460 .await
461 .unwrap();
462 assert_eq!(page.blobs.len(), 2);
463 assert_eq!(
464 page.total,
465 ByteCount {
466 count: 3,
467 bytes: 12
468 }
469 );
470 }
471}