1use std::collections::HashSet;
2use std::path::Path;
3
4use anyhow::{Context, Result, bail};
5use sqlx::{Row, SqliteConnection};
6
7use super::wire::{BlobUploadContract, ChangeWire};
8use crate::attachments::decode::validate_image;
9use crate::attachments::lifecycle::{
10 ByteCount, LifecyclePolicy, SystemClock, acquire_lease, ensure_local_capacity, prune,
11 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 PreparedBlobUpload {
30 pub bytes: Vec<u8>,
31 pub lease_id: String,
32}
33
34#[derive(Debug)]
35pub struct ServerBlobDownload {
36 pub bytes: Vec<u8>,
37}
38
39pub(super) fn attachment_blob_contract(change: &ChangeWire) -> Result<Option<BlobUploadContract>> {
40 if change.op_type != op_type::ATTACHMENT_ADD {
41 return Ok(None);
42 }
43 Ok(Some(BlobUploadContract {
44 workspace_id: payload_string(change, "workspace_id")?,
45 sha256: payload_string(change, "sha256")?,
46 byte_size: payload_i64(change, "byte_size")?,
47 media_type: payload_string(change, "media_type")?,
48 width: payload_i64(change, "width")?,
49 height: payload_i64(change, "height")?,
50 }))
51}
52
53pub(super) fn unique_blob_contracts(changes: &[ChangeWire]) -> Result<Vec<BlobUploadContract>> {
54 let mut seen = HashSet::new();
55 let mut contracts = Vec::new();
56 for change in changes {
57 if let Some(contract) = attachment_blob_contract(change)?
58 && seen.insert((contract.workspace_id.clone(), contract.sha256.clone()))
59 {
60 contracts.push(contract);
61 }
62 }
63 Ok(contracts)
64}
65
66pub(super) fn contract_by_hash(
67 contracts: &[BlobUploadContract],
68 sha256: &str,
69) -> Result<BlobUploadContract> {
70 contracts
71 .iter()
72 .find(|contract| contract.sha256 == sha256)
73 .cloned()
74 .context("payload missing attachment blob contract")
75}
76
77fn payload_string(change: &ChangeWire, key: &str) -> Result<String> {
78 change
79 .payload
80 .get(key)
81 .and_then(serde_json::Value::as_str)
82 .map(str::to_string)
83 .with_context(|| format!("payload missing {key}"))
84}
85
86fn payload_i64(change: &ChangeWire, key: &str) -> Result<i64> {
87 change
88 .payload
89 .get(key)
90 .and_then(serde_json::Value::as_i64)
91 .with_context(|| format!("payload missing {key}"))
92}
93
94impl Database {
95 pub(super) async fn prepare_blob_upload(
96 &self,
97 blob_dir: &Path,
98 contract: &BlobUploadContract,
99 ) -> Result<PreparedBlobUpload> {
100 let mut conn = self.acquire().await?;
101 let row = blob_inventory_row(&mut conn, &contract.sha256)
102 .await?
103 .filter(|row| row.available)
104 .context("error attachment-blob-local-missing")?;
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 != i64::try_from(bytes.len()).context("attachment bytes exceed i64")?
110 {
111 bail!("error attachment-blob-local-invalid");
112 }
113 let validated = validate_image(bytes.clone(), Some(row.media_type.clone())).await?;
114 if (validated.facts.width, validated.facts.height) != (contract.width, contract.height)
115 || validated.facts.media_type != contract.media_type
116 || contract.byte_size != row.byte_size
117 {
118 bail!("error attachment-blob-local-invalid");
119 }
120 let lease_id = acquire_lease(&mut conn, &contract.sha256, "transfer", &SystemClock).await?;
121 Ok(PreparedBlobUpload { bytes, lease_id })
122 }
123
124 pub(super) async fn finish_blob_upload(&self, lease_id: &str) -> Result<()> {
125 let mut conn = self.acquire().await?;
126 release_lease(&mut conn, lease_id).await
127 }
128
129 pub(super) async fn missing_local_blobs(
130 &self,
131 blob_dir: &Path,
132 ) -> Result<Vec<MissingLocalBlob>> {
133 let mut conn = self.acquire().await?;
134 missing_local_blobs(&mut conn, blob_dir).await
135 }
136
137 pub(super) async fn store_downloaded_blob(
138 &self,
139 blob_dir: &Path,
140 policy: LifecyclePolicy,
141 blob: &MissingLocalBlob,
142 bytes: Vec<u8>,
143 ) -> Result<()> {
144 let expected = usize::try_from(blob.byte_size).context("attachment bytes exceed usize")?;
145 if bytes.len() != expected || sha256_hex(&bytes) != blob.sha256 {
146 bail!("error attachment-blob-remote-invalid");
147 }
148 let validated = validate_image(bytes, Some(blob.media_type.clone()))
149 .await
150 .context("error attachment-blob-remote-invalid")?;
151 if (blob.width, blob.height) != (Some(validated.facts.width), Some(validated.facts.height))
152 {
153 bail!("error attachment-blob-remote-invalid");
154 }
155 let mut conn = self.acquire().await?;
156 let reservation = ensure_local_capacity(
157 &mut conn,
158 blob_dir,
159 &blob.sha256,
160 blob.byte_size,
161 policy,
162 &SystemClock,
163 )
164 .await?;
165 let result = store_validated_blob(&mut conn, blob_dir, validated).await;
166 if let Some(reservation) = reservation {
167 release_reservation(&mut conn, &reservation).await?;
168 }
169 result.map(|_| ())
170 }
171
172 pub async fn prepare_server_blob_uploads(
173 &self,
174 blob_dir: &Path,
175 policy: LifecyclePolicy,
176 blobs: &[BlobUploadContract],
177 ) -> Result<Vec<String>> {
178 super::wire::validate_blob_contracts(blobs)?;
179 let mut conn = self.acquire().await?;
180 prune(&mut conn, blob_dir, policy, true, &SystemClock).await?;
181 let mut missing = Vec::new();
182 let mut missing_hashes = HashSet::new();
183 for blob in blobs {
184 if blob_available(&mut conn, blob_dir, &blob.sha256).await? {
185 reserve_upload(
186 &mut conn,
187 &blob.workspace_id,
188 &blob.sha256,
189 blob.byte_size,
190 policy.quota_bytes,
191 &SystemClock,
192 )
193 .await?;
194 } else if missing_hashes.insert(blob.sha256.as_str()) {
195 missing.push(blob.sha256.clone());
196 }
197 }
198 Ok(missing)
199 }
200
201 pub async fn store_server_blob(
202 &self,
203 blob_dir: &Path,
204 policy: LifecyclePolicy,
205 contract: &BlobUploadContract,
206 bytes: Vec<u8>,
207 ) -> Result<()> {
208 super::wire::validate_blob_contracts(std::slice::from_ref(contract))?;
209 if sha256_hex(&bytes) != contract.sha256
210 || usize::try_from(contract.byte_size).ok() != Some(bytes.len())
211 {
212 bail!("error blob-hash-or-size-mismatch");
213 }
214 let validated = validate_image(bytes, Some(contract.media_type.clone()))
215 .await
216 .context("error blob-validation-failed")?;
217 if (validated.facts.width, validated.facts.height) != (contract.width, contract.height) {
218 bail!("error blob-validation-failed");
219 }
220 let mut conn = self.acquire().await?;
221 prune(&mut conn, blob_dir, policy, true, &SystemClock).await?;
222 let reservation = reserve_upload(
223 &mut conn,
224 &contract.workspace_id,
225 &contract.sha256,
226 contract.byte_size,
227 policy.quota_bytes,
228 &SystemClock,
229 )
230 .await?;
231 let result = store_validated_blob(&mut conn, blob_dir, validated).await;
232 if result.is_err()
233 && let Some(reservation) = reservation
234 {
235 release_reservation(&mut conn, &reservation).await?;
236 }
237 result.map(|_| ())
238 }
239
240 pub async fn read_server_blob(
241 &self,
242 blob_dir: &Path,
243 sha256: &str,
244 ) -> Result<Option<ServerBlobDownload>> {
245 super::wire::validate_blob_hashes(&[sha256.to_string()])?;
246 let mut conn = self.acquire().await?;
247 if !blob_available(&mut conn, blob_dir, sha256).await? {
248 return Ok(None);
249 }
250 let lease = acquire_lease(&mut conn, sha256, "transfer", &SystemClock).await?;
251 let result = tokio::fs::read(object_path(blob_dir, sha256)?).await;
252 release_lease(&mut conn, &lease).await?;
253 Ok(Some(ServerBlobDownload {
254 bytes: result.context("error blob-read-failed")?,
255 }))
256 }
257
258 pub async fn maintain_server_blobs(
259 &self,
260 blob_dir: &Path,
261 policy: LifecyclePolicy,
262 ) -> Result<()> {
263 let mut conn = self.acquire().await?;
264 prune(&mut conn, blob_dir, policy, true, &SystemClock)
265 .await
266 .map(|_| ())
267 }
268}
269
270async fn missing_local_blobs(
271 conn: &mut SqliteConnection,
272 blob_dir: &Path,
273) -> Result<Vec<MissingLocalBlob>> {
274 let rows = sqlx::query(
275 "SELECT ta.sha256, MAX(ta.byte_size) AS byte_size,
276 MAX(ta.media_type) AS media_type, MAX(ta.width) AS width,
277 MAX(ta.height) AS height, MAX(COALESCE(bi.available, 0)) AS available
278 FROM task_attachments ta
279 JOIN tasks t ON t.workspace_id = ta.workspace_id AND t.id = ta.task_id
280 LEFT JOIN blob_inventory bi ON bi.sha256 = ta.sha256
281 WHERE ta.deleted = 0 AND t.deleted = 0
282 GROUP BY ta.sha256
283 ORDER BY ta.sha256",
284 )
285 .fetch_all(&mut *conn)
286 .await?;
287 let mut missing = Vec::new();
288 for row in rows {
289 let sha256: String = row.get("sha256");
290 let available: i64 = row.get("available");
291 if available == 0 || !object_path(blob_dir, &sha256)?.exists() {
292 missing.push(MissingLocalBlob {
293 sha256,
294 byte_size: row.get("byte_size"),
295 media_type: row.get("media_type"),
296 width: row.get("width"),
297 height: row.get("height"),
298 });
299 }
300 }
301 Ok(missing)
302}
303
304pub(super) async fn blob_available(
305 conn: &mut SqliteConnection,
306 blob_dir: &Path,
307 sha256: &str,
308) -> Result<bool> {
309 let Some(row) = blob_inventory_row(conn, sha256).await? else {
310 return Ok(false);
311 };
312 Ok(row.available && object_path(blob_dir, sha256)?.exists())
313}
314
315pub(super) fn missing_counts(blobs: &[MissingLocalBlob]) -> ByteCount {
316 ByteCount {
317 count: blobs.len() as u64,
318 bytes: blobs
319 .iter()
320 .map(|blob| u64::try_from(blob.byte_size).unwrap_or(0))
321 .sum(),
322 }
323}