Skip to main content

max / makenotwork

11.2 KB · 233 lines History Blame Raw
1 //! The storage trait every backend implements, and the capped read every
2 //! caller of it goes through.
3
4 use super::bucket::S3DeleteAuthority;
5 use super::key::S3Key;
6 use crate::error::{AppError, Result};
7
8 /// Aggregate a `ByteStream` into memory, aborting once more than `max_bytes`
9 /// have been read. Backs [`StorageBackend::download_object_buf_capped`]; factored
10 /// out as a free function so the cap logic is unit-testable without a full
11 /// backend. `label` is only used in the error message (the object key).
12 pub(crate) async fn read_bytestream_capped(
13 mut stream: s3_storage::ByteStream,
14 label: &str,
15 max_bytes: u64,
16 ) -> Result<bytes::Bytes> {
17 let mut buf = bytes::BytesMut::new();
18 let mut read: u64 = 0;
19 loop {
20 match stream.try_next().await {
21 Ok(Some(chunk)) => {
22 read += chunk.len() as u64;
23 if read > max_bytes {
24 return Err(AppError::Storage(format!(
25 "object {label} exceeds scan in-memory cap ({read} > {max_bytes} bytes); \
26 recorded size under-reported the real object"
27 )));
28 }
29 buf.extend_from_slice(&chunk);
30 }
31 Ok(None) => break,
32 Err(e) => return Err(AppError::Storage(format!("read object from S3: {e}"))),
33 }
34 }
35 Ok(buf.freeze())
36 }
37
38 /// Abstract storage backend, implemented by `S3Client` (production) and
39 /// `InMemoryStorage` (tests). Routes access storage through this trait.
40 #[async_trait::async_trait]
41 pub trait StorageBackend: Send + Sync {
42 /// Generate a presigned upload URL. `max_bytes`, when set, is signed into
43 /// the URL as `Content-Length` so S3 itself enforces the size cap at the
44 /// protocol level (prevents oversized PUTs from burning bandwidth before
45 /// hitting the post-PUT delete-and-charge fallback).
46 async fn presign_upload(
47 &self,
48 s3_key: &S3Key,
49 content_type: &str,
50 expiry_secs: Option<u64>,
51 cache_control: Option<&str>,
52 max_bytes: Option<i64>,
53 ) -> Result<String>;
54 async fn presign_download(&self, s3_key: &S3Key, expiry_secs: Option<u64>) -> Result<String>;
55 async fn object_exists(&self, s3_key: &str) -> Result<bool>;
56 async fn object_size(&self, s3_key: &str) -> Result<Option<i64>>;
57 async fn download_object(&self, s3_key: &str) -> Result<Vec<u8>>;
58 /// Download as `bytes::Bytes`, no `to_vec` copy of the aggregated body.
59 /// Memory-sensitive callers (the scanner's buffered branch) use this so the
60 /// payload isn't transiently doubled.
61 async fn download_object_buf(&self, s3_key: &str) -> Result<bytes::Bytes>;
62 /// Stream the object body without buffering the whole payload. Callers
63 /// drive the stream to disk (scanner spool) or to a layer that consumes
64 /// chunks directly.
65 async fn download_stream(&self, s3_key: &str) -> Result<s3_storage::ByteStream>;
66 /// Download into memory like [`download_object_buf`], but abort if the body
67 /// exceeds `max_bytes`. The scanner routes files it *recorded* as small to an
68 /// in-memory branch, but `file_size_bytes` is asserted at upload time and can
69 /// under-report the real object; this bounds the aggregation so a mis-recorded
70 /// or abusive object can't pull an unbounded body into RAM, the independent
71 /// ceiling the spool path already enforces. Streams via
72 /// `download_stream`, so no backend can hand back the whole body up front.
73 async fn download_object_buf_capped(
74 &self,
75 s3_key: &str,
76 max_bytes: u64,
77 ) -> Result<bytes::Bytes> {
78 let stream = self.download_stream(s3_key).await?;
79 read_bytestream_capped(stream, s3_key, max_bytes).await
80 }
81 /// Read up to the first `len` bytes of an object. Production overrides this
82 /// with a ranged `GetObject` so a content sniff transfers only the header,
83 /// not the whole object. The default streams and stops early, correct, but
84 /// it still initiates a full GET, which is fine for in-memory test backends.
85 async fn download_head(&self, s3_key: &str, len: usize) -> Result<Vec<u8>> {
86 let mut stream = self.download_stream(s3_key).await?;
87 let mut head = Vec::with_capacity(len.min(64 * 1024));
88 while head.len() < len {
89 match stream.try_next().await {
90 Ok(Some(chunk)) => head.extend_from_slice(&chunk),
91 Ok(None) => break,
92 Err(e) => return Err(AppError::Storage(format!("read object head from S3: {e}"))),
93 }
94 }
95 head.truncate(len);
96 Ok(head)
97 }
98 async fn upload_object(
99 &self,
100 s3_key: &S3Key,
101 content_type: &str,
102 data: Vec<u8>,
103 cache_control: Option<&str>,
104 ) -> Result<()>;
105 /// Delete an object. Requires an [`S3DeleteAuthority`], route handlers
106 /// cannot mint one, so they must enqueue through `pending_s3_deletions`
107 /// instead of deleting directly.
108 async fn delete_object(&self, auth: &S3DeleteAuthority, s3_key: &S3Key) -> Result<()>;
109 /// Delete a batch of objects in a single S3 `DeleteObjects` request
110 /// (up to 1000 keys/call). Default loops `delete_object` so test backends
111 /// don't have to implement it, but production should override.
112 async fn delete_objects(&self, auth: &S3DeleteAuthority, keys: &[S3Key]) -> Result<()> {
113 let mut failed = 0usize;
114 for k in keys {
115 if let Err(e) = self.delete_object(auth, k).await {
116 failed += 1;
117 tracing::warn!(key = %k, error = ?e, "delete_objects: per-key delete failed");
118 }
119 }
120 // Don't report success when every key failed, a total failure must
121 // surface so the caller can fall back (Run #2 Storage MINOR). Partial
122 // failures stay logged; callers pre-enqueue to pending_s3_deletions.
123 if !keys.is_empty() && failed == keys.len() {
124 return Err(AppError::Storage(format!(
125 "delete_objects: all {failed} keys failed"
126 )));
127 }
128 Ok(())
129 }
130 /// Delete all objects under a key prefix. Default logs a warning (no-op).
131 async fn delete_prefix(&self, _auth: &S3DeleteAuthority, _prefix: &str) -> Result<()> {
132 tracing::warn!("delete_prefix called on a storage backend that does not implement it");
133 Ok(())
134 }
135 /// Upload a file via S3 multipart upload. Required (not defaulted): a
136 /// default that `tokio::fs::read`s the whole file into RAM + single PUT
137 /// silently defeats streaming, so a future backend that forgot to override
138 /// it would quietly lose multipart. Every backend must declare its strategy.
139 async fn upload_multipart(
140 &self,
141 s3_key: &S3Key,
142 content_type: &str,
143 file_path: &std::path::Path,
144 ) -> Result<()>;
145 /// Server-side copy `src_key` to `dst_key` within this backend's bucket
146 /// (no bytes transit the process). The scan-then-promote primitive: a Clean
147 /// staging object is copied to the served key the client holds no presign
148 /// for, so served bytes are provably the scanned bytes. Required (not
149 /// defaulted): a silent no-op default would make a
150 /// promote "succeed" while the served key stays empty.
151 async fn copy_object(&self, src_key: &S3Key, dst_key: &S3Key) -> Result<()>;
152 /// Server-side copy from `src_bucket` into THIS backend's bucket. The
153 /// cross-bucket half of scan-then-promote: a Clean staging object in the
154 /// private bucket is lifted into the public (CDN-served) bucket. Call on the
155 /// public backend with the private bucket name as `src_bucket`. Required
156 /// (not defaulted) for the same reason as `copy_object`.
157 async fn copy_object_from(
158 &self,
159 src_bucket: &str,
160 src_key: &S3Key,
161 dst_key: &S3Key,
162 ) -> Result<()>;
163 /// Server-side multipart copy (`UploadPartCopy`) for sources over the 5 GiB
164 /// single-part `CopyObject` limit, the >5 GiB half of scan-then-promote.
165 /// Always takes `src_bucket` explicitly, collapsing the
166 /// `copy_object`/`copy_object_from` pair into one method (pass this
167 /// backend's own bucket for a same-bucket promote). `content_type` sets the
168 /// destination's type, since a fresh multipart upload does not inherit the
169 /// source's metadata the way `CopyObject` does. Required (not defaulted) for
170 /// the same reason as `copy_object`.
171 async fn copy_object_multipart(
172 &self,
173 src_bucket: &str,
174 src_key: &S3Key,
175 dst_key: &S3Key,
176 content_type: &str,
177 src_size: u64,
178 part_size: Option<usize>,
179 ) -> Result<()>;
180
181 // Client-direct multipart sessions
182 //
183 // The counterpart to `upload_multipart`, which drives a whole transfer
184 // server-side from a local file. Here the server only mints the session and
185 // the per-part presigned URLs; the client streams parts straight to S3, so
186 // no object bytes transit the server. This is the path large CLI/desktop
187 // uploads take (a browser stays on the single-PUT `presign_upload`).
188 //
189 // All four are required (not defaulted): a no-op default would mint a
190 // session no client could complete, or silently drop the cleanup that
191 // stops orphaned parts billing forever.
192
193 /// Begin a client-direct multipart upload, returning the `upload_id` the
194 /// part/complete/abort calls key on.
195 async fn create_multipart_upload(&self, s3_key: &S3Key, content_type: &str) -> Result<String>;
196 /// Presign an `UploadPart` request for one part (1-based `part_number`).
197 /// `max_bytes`, when set, is signed as `Content-Length`, the same
198 /// defense-in-depth as [`Self::presign_upload`], the authoritative size
199 /// check still happens at confirm time.
200 ///
201 /// `checksum_sha256` (base64 of the raw digest), when set, is signed as
202 /// `x-amz-checksum-sha256` and IS enforced: S3 rehashes the part and
203 /// rejects a mismatch before the bytes are durable.
204 async fn presign_upload_part(
205 &self,
206 s3_key: &S3Key,
207 upload_id: &str,
208 part_number: i32,
209 expiry_secs: Option<u64>,
210 max_bytes: Option<i64>,
211 checksum_sha256: Option<&str>,
212 ) -> Result<String>;
213 /// Complete a multipart upload from the collected `(part_number, etag)`
214 /// pairs. Parts may be passed in any order; the backend sorts them.
215 async fn complete_multipart_upload(
216 &self,
217 s3_key: &S3Key,
218 upload_id: &str,
219 parts: &[(i32, String)],
220 ) -> Result<()>;
221 /// Abort a multipart upload, releasing its uploaded parts. The
222 /// pending-upload reaper calls this on sessions that were never confirmed,
223 /// incomplete multipart uploads bill for their parts indefinitely.
224 async fn abort_multipart_upload(&self, s3_key: &S3Key, upload_id: &str) -> Result<()>;
225 /// Upload ids of the in-progress multipart sessions for exactly `s3_key`.
226 /// The reaper recovers them from S3 rather than the database, so a session
227 /// whose tracking row was lost is still cleaned up. Required (not defaulted):
228 /// an empty default would silently strand billed parts.
229 async fn list_multipart_uploads_for_key(&self, s3_key: &str) -> Result<Vec<String>>;
230 async fn check_connectivity(&self) -> std::result::Result<(), String>;
231 fn bucket(&self) -> &str;
232 }
233