pierrelf.com
Rust implementation of OCI Distribution Spec with granular access control
1use std::net::TcpListener;
2use std::path::PathBuf;
3use std::process::{Child, Command};
4use std::thread;
5use std::time::Duration;
6use tempfile::TempDir;
7
8#[allow(dead_code)]
9pub struct TestServer {
10 pub base_url: String,
11 pub host: String,
12 pub port: u16,
13 pub temp_dir: TempDir,
14 pub users_file: PathBuf,
15 process: Option<Child>,
16}
17
18impl TestServer {
19 pub fn new() -> Self {
20 Self::new_with_users(default_test_users())
21 }
22
23 pub fn new_with_users(users_json: serde_json::Value) -> Self {
24 // Find available port
25 let listener = TcpListener::bind("127.0.0.1:0").expect("Failed to bind to random port");
26 let port = listener.local_addr().unwrap().port();
27 drop(listener);
28
29 let host = format!("127.0.0.1:{}", port);
30 let base_url = format!("http://{}", host);
31
32 // Create isolated temp directory
33 let temp_dir = TempDir::new().expect("Failed to create temp dir");
34 let temp_path = temp_dir.path();
35
36 // Setup storage directories
37 std::fs::create_dir_all(temp_path.join("tmp/blobs")).unwrap();
38 std::fs::create_dir_all(temp_path.join("tmp/manifests")).unwrap();
39 std::fs::create_dir_all(temp_path.join("tmp/uploads")).unwrap();
40
41 // Create users.json
42 let users_file = temp_path.join("users.json");
43 std::fs::write(
44 &users_file,
45 serde_json::to_string_pretty(&users_json).unwrap(),
46 )
47 .expect("Failed to write users.json");
48
49 TestServer {
50 base_url,
51 host,
52 port,
53 temp_dir,
54 users_file,
55 process: None,
56 }
57 }
58
59 pub fn start(&mut self) {
60 // Get the workspace root directory
61 let workspace_root = std::env::current_dir().expect("Failed to get current directory");
62
63 // Build if not already built
64 let build_status = Command::new("cargo")
65 .args(["build", "--bin", "grain"])
66 .current_dir(&workspace_root)
67 .status()
68 .expect("Failed to build grain");
69
70 assert!(build_status.success(), "Failed to build grain binary");
71
72 // Path to the binary
73 let binary_path = workspace_root.join("target/debug/grain");
74 assert!(
75 binary_path.exists(),
76 "grain binary not found at {:?}",
77 binary_path
78 );
79
80 // Change to temp directory for storage
81 let temp_path = self.temp_dir.path();
82
83 // Start server process
84 let mut child = Command::new(binary_path)
85 .args([
86 "--host",
87 &self.host,
88 "--users-file",
89 self.users_file.to_str().unwrap(),
90 ])
91 .current_dir(temp_path)
92 .spawn()
93 .expect("Failed to start grain server");
94
95 // Wait for server to be ready
96 let client = reqwest::blocking::Client::new();
97 let url = format!("{}/v2/", self.base_url);
98
99 for _ in 0..50 {
100 thread::sleep(Duration::from_millis(100));
101
102 // Check if process is still running
103 if let Ok(Some(_)) = child.try_wait() {
104 panic!("Server process exited prematurely");
105 }
106
107 // Try to connect
108 if client
109 .get(&url)
110 .basic_auth("admin", Some("admin"))
111 .send()
112 .is_ok()
113 {
114 self.process = Some(child);
115 return;
116 }
117 }
118
119 // Kill the process if startup failed
120 let _ = child.kill();
121 panic!("Server failed to start within timeout");
122 }
123
124 pub fn stop(&mut self) {
125 if let Some(mut process) = self.process.take() {
126 let _ = process.kill();
127 let _ = process.wait();
128 }
129 }
130
131 pub fn client(&self) -> TestClient {
132 TestClient {
133 base_url: self.base_url.clone(),
134 client: reqwest::blocking::Client::new(),
135 }
136 }
137}
138
139impl Drop for TestServer {
140 fn drop(&mut self) {
141 self.stop();
142 }
143}
144
145pub struct TestClient {
146 pub base_url: String,
147 client: reqwest::blocking::Client,
148}
149
150#[allow(dead_code)]
151impl TestClient {
152 pub fn get(&self, path: &str) -> reqwest::blocking::RequestBuilder {
153 self.client.get(format!("{}{}", self.base_url, path))
154 }
155
156 pub fn head(&self, path: &str) -> reqwest::blocking::RequestBuilder {
157 self.client.head(format!("{}{}", self.base_url, path))
158 }
159
160 pub fn post(&self, path: &str) -> reqwest::blocking::RequestBuilder {
161 self.client.post(format!("{}{}", self.base_url, path))
162 }
163
164 pub fn put(&self, path: &str) -> reqwest::blocking::RequestBuilder {
165 self.client.put(format!("{}{}", self.base_url, path))
166 }
167
168 pub fn patch(&self, path: &str) -> reqwest::blocking::RequestBuilder {
169 self.client.patch(format!("{}{}", self.base_url, path))
170 }
171
172 pub fn delete(&self, path: &str) -> reqwest::blocking::RequestBuilder {
173 self.client.delete(format!("{}{}", self.base_url, path))
174 }
175}
176
177pub fn default_test_users() -> serde_json::Value {
178 serde_json::json!({
179 "users": [
180 {
181 "username": "admin",
182 "password": "admin",
183 "permissions": [
184 {
185 "repository": "*",
186 "tag": "*",
187 "actions": ["pull", "push", "delete"]
188 }
189 ]
190 },
191 {
192 "username": "reader",
193 "password": "reader",
194 "permissions": [
195 {
196 "repository": "test/*",
197 "tag": "*",
198 "actions": ["pull"]
199 }
200 ]
201 },
202 {
203 "username": "writer",
204 "password": "writer",
205 "permissions": [
206 {
207 "repository": "test/*",
208 "tag": "*",
209 "actions": ["pull", "push"]
210 }
211 ]
212 },
213 {
214 "username": "limited",
215 "password": "limited",
216 "permissions": [
217 {
218 "repository": "myorg/myrepo",
219 "tag": "v*",
220 "actions": ["pull", "push"]
221 }
222 ]
223 }
224 ]
225 })
226}
227
228pub fn sample_blob() -> Vec<u8> {
229 b"This is a test blob content".to_vec()
230}
231
232pub fn sample_blob_digest() -> String {
233 format!("sha256:{}", sha256::digest("This is a test blob content"))
234}
235
236pub fn sample_manifest() -> serde_json::Value {
237 let blob_digest = sample_blob_digest();
238 serde_json::json!({
239 "schemaVersion": 2,
240 "mediaType": "application/vnd.oci.image.manifest.v1+json",
241 "config": {
242 "mediaType": "application/vnd.oci.image.config.v1+json",
243 "size": 27,
244 "digest": blob_digest
245 },
246 "layers": [
247 {
248 "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
249 "size": 27,
250 "digest": blob_digest
251 }
252 ]
253 })
254}
255
256pub fn sample_manifest_digest(manifest: &serde_json::Value) -> String {
257 let manifest_bytes = serde_json::to_vec(manifest).unwrap();
258 format!("sha256:{}", sha256::digest(&manifest_bytes))
259}
260
261pub fn sample_image_index() -> serde_json::Value {
262 let manifest_digest = sample_manifest_digest(&sample_manifest());
263 serde_json::json!({
264 "schemaVersion": 2,
265 "mediaType": "application/vnd.oci.image.index.v1+json",
266 "manifests": [
267 {
268 "mediaType": "application/vnd.oci.image.manifest.v1+json",
269 "size": 500,
270 "digest": manifest_digest,
271 "platform": {
272 "architecture": "amd64",
273 "os": "linux"
274 }
275 }
276 ]
277 })
278}
279
280#[cfg(test)]
281mod tests {
282 use super::*;
283
284 #[test]
285 fn test_server_lifecycle() {
286 let mut server = TestServer::new();
287 server.start();
288
289 let client = server.client();
290 let resp = client
291 .get("/v2/")
292 .basic_auth("admin", Some("admin"))
293 .send()
294 .unwrap();
295
296 assert_eq!(resp.status(), 200);
297
298 server.stop();
299 }
300
301 #[test]
302 fn test_sample_data() {
303 let blob = sample_blob();
304 assert!(!blob.is_empty());
305
306 let digest = sample_blob_digest();
307 assert!(digest.starts_with("sha256:"));
308
309 let manifest = sample_manifest();
310 assert_eq!(manifest["schemaVersion"], 2);
311
312 let index = sample_image_index();
313 assert_eq!(index["schemaVersion"], 2);
314 }
315}