Skip to main content

max / makenotwork

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