Skip to main content

max / makenotwork

10.5 KB · 305 lines History Blame Raw
1 //! In-memory storage backend for integration tests.
2
3 use super::faults::Faults;
4 use makenotwork::error::{AppError, Result};
5 use makenotwork::storage::{S3DeleteAuthority, S3Key, StorageBackend};
6 use std::collections::HashMap;
7 use std::sync::Mutex;
8
9 /// In-memory implementation of `StorageBackend` for tests.
10 /// Files are stored in a `HashMap<String, Vec<u8>>` behind a `Mutex`.
11 pub(crate) struct InMemoryStorage {
12 objects: Mutex<HashMap<String, Vec<u8>>>,
13 /// Open multipart sessions: `upload_id -> s3_key`. Tracked for real (not
14 /// stubbed) so the orphan reaper's abort path is observable in tests, an
15 /// unaborted session is exactly the leak this models.
16 multipart: Mutex<HashMap<String, String>>,
17 bucket: String,
18 /// Injected transport failures. Empty by default, so a backend with no
19 /// policy installed behaves exactly as it did before this existed.
20 faults: Faults,
21 }
22
23 #[allow(dead_code)]
24 impl InMemoryStorage {
25 pub(crate) fn new() -> Self {
26 InMemoryStorage {
27 objects: Mutex::new(HashMap::new()),
28 multipart: Mutex::new(HashMap::new()),
29 bucket: "test-bucket".to_string(),
30 faults: Faults::new(),
31 }
32 }
33
34 /// The failure policy. Install rules on it to reach the retry and
35 /// compensation paths that no test could otherwise enter.
36 pub(crate) fn faults(&self) -> &Faults {
37 &self.faults
38 }
39
40 /// Number of multipart sessions still open. Zero after a clean reap.
41 pub(crate) fn open_multipart_count(&self) -> usize {
42 self.multipart.lock().unwrap().len()
43 }
44
45 /// Open a multipart session directly, standing in for one a client started
46 /// and abandoned.
47 pub(crate) fn put_open_multipart(&self, upload_id: &str, s3_key: &str) {
48 self.multipart
49 .lock()
50 .unwrap()
51 .insert(upload_id.to_string(), s3_key.to_string());
52 }
53
54 /// Pre-populate a file so that subsequent `object_exists` / `download_object`
55 /// calls see it. Useful for testing confirm_upload flows.
56 pub(crate) fn put(&self, key: &str, data: Vec<u8>) {
57 self.objects.lock().unwrap().insert(key.to_string(), data);
58 }
59
60 /// Retrieve stored bytes for a key. Panics if not found.
61 pub(crate) fn get(&self, key: &str) -> Vec<u8> {
62 self.objects
63 .lock()
64 .unwrap()
65 .get(key)
66 .cloned()
67 .expect("key not found in storage")
68 }
69 }
70
71 #[async_trait::async_trait]
72 impl StorageBackend for InMemoryStorage {
73 async fn presign_upload(
74 &self,
75 s3_key: &S3Key,
76 _content_type: &str,
77 _expiry_secs: Option<u64>,
78 _cache_control: Option<&str>,
79 _max_bytes: Option<i64>,
80 ) -> Result<String> {
81 self.faults.check("presign_upload")?;
82 Ok(format!("http://test-storage/{s3_key}"))
83 }
84
85 async fn presign_download(&self, s3_key: &S3Key, _expiry_secs: Option<u64>) -> Result<String> {
86 self.faults.check("presign_download")?;
87 if self.objects.lock().unwrap().contains_key(s3_key.as_str()) {
88 Ok(format!("http://test-storage/{s3_key}"))
89 } else {
90 Err(AppError::Storage(format!("Object not found: {s3_key}")))
91 }
92 }
93
94 async fn object_exists(&self, s3_key: &str) -> Result<bool> {
95 self.faults.check("object_exists")?;
96 Ok(self.objects.lock().unwrap().contains_key(s3_key))
97 }
98
99 async fn object_size(&self, s3_key: &str) -> Result<Option<i64>> {
100 self.faults.check("object_size")?;
101 Ok(self
102 .objects
103 .lock()
104 .unwrap()
105 .get(s3_key)
106 .map(|v| v.len() as i64))
107 }
108
109 // `download_object_buf` and the two `copy_object_*` wrappers delegate to the
110 // methods below them, so a policy on the delegate is what fires and the call
111 // count is recorded once. Name the underlying operation in a rule, not the
112 // wrapper.
113 async fn download_object(&self, s3_key: &str) -> Result<Vec<u8>> {
114 self.faults.check("download_object")?;
115 self.objects
116 .lock()
117 .unwrap()
118 .get(s3_key)
119 .cloned()
120 .ok_or_else(|| AppError::Storage(format!("Object not found: {s3_key}")))
121 }
122
123 async fn download_object_buf(&self, s3_key: &str) -> Result<bytes::Bytes> {
124 self.download_object(s3_key).await.map(bytes::Bytes::from)
125 }
126
127 async fn download_stream(&self, s3_key: &str) -> Result<s3_storage::ByteStream> {
128 self.faults.check("download_stream")?;
129 let bytes = self
130 .objects
131 .lock()
132 .unwrap()
133 .get(s3_key)
134 .cloned()
135 .ok_or_else(|| AppError::Storage(format!("Object not found: {s3_key}")))?;
136 Ok(s3_storage::ByteStream::from(bytes))
137 }
138
139 async fn upload_object(
140 &self,
141 s3_key: &S3Key,
142 _content_type: &str,
143 data: Vec<u8>,
144 _cache_control: Option<&str>,
145 ) -> Result<()> {
146 self.faults.check("upload_object")?;
147 self.objects
148 .lock()
149 .unwrap()
150 .insert(s3_key.as_str().to_string(), data);
151 Ok(())
152 }
153
154 async fn upload_multipart(
155 &self,
156 s3_key: &S3Key,
157 _content_type: &str,
158 file_path: &std::path::Path,
159 ) -> Result<()> {
160 self.faults.check("upload_multipart")?;
161 let data = tokio::fs::read(file_path)
162 .await
163 .map_err(|e| AppError::Storage(format!("read multipart source: {e}")))?;
164 self.objects
165 .lock()
166 .unwrap()
167 .insert(s3_key.as_str().to_string(), data);
168 Ok(())
169 }
170
171 async fn delete_object(&self, _auth: &S3DeleteAuthority, s3_key: &S3Key) -> Result<()> {
172 self.faults.check("delete_object")?;
173 self.objects.lock().unwrap().remove(s3_key.as_str());
174 Ok(())
175 }
176
177 async fn copy_object(&self, src_key: &S3Key, dst_key: &S3Key) -> Result<()> {
178 self.faults.check("copy_object")?;
179 let mut objects = self.objects.lock().unwrap();
180 let bytes = objects
181 .get(src_key.as_str())
182 .cloned()
183 .ok_or_else(|| AppError::Storage(format!("copy source not found: {src_key}")))?;
184 objects.insert(dst_key.as_str().to_string(), bytes);
185 Ok(())
186 }
187
188 async fn copy_object_from(
189 &self,
190 _src_bucket: &str,
191 src_key: &S3Key,
192 dst_key: &S3Key,
193 ) -> Result<()> {
194 // The in-memory store is one flat map with no bucket isolation, so the
195 // cross-bucket promote is just a same-map copy. Tests wire `public_s3` to
196 // the same backend as `s3` so the staging source resolves here.
197 self.copy_object(src_key, dst_key).await
198 }
199
200 async fn copy_object_multipart(
201 &self,
202 _src_bucket: &str,
203 src_key: &S3Key,
204 dst_key: &S3Key,
205 _content_type: &str,
206 _src_size: u64,
207 _part_size: Option<usize>,
208 ) -> Result<()> {
209 // The flat map has no 5 GiB single-copy limit, so the multipart promote
210 // is the same copy. Implemented for real (not stubbed) so a promote test
211 // gets the same observable result down either branch.
212 self.copy_object(src_key, dst_key).await
213 }
214
215 // Client-direct multipart sessions
216 //
217 // Stubbed the same way presigned single-PUT uploads already are: the URLs
218 // this backend hands out are fake, so no client bytes can flow back into the
219 // map. Tests simulate the finished object with `put()` and then exercise the
220 // confirm path, exactly as they do for `presign_upload` + confirm.
221
222 async fn create_multipart_upload(&self, s3_key: &S3Key, _content_type: &str) -> Result<String> {
223 self.faults.check("create_multipart_upload")?;
224 let upload_id = format!("test-upload-id/{s3_key}");
225 self.put_open_multipart(&upload_id, s3_key.as_str());
226 Ok(upload_id)
227 }
228
229 async fn presign_upload_part(
230 &self,
231 s3_key: &S3Key,
232 upload_id: &str,
233 part_number: i32,
234 _expiry_secs: Option<u64>,
235 _max_bytes: Option<i64>,
236 checksum_sha256: Option<&str>,
237 ) -> Result<String> {
238 self.faults.check("presign_upload_part")?;
239 // Mirror the production range check so a bad part number fails in tests
240 // the same way it would against S3.
241 if !(1..=s3_storage::MULTIPART_MAX_PARTS as i32).contains(&part_number) {
242 return Err(AppError::Storage(format!(
243 "part number {part_number} out of range 1..={}",
244 s3_storage::MULTIPART_MAX_PARTS
245 )));
246 }
247 // Echo the bound checksum into the URL so tests can assert it reached
248 // the signer, standing in for the SignedHeaders a real presign carries.
249 let checksum = checksum_sha256
250 .map(|c| format!("&checksum={c}"))
251 .unwrap_or_default();
252 Ok(format!(
253 "http://test-storage/{s3_key}?uploadId={upload_id}&partNumber={part_number}{checksum}"
254 ))
255 }
256
257 async fn complete_multipart_upload(
258 &self,
259 _s3_key: &S3Key,
260 upload_id: &str,
261 parts: &[(i32, String)],
262 ) -> Result<()> {
263 self.faults.check("complete_multipart_upload")?;
264 if parts.is_empty() {
265 return Err(AppError::Storage(
266 "cannot complete a multipart upload with no parts".to_string(),
267 ));
268 }
269 // Completing closes the session, so it no longer holds billed parts.
270 self.multipart.lock().unwrap().remove(upload_id);
271 Ok(())
272 }
273
274 async fn abort_multipart_upload(&self, _s3_key: &S3Key, upload_id: &str) -> Result<()> {
275 self.faults.check("abort_multipart_upload")?;
276 // Idempotent: aborting an unknown session is fine.
277 self.multipart.lock().unwrap().remove(upload_id);
278 Ok(())
279 }
280
281 async fn list_multipart_uploads_for_key(&self, s3_key: &str) -> Result<Vec<String>> {
282 self.faults.check("list_multipart_uploads_for_key")?;
283 Ok(self
284 .multipart
285 .lock()
286 .unwrap()
287 .iter()
288 .filter(|(_, key)| key.as_str() == s3_key)
289 .map(|(id, _)| id.clone())
290 .collect())
291 }
292
293 async fn check_connectivity(&self) -> std::result::Result<(), String> {
294 // Health checks report a string, not an `AppError`, so the injected
295 // error is rendered rather than propagated.
296 self.faults
297 .check("check_connectivity")
298 .map_err(|e| e.to_string())
299 }
300
301 fn bucket(&self) -> &str {
302 &self.bucket
303 }
304 }
305