Line data Source code
1 : //! An S3 server in memory for the tests of the commands that talk to the release bucket, and the files they read.
2 :
3 : use std::collections::{BTreeMap, HashSet};
4 : use std::path::{Path, PathBuf};
5 : use std::sync::{Arc, Mutex};
6 :
7 : use tokio::io::{AsyncReadExt, AsyncWriteExt};
8 : use tokio::net::{TcpListener, TcpStream};
9 :
10 : use crate::remote_access::RemoteAccess;
11 :
12 : /// Held by the tests that read or change the R2 variables of the environment
13 : pub static ENV_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
14 :
15 : /// A request the server received
16 : #[derive(Clone, Debug)]
17 : pub struct Request {
18 : pub method: String,
19 : /// The object key, empty for requests to the bucket
20 : pub key: String,
21 : pub query: BTreeMap<String, String>,
22 : /// The headers with lowercase names
23 : pub headers: BTreeMap<String, String>,
24 : /// The body without its chunked encoding
25 : pub body: Vec<u8>,
26 : }
27 :
28 : #[derive(Default)]
29 : struct State {
30 : objects: BTreeMap<String, Vec<u8>>,
31 : requests: Vec<Request>,
32 : /// Keys that uploads are refused for
33 : denied_uploads: HashSet<String>,
34 : /// Keys that a delete reports an error for and keeps
35 : kept_objects: HashSet<String>,
36 : deny_listing: bool,
37 : deny_deletes: bool,
38 : page_size: Option<usize>,
39 : }
40 :
41 : /// The server. It answers path-style requests to the bucket `dashql-get`, which is what the SDK sends to an IP address.
42 : pub struct FakeS3 {
43 : endpoint: String,
44 : state: Arc<Mutex<State>>,
45 : }
46 :
47 : impl FakeS3 {
48 26 : pub async fn start() -> FakeS3 {
49 26 : let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
50 26 : let endpoint = format!("http://{}", listener.local_addr().unwrap());
51 26 : let state = Arc::new(Mutex::new(State::default()));
52 26 : let server_state = state.clone();
53 26 : tokio::spawn(async move {
54 71 : while let Ok((stream, _)) = listener.accept().await {
55 50 : tokio::spawn(serve_connection(stream, server_state.clone()));
56 50 : }
57 0 : });
58 26 : FakeS3 { endpoint, state }
59 26 : }
60 :
61 9 : pub fn endpoint(&self) -> String {
62 9 : self.endpoint.clone()
63 9 : }
64 :
65 17 : pub fn client(&self) -> aws_sdk_s3::Client {
66 17 : RemoteAccess {
67 17 : r2_endpoint: self.endpoint.clone(),
68 17 : r2_access_key_id: "key".to_string(),
69 17 : r2_secret_access_key: "secret".to_string(),
70 17 : }.r2_client()
71 17 : }
72 :
73 9 : pub fn put_objects(&self, keys: Vec<String>) {
74 9 : let mut state = self.state.lock().unwrap();
75 4039 : for key in keys {
76 4039 : state.objects.insert(key, Vec::new());
77 4039 : }
78 9 : }
79 :
80 3 : pub fn set_page_size(&self, size: usize) {
81 3 : self.state.lock().unwrap().page_size = Some(size);
82 3 : }
83 :
84 1 : pub fn deny_listing(&self) {
85 1 : self.state.lock().unwrap().deny_listing = true;
86 1 : }
87 :
88 1 : pub fn deny_deletes(&self) {
89 1 : self.state.lock().unwrap().deny_deletes = true;
90 1 : }
91 :
92 3 : pub fn deny_upload(&self, key: &str) {
93 3 : self.state.lock().unwrap().denied_uploads.insert(key.to_string());
94 3 : }
95 :
96 1 : pub fn keep_object(&self, key: &str) {
97 1 : self.state.lock().unwrap().kept_objects.insert(key.to_string());
98 1 : }
99 :
100 : /// The keys in the bucket, sorted
101 12 : pub fn object_keys(&self) -> Vec<String> {
102 12 : self.state.lock().unwrap().objects.keys().cloned().collect()
103 12 : }
104 :
105 18 : pub fn object(&self, key: &str) -> Option<Vec<u8>> {
106 18 : self.state.lock().unwrap().objects.get(key).cloned()
107 18 : }
108 :
109 20 : pub fn requests(&self) -> Vec<Request> {
110 20 : self.state.lock().unwrap().requests.clone()
111 20 : }
112 : }
113 :
114 50 : async fn serve_connection(mut stream: TcpStream, state: Arc<Mutex<State>>) {
115 50 : let mut buffer = Vec::new();
116 : loop {
117 251 : let header_end = loop {
118 229497 : if let Some(position) = buffer.windows(4).position(|window| window == b"\r\n\r\n") {
119 251 : break position;
120 301 : }
121 301 : let mut chunk = [0; 8192];
122 301 : match stream.read(&mut chunk).await {
123 1 : Ok(0) | Err(_) => return,
124 251 : Ok(read) => buffer.extend_from_slice(&chunk[..read]),
125 : }
126 : };
127 251 : let head = String::from_utf8_lossy(&buffer[..header_end]).to_string();
128 251 : buffer.drain(..header_end + 4);
129 :
130 251 : let mut lines = head.split("\r\n");
131 251 : let mut request_line = lines.next().unwrap().split(' ');
132 251 : let method = request_line.next().unwrap().to_string();
133 251 : let target = request_line.next().unwrap().to_string();
134 251 : let headers: BTreeMap<String, String> = lines
135 2769 : .filter_map(|line| line.split_once(':'))
136 2769 : .map(|(name, value)| (name.trim().to_lowercase(), value.trim().to_string()))
137 251 : .collect();
138 :
139 251 : if headers.contains_key("expect") {
140 0 : let _ = stream.write_all(b"HTTP/1.1 100 Continue\r\n\r\n").await;
141 251 : }
142 :
143 251 : let raw = if headers.get("transfer-encoding").is_some_and(|value| value.contains("chunked")) {
144 : loop {
145 0 : if let Some(length) = chunked_length(&buffer) {
146 0 : break buffer.drain(..length).collect::<Vec<u8>>();
147 0 : }
148 0 : let mut chunk = [0; 65536];
149 0 : match stream.read(&mut chunk).await {
150 0 : Ok(0) | Err(_) => return,
151 0 : Ok(read) => buffer.extend_from_slice(&chunk[..read]),
152 : }
153 : }
154 : } else {
155 251 : let length: usize = headers.get("content-length").map_or(0, |value| value.parse().unwrap());
156 340 : while buffer.len() < length {
157 89 : let mut chunk = [0; 65536];
158 89 : match stream.read(&mut chunk).await {
159 0 : Ok(0) | Err(_) => return,
160 89 : Ok(read) => buffer.extend_from_slice(&chunk[..read]),
161 : }
162 : }
163 251 : buffer.drain(..length).collect()
164 : };
165 251 : let body = if headers.get("content-encoding").is_some_and(|value| value.contains("aws-chunked"))
166 251 : || headers.get("transfer-encoding").is_some_and(|value| value.contains("chunked")) {
167 0 : decode_chunks(&raw)
168 : } else {
169 251 : raw
170 : };
171 :
172 251 : let (path, query) = target.split_once('?').unwrap_or((&target, ""));
173 251 : let query: BTreeMap<String, String> = query
174 251 : .split('&')
175 389 : .filter(|pair| !pair.is_empty())
176 389 : .map(|pair| {
177 389 : let (name, value) = pair.split_once('=').unwrap_or((pair, ""));
178 389 : (percent_decode(name), percent_decode(value))
179 389 : })
180 251 : .collect();
181 251 : let relative = path.strip_prefix("/dashql-get").unwrap_or_else(|| panic!("request outside of the bucket: {path}"));
182 251 : let key = percent_decode(relative.strip_prefix('/').unwrap_or(relative));
183 :
184 251 : let request = Request { method, key, query, headers, body };
185 251 : let (status, response) = respond(&mut state.lock().unwrap(), &request);
186 251 : state.lock().unwrap().requests.push(request);
187 251 : let reply = format!(
188 : "HTTP/1.1 {status}\r\nContent-Type: application/xml\r\nETag: \"etag\"\r\nContent-Length: {}\r\n\r\n{response}",
189 251 : response.len(),
190 : );
191 251 : if stream.write_all(reply.as_bytes()).await.is_err() {
192 0 : return;
193 251 : }
194 : }
195 1 : }
196 :
197 : const ACCESS_DENIED: (&str, &str) = (
198 : "403 Forbidden",
199 : r#"<?xml version="1.0"?><Error><Code>AccessDenied</Code><Message>Access Denied</Message></Error>"#,
200 : );
201 :
202 251 : fn respond(state: &mut State, request: &Request) -> (&'static str, String) {
203 : const OK: &str = "200 OK";
204 251 : match (request.method.as_str(), request.key.as_str()) {
205 251 : ("GET", "") => {
206 16 : if state.deny_listing {
207 1 : return (ACCESS_DENIED.0, ACCESS_DENIED.1.to_string());
208 15 : }
209 15 : let prefix = request.query.get("prefix").cloned().unwrap_or_default();
210 15 : let start: usize = request.query.get("continuation-token").map_or(0, |token| token.parse().unwrap());
211 10569 : let matching: Vec<&String> = state.objects.keys().filter(|key| key.starts_with(&prefix)).collect();
212 15 : let end = state.page_size.map_or(matching.len(), |size| (start + size).min(matching.len()));
213 4034 : let contents: String = matching[start..end].iter().map(|key| format!("<Contents><Key>{key}</Key></Contents>")).collect();
214 15 : let next = if end < matching.len() {
215 6 : format!("<IsTruncated>true</IsTruncated><NextContinuationToken>{end}</NextContinuationToken>")
216 : } else {
217 9 : "<IsTruncated>false</IsTruncated>".to_string()
218 : };
219 15 : (OK, format!(r#"<?xml version="1.0"?><ListBucketResult><Name>dashql-get</Name>{next}{contents}</ListBucketResult>"#))
220 : }
221 235 : ("POST", "") => {
222 8 : if state.deny_deletes {
223 1 : return (ACCESS_DENIED.0, ACCESS_DENIED.1.to_string());
224 7 : }
225 7 : let body = String::from_utf8_lossy(&request.body).to_string();
226 7 : let mut result = String::new();
227 2516 : for part in body.split("<Key>").skip(1) {
228 2516 : let key = part.split("</Key>").next().unwrap();
229 2516 : if state.kept_objects.contains(key) {
230 1 : result.push_str(&format!("<Error><Key>{key}</Key><Code>AccessDenied</Code><Message>denied</Message></Error>"));
231 2515 : } else {
232 2515 : state.objects.remove(key);
233 2515 : result.push_str(&format!("<Deleted><Key>{key}</Key></Deleted>"));
234 2515 : }
235 : }
236 7 : (OK, format!(r#"<?xml version="1.0"?><DeleteResult>{result}</DeleteResult>"#))
237 : }
238 227 : (method, key) if state.denied_uploads.contains(key) && (method == "PUT" || method == "POST") => {
239 3 : (ACCESS_DENIED.0, ACCESS_DENIED.1.to_string())
240 : }
241 224 : ("POST", key) if request.query.contains_key("uploads") => {
242 58 : (OK, format!(r#"<?xml version="1.0"?><InitiateMultipartUploadResult><Bucket>dashql-get</Bucket><Key>{key}</Key><UploadId>upload</UploadId></InitiateMultipartUploadResult>"#))
243 : }
244 109 : ("PUT", _) if request.query.contains_key("partNumber") => (OK, String::new()),
245 57 : ("POST", key) if request.query.contains_key("uploadId") => {
246 57 : let parts = state.requests.iter()
247 523 : .filter(|part| part.key == key && part.query.contains_key("partNumber"))
248 58 : .flat_map(|part| part.body.clone())
249 57 : .collect();
250 57 : state.objects.insert(key.to_string(), parts);
251 57 : (OK, format!(r#"<?xml version="1.0"?><CompleteMultipartUploadResult><Key>{key}</Key><ETag>"etag"</ETag></CompleteMultipartUploadResult>"#))
252 : }
253 51 : ("PUT", key) => {
254 51 : state.objects.insert(key.to_string(), request.body.clone());
255 51 : (OK, String::new())
256 : }
257 0 : _ => panic!("unexpected request: {} {}", request.method, request.key),
258 : }
259 251 : }
260 :
261 : /// The length of a complete chunked body at the begin of the buffer
262 0 : fn chunked_length(buffer: &[u8]) -> Option<usize> {
263 0 : let mut position = 0;
264 : loop {
265 0 : let line_end = position + buffer[position..].windows(2).position(|window| window == b"\r\n")?;
266 0 : let line = std::str::from_utf8(&buffer[position..line_end]).ok()?;
267 0 : let size = usize::from_str_radix(line.split(';').next()?.trim(), 16).ok()?;
268 0 : position = line_end + 2;
269 0 : if size == 0 {
270 : // The trailers end with an empty line
271 0 : let end = position + buffer[position..].windows(2).position(|window| window == b"\r\n")?;
272 0 : if end == position {
273 0 : return Some(end + 2);
274 0 : }
275 0 : let trailers_end = position + buffer[position..].windows(4).position(|window| window == b"\r\n\r\n")?;
276 0 : return Some(trailers_end + 4);
277 0 : }
278 0 : if buffer.len() < position + size + 2 {
279 0 : return None;
280 0 : }
281 0 : position += size + 2;
282 : }
283 0 : }
284 :
285 : /// Remove the chunk framing of a body
286 0 : fn decode_chunks(raw: &[u8]) -> Vec<u8> {
287 0 : let mut decoded = Vec::new();
288 0 : let mut position = 0;
289 0 : while let Some(offset) = raw[position..].windows(2).position(|window| window == b"\r\n") {
290 0 : let line = String::from_utf8_lossy(&raw[position..position + offset]).to_string();
291 0 : let Ok(size) = usize::from_str_radix(line.split(';').next().unwrap().trim(), 16) else { break };
292 0 : position += offset + 2;
293 0 : if size == 0 {
294 0 : break;
295 0 : }
296 0 : decoded.extend_from_slice(&raw[position..position + size]);
297 0 : position += size + 2;
298 : }
299 0 : decoded
300 0 : }
301 :
302 1029 : fn percent_decode(text: &str) -> String {
303 1029 : let bytes = text.as_bytes();
304 1029 : let mut decoded = Vec::with_capacity(bytes.len());
305 1029 : let mut i = 0;
306 14719 : while i < bytes.len() {
307 13690 : if bytes[i] == b'%' && i + 2 < bytes.len() {
308 16 : if let Ok(value) = u8::from_str_radix(&text[i + 1..i + 3], 16) {
309 16 : decoded.push(value);
310 16 : i += 3;
311 16 : continue;
312 0 : }
313 13674 : }
314 13674 : decoded.push(bytes[i]);
315 13674 : i += 1;
316 : }
317 1029 : String::from_utf8(decoded).unwrap()
318 1029 : }
319 :
320 : /// A directory that is removed at the end of a test
321 : pub struct TempDir(PathBuf);
322 :
323 : impl TempDir {
324 41 : pub fn new() -> TempDir {
325 41 : let path = std::env::temp_dir().join(format!("dashql-pack-test-{}", uuid::Uuid::new_v4()));
326 41 : std::fs::create_dir_all(&path).unwrap();
327 41 : TempDir(path)
328 41 : }
329 :
330 33 : pub fn path(&self) -> &Path {
331 33 : &self.0
332 33 : }
333 :
334 181 : pub fn write(&self, name: &str, content: &[u8]) -> PathBuf {
335 181 : let path = self.0.join(name);
336 181 : std::fs::write(&path, content).unwrap();
337 181 : path
338 181 : }
339 : }
340 :
341 : impl Drop for TempDir {
342 41 : fn drop(&mut self) {
343 41 : let _ = std::fs::remove_dir_all(&self.0);
344 41 : }
345 : }
346 :
347 : /// Create a repository with a tagged commit and a number of commits after it
348 12 : pub fn git_repository(path: &Path, tag: &str, commits_after_tag: usize) -> git2::Repository {
349 12 : let repo = git2::Repository::init(path).unwrap();
350 12 : let signature = git2::Signature::now("test", "test@example.com").unwrap();
351 23 : let commit = |parent: Option<git2::Oid>| {
352 23 : let tree_id = repo.treebuilder(None).unwrap().write().unwrap();
353 23 : let tree = repo.find_tree(tree_id).unwrap();
354 23 : let parents: Vec<git2::Commit> = parent.iter().map(|id| repo.find_commit(*id).unwrap()).collect();
355 23 : let parents: Vec<&git2::Commit> = parents.iter().collect();
356 23 : repo.commit(Some("HEAD"), &signature, &signature, "commit", &tree, &parents).unwrap()
357 23 : };
358 12 : let mut head = commit(None);
359 12 : repo.tag_lightweight(tag, &repo.find_object(head, None).unwrap(), false).unwrap();
360 12 : for _ in 0..commits_after_tag {
361 11 : head = commit(Some(head));
362 11 : }
363 12 : repo
364 12 : }
|
