1use crate::{parse_digest, Manifest};
2use hyper::body::HttpBody;
3use regex::Regex;
4use saphir::hyper::body::Buf;
5use saphir::prelude::*;
6use serde::Deserialize;
7use slog_scope::{debug, error};
8use std::collections::HashMap;
9use std::fs::create_dir_all;
10use std::io;
11use std::io::Write;
12use std::path::Path;
13use tokio_02::io::AsyncWriteExt;
14
15const REPOSITORY: &str = "repository";
16const IMAGE_NAME: &str = "image_name";
17const DIGEST: &str = "digest";
18const TAG: &str = "tag";
19pub const ARTIFACTS_DIR: &str = "artifacts";
20const ARTIFACTS_CONTENT: &str = "content.yaml";
21
22pub const BLOB_PATH: &str = "/registry/v2/<repository>/<image_name>/blobs/<digest>";
23pub const BLOB_GET_LOCATION_PATH: &str = "/registry/v2/<repository>/<image_name>/blobs/uploads";
24pub const UPLOAD_BLOB_PATH: &str = "/registry/<repository>/<image_name>";
25pub const MANIFEST_PATH: &str = "/registry/v2/<repository>/<image_name>/manifests/<tag>";
26
27pub const BLOB_EXIST_ENDPOINT: &str = "is_blob_exist";
28pub const BLOB_GET_LOCATION_ENDPOINT: &str = "get_blob_location";
29pub const BLOB_UPLOAD_ENDPOINT: &str = "save_blob";
30pub const BLOB_DOWNLOAD_ENDPOINT: &str = "get_blob";
31pub const MANIFEST_EXIST_ENDPOINT: &str = "is_manifest_exist";
32pub const MANIFEST_UPLOAD_ENDPOINT: &str = "save_manifest";
33pub const MANIFEST_DOWNLOAD_ENDPOINT: &str = "get_manifest";
34
35const CONTENT_TYPE_HEADER: &str = "content-type";
36const ACCEPT_HEADER: &str = "accept";
37
38pub struct SogarCustomResponse {
39 headers: HashMap<String, String>,
40 status: StatusCode,
41 file: Option<File>,
42}
43
44impl SogarCustomResponse {
45 pub fn new(status: StatusCode) -> Self {
46 SogarCustomResponse {
47 headers: HashMap::new(),
48 status,
49 file: None,
50 }
51 }
52
53 pub fn header(&mut self, key: &str, value: &str) {
54 self.headers.insert(key.to_string(), value.to_string());
55 }
56
57 pub fn file(&mut self, file: File) {
58 self.file = Some(file);
59 }
60}
61
62impl Responder for SogarCustomResponse {
63 fn respond_with_builder(self, builder: Builder, _ctx: &HttpContext) -> Builder {
64 let mut builder_copy = builder;
65 for (key, val) in self.headers {
66 builder_copy = builder_copy.header(&key, val);
67 }
68
69 builder_copy = builder_copy.status(self.status);
70
71 if let Some(file) = self.file {
72 builder_copy = builder_copy.file(file);
73 }
74
75 builder_copy
76 }
77}
78
79pub struct SogarController {
80 _priv: (),
81}
82
83impl SogarController {
84 pub fn new(registry: &str, image_name: &str) -> Self {
85 use std::fs::File;
86
87 let path = Path::new(registry).join(image_name);
88 if let Err(e) = create_dir_all(path.join(ARTIFACTS_DIR)) {
89 error!("Failed to create registry! {}", e);
90 }
91
92 let content = path.join(ARTIFACTS_CONTENT);
93 if !content.exists() {
94 match File::create(path.join(ARTIFACTS_CONTENT)) {
95 Ok(mut file) => {
96 if let Err(e) = writeln!(file, "artifacts:") {
97 error!("Couldn't write to file: {}", e);
98 }
99 }
100 Err(e) => {
101 error!("Failed to create file with artifacts content! {}", e);
102 }
103 }
104 }
105
106 Self { _priv: () }
107 }
108}
109
110#[controller(name = "registry")]
111impl SogarController {
112 #[head("/v2/<repository>/<image_name>/blobs/<digest>")]
113 async fn is_blob_exist(&self, req: Request) -> impl Responder {
114 debug!("Head request for blob");
115 let map = req.captures();
116
117 if let (Some(repository), Some(image_name), Some(digest)) =
118 (map.get(REPOSITORY), map.get(IMAGE_NAME), map.get(DIGEST))
119 {
120 if let Some(digest) = parse_digest(digest) {
121 let path = Path::new(repository)
122 .join(image_name)
123 .join(ARTIFACTS_DIR)
124 .join(&digest.digest_type)
125 .join(&digest.value);
126
127 if path.exists() {
128 let mut response = SogarCustomResponse::new(StatusCode::OK);
129 response.header(
130 "Docker-Content-Digest",
131 format!("{}:{}", digest.digest_type, digest.value).as_str(),
132 );
133 return response;
134 }
135 }
136 }
137
138 SogarCustomResponse::new(StatusCode::NOT_FOUND)
139 }
140
141 #[post("/v2/<repository>/<image_name>/blobs/uploads/")]
142 async fn get_blob_location(&self, req: Request) -> impl Responder {
143 debug!("Post request for blob");
144 let map = req.captures();
145
146 if let (Some(repository), Some(image_name)) = (map.get(REPOSITORY), map.get(IMAGE_NAME)) {
147 let path = Path::new(repository).join(image_name);
148 if path.exists() {
149 let mut response = SogarCustomResponse::new(StatusCode::ACCEPTED);
150 response.header("Location", format!("/{}/{}", repository, image_name).as_str());
151 return response;
152 }
153 }
154
155 SogarCustomResponse::new(StatusCode::NOT_FOUND)
156 }
157
158 #[put("/<repository>/<image_name>")]
160 async fn save_blob(&self, mut req: Request) -> impl Responder {
161 debug!("Put request for blob");
162
163 let body: hyper::Body = req.body_mut().take().into();
164
165 let digest = parse_digest_uri(req.uri());
166 let map = req.captures();
167
168 if let (Some(repository), Some(image_name), Some(digest)) = (map.get(REPOSITORY), map.get(IMAGE_NAME), digest) {
169 let path = Path::new(repository).join(image_name).join(ARTIFACTS_DIR);
170 if path.exists() {
171 if let Some(blob_digest) = parse_digest(&digest) {
172 let path = path.join(blob_digest.digest_type.as_str());
173 if !path.exists() {
174 if let Err(e) = create_dir_all(path.as_path()) {
175 error!("Failed to create directory for saving blob {}", e);
176 return SogarCustomResponse::new(StatusCode::BAD_REQUEST);
177 }
178 }
179
180 let path = path.join(blob_digest.value.as_str());
181
182 if let Err(error_response) = remove_file_if_exists(path.as_path()) {
183 return error_response;
184 }
185
186 if let Err(e) = write_body_to_file(body, path.as_path()).await {
187 error!("Failed to write data to the file {}", e);
188 return SogarCustomResponse::new(StatusCode::BAD_REQUEST);
189 }
190
191 let mut response = SogarCustomResponse::new(StatusCode::CREATED);
192 response.header(
193 "Location",
194 format!("/v2/{}/{}/blobs/{}", repository, image_name, digest).as_str(),
195 );
196
197 return response;
198 }
199 }
200 }
201
202 SogarCustomResponse::new(StatusCode::NOT_FOUND)
203 }
204
205 #[head("/v2/<repository>/<image_name>/manifests/<tag>")]
206 async fn is_manifest_exist(&self, req: Request) -> (StatusCode, ()) {
207 debug!("Head request for manifest");
208
209 let map = req.captures();
210
211 if let (Some(repository), Some(image_name), Some(tag)) =
212 (map.get(REPOSITORY), map.get(IMAGE_NAME), map.get(TAG))
213 {
214 let path = Path::new(repository).join(image_name).join(ARTIFACTS_DIR).join(tag);
215 if path.exists() {
216 return (StatusCode::OK, ());
217 }
218 }
219
220 (StatusCode::NOT_FOUND, ())
221 }
222
223 #[put("/v2/<repository>/<image_name>/manifests/<tag>")]
224 async fn save_manifest(&self, mut req: Request) -> impl Responder {
225 debug!("Put request for manifest");
226 let body: hyper::Body = req.body_mut().take().into();
227
228 let map = req.captures();
229 let headers = req.headers();
230
231 if let (Some(repository), Some(image_name), Some(tag)) =
232 (map.get(REPOSITORY), map.get(IMAGE_NAME), map.get(TAG))
233 {
234 let image_path = Path::new(repository).join(image_name);
235 let path = image_path.join(ARTIFACTS_DIR);
236 if path.exists() {
237 let path = path.join(tag);
238
239 if let Err(error_response) = remove_file_if_exists(path.as_path()) {
240 return error_response;
241 }
242
243 if let Err(e) = write_body_to_file(body, path.as_path()).await {
244 error!("Failed to write data to the file {}", e);
245 return SogarCustomResponse::new(StatusCode::BAD_REQUEST);
246 }
247
248 let mut manifest_mime_type = None;
249 if headers.contains_key(CONTENT_TYPE_HEADER) {
250 manifest_mime_type = headers
251 .get(CONTENT_TYPE_HEADER)
252 .and_then(|header| header.to_str().map_or(None, |result| Some(result.to_string())));
253 }
254
255 add_artifacts_info(tag, manifest_mime_type, image_path.as_path());
256
257 let mut response = SogarCustomResponse::new(StatusCode::CREATED);
258 response.header(
259 "Location",
260 format!("/v2/{}/{}/manifests/{}", repository, image_name, tag).as_str(),
261 );
262
263 return response;
264 }
265 }
266
267 SogarCustomResponse::new(StatusCode::BAD_REQUEST)
268 }
269
270 #[get("/v2/<repository>/<image_name>/manifests/<tag>")]
271 async fn get_manifest(&self, req: Request) -> (StatusCode, Option<File>) {
272 debug!("Get request for manifest");
273 let map = req.captures();
274 let headers = req.headers();
275
276 if let (Some(repository), Some(image_name), Some(tag)) =
277 (map.get(REPOSITORY), map.get(IMAGE_NAME), map.get(TAG))
278 {
279 let image_path = Path::new(repository).join(image_name);
280 let path = image_path.join(ARTIFACTS_DIR).join(tag);
281 let content_type = read_artifact_info(tag, image_path.as_path());
282
283 if headers.contains_key(ACCEPT_HEADER) {
284 if let (Some(accept_value), Some(content_type)) = (
285 headers.get(ACCEPT_HEADER).and_then(|header| header.to_str().ok()),
286 content_type,
287 ) {
288 if accept_value.contains(&content_type) || accept_value.contains("*/*") {
289 return get_file_if_exists(path.as_path()).await;
290 }
291 }
292 } else {
293 return get_file_if_exists(path.as_path()).await;
294 }
295 }
296
297 (StatusCode::NOT_FOUND, None)
298 }
299
300 #[get("/v2/<repository>/<image_name>/blobs/<digest>")]
301 async fn get_blob(&self, req: Request) -> impl Responder {
302 debug!("Get request for blob");
303
304 let map = req.captures();
305
306 if let (Some(repository), Some(image_name), Some(digest)) =
307 (map.get(REPOSITORY), map.get(IMAGE_NAME), map.get(DIGEST))
308 {
309 if let Some(digest) = parse_digest(digest) {
310 let image_path = Path::new(repository).join(image_name);
311 let path = image_path
312 .join(ARTIFACTS_DIR)
313 .join(&digest.digest_type)
314 .join(&digest.value);
315 let (status_code, file_result) = get_file_if_exists(path.as_path()).await;
316
317 if let Some(file) = file_result {
318 let content_type = read_artifact_info(&digest.value, image_path.as_path());
319
320 let mut response = SogarCustomResponse::new(status_code);
321 response.file(file);
322
323 if let Some(content_type) = content_type {
324 response.header(CONTENT_TYPE_HEADER, content_type.as_str());
325 }
326
327 return response;
328 }
329 }
330 }
331
332 SogarCustomResponse::new(StatusCode::NOT_FOUND)
333 }
334}
335
336async fn write_body_to_file(mut body: hyper::Body, path: &Path) -> io::Result<()> {
337 use tokio_02::fs::File;
338
339 let mut file = File::create(path).await?;
340
341 while let Some(chunk_res) = body.data().await {
342 match chunk_res {
343 Ok(chunk) => {
344 debug!("Got a chunk [len = {}]", chunk.len());
345 file.write_all(chunk.bytes()).await?
346 }
347
348 Err(e) => {
349 return Err(io::Error::new(
350 io::ErrorKind::InvalidData,
351 format!("Failed to read chunk! Error is {}", e),
352 ));
353 }
354 }
355 }
356
357 file.flush().await?;
358
359 Ok(())
360}
361
362fn parse_digest_uri(uri: &Uri) -> Option<String> {
363 if let Some(digest_uri) = uri.query() {
364 let digest_index = 1;
365 let digest_re = Regex::new(r"(.*)=(.*)").unwrap();
366
367 if digest_re.is_match(digest_uri) {
368 let split = digest_uri.split('=');
369 let result = split.into_iter().map(ToString::to_string).collect::<Vec<String>>();
370 return Some(result[digest_index].clone());
371 }
372 }
373
374 None
375}
376
377fn remove_file_if_exists(path: &Path) -> std::result::Result<(), SogarCustomResponse> {
378 if path.exists() {
379 if let Err(e) = std::fs::remove_file(path) {
380 error!("Failed to delete existed file {:?} with error {}", path, e);
381 return Err(SogarCustomResponse::new(StatusCode::BAD_REQUEST));
382 }
383 }
384
385 Ok(())
386}
387
388async fn get_file_if_exists(path: &Path) -> (StatusCode, Option<File>) {
389 if path.exists() && path.to_str().is_some() {
390 if let Ok(file) = File::open(path.to_str().unwrap()).await {
391 return (StatusCode::OK, Some(file));
392 }
393 }
394
395 (StatusCode::NOT_FOUND, None)
396}
397
398pub fn add_artifacts_info(filename: &str, manifest_mime: Option<String>, image_path: &Path) {
399 use std::fs::{File, OpenOptions};
400
401 let content_path = image_path.join(ARTIFACTS_CONTENT);
402 let filepath = image_path.join(ARTIFACTS_DIR).join(&filename);
403
404 let file = File::open(filepath);
405
406 if let Ok(file) = file {
407 let json = serde_json::from_reader(file);
408 if let Err(e) = json {
409 error!("Failed to convert manifest data to json with error: {}", e);
410 return;
411 }
412
413 let artifacts_content_file = OpenOptions::new().write(true).append(true).open(content_path);
414
415 let manifest: Manifest = json.unwrap();
416
417 if let Ok(mut file) = artifacts_content_file {
418 for layer in manifest.layers {
419 if let Some(digest) = parse_digest(&layer.digest) {
420 if let Err(e) = writeln!(file, "{}", format!(" {}: {}", digest.value, layer.media_type)) {
421 error!("Couldn't write to file: {}", e);
422 }
423 }
424 }
425
426 if let Some(manifest_mime) = manifest_mime {
427 if let Err(e) = writeln!(file, "{}", format!(" {}: {}", filename, manifest_mime)) {
428 error!("Couldn't write to file: {}", e);
429 }
430 }
431 }
432 }
433}
434
435fn read_artifact_info(digest_value: &str, image_path: &Path) -> Option<String> {
436 use std::fs::File;
437
438 let content_path = image_path.join(ARTIFACTS_CONTENT);
439 match File::open(&content_path) {
440 Ok(file) => {
441 #[derive(Deserialize)]
442 struct ArtifactsData {
443 artifacts: HashMap<String, String>,
444 }
445
446 let yaml = serde_yaml::from_reader(file);
447 if let Err(e) = yaml {
448 error!("Failed to convert manifest data to yaml with error: {}", e);
449 return None;
450 }
451
452 let blobs_data: ArtifactsData = yaml.unwrap();
453 if blobs_data.artifacts.contains_key(digest_value) {
454 return blobs_data
455 .artifacts
456 .get(digest_value)
457 .map(|mime_type| mime_type.to_string());
458 }
459 }
460 Err(e) => {
461 error!("Content file ({}) can't be opened: {}", content_path.display(), e);
462 }
463 }
464
465 None
466}