1  use rocket::{
  2      fs::NamedFile,
  3      http::{ContentType, Header, Status},
  4      request::{self, FromRequest},
  5      Request, State,
  6  };
  7  use sqlx::Pool;
  8  
  9  use crate::{
 10      config::Config,
 11      db::DB,
 12      services::{self, get_manifest_service},
 13  };
 14  
 15  use super::{
 16      CONTENT_LENGTH_HEADER_NAME, CONTENT_TYPE_HEADER_NAME, DOCKER_CONTENT_DIGEST_HEADER_NAME,
 17      LOCATION_HEADER_NAME,
 18  };
 19  
 20  #[derive(Responder, Debug)]
 21  pub struct GetManifestResponseData<'a> {
 22      file: NamedFile,
 23      content_type: ContentType,
 24      docker_digest: Header<'a>,
 25  }
 26  
 27  #[derive(Responder, Debug)]
 28  pub enum GetManifestResponse<'a> {
 29      #[response(status = 200)]
 30      Success(GetManifestResponseData<'a>),
 31      #[response(status = 404)]
 32      FileNotFound(&'a str),
 33      #[response(status = 500)]
 34      Failure(&'a str),
 35  }
 36  
 37  #[get("/v2/<name>/manifests/<reference>")]
 38  pub async fn get_manifest<'a>(
 39      name: &str,
 40      reference: &str,
 41      db_pool: &State<Pool<DB>>,
 42      config: &State<Config>,
 43  ) -> GetManifestResponse<'a> {
 44      match get_manifest_service::find_manifest(
 45          db_pool,
 46          name.to_string(),
 47          reference.to_string(),
 48          config,
 49      )
 50      .await
 51      {
 52          Ok(Some((manifest, blob, file))) => {
 53              println!("Manifest found for {name}/{reference}");
 54              GetManifestResponse::Success(GetManifestResponseData {
 55                  file: file,
 56                  content_type: ContentType::new(
 57                      manifest.content_type_top,
 58                      manifest.content_type_sub,
 59                  ),
 60                  docker_digest: Header::new(DOCKER_CONTENT_DIGEST_HEADER_NAME, blob.digest),
 61              })
 62          }
 63          Ok(None) => {
 64              println!("Failed to find manifest {name}/{reference}");
 65              GetManifestResponse::FileNotFound("File not found")
 66          }
 67          Err(e) => {
 68              error!("Failed to get manifest, err: {e:?}");
 69              GetManifestResponse::Failure("An error occurred")
 70          }
 71      }
 72  }
 73  
 74  #[derive(Responder, Debug)]
 75  pub struct PutManifestResponseData<'a> {
 76      response: &'a str,
 77      location: Header<'a>,
 78      docker_content_digest: Header<'a>,
 79  }
 80  
 81  #[derive(Responder, Debug)]
 82  pub enum PutManifestResponse<'a> {
 83      #[response(status = 201)]
 84      Success(PutManifestResponseData<'a>),
 85      #[response(status = 400)]
 86      BadRequest(String),
 87      #[response(status = 500)]
 88      Failure(&'a str),
 89  }
 90  
 91  pub struct ManifestHeaders {
 92      content_length: usize,
 93      content_type: String,
 94  }
 95  
 96  #[rocket::async_trait]
 97  impl<'r> FromRequest<'r> for ManifestHeaders {
 98      type Error = PutManifestResponse<'r>;
 99  
100      async fn from_request(req: &'r Request<'_>) -> request::Outcome<Self, Self::Error> {
101          let content_length: usize = match req.headers().get_one(CONTENT_LENGTH_HEADER_NAME) {
102              Some(h) => match h.parse() {
103                  Ok(val) => val,
104                  Err(e) => {
105                      error!("Failed to parse content_length header, err: {e:?}");
106                      return request::Outcome::Failure((
107                          Status::BadRequest,
108                          PutManifestResponse::BadRequest(format!(
109                              "Invalid header {CONTENT_LENGTH_HEADER_NAME}"
110                          )),
111                      ));
112                  }
113              },
114              None => {
115                  error!("Missing content_length header");
116                  return request::Outcome::Failure((
117                      Status::BadRequest,
118                      PutManifestResponse::BadRequest(format!(
119                          "Missing header {CONTENT_LENGTH_HEADER_NAME}"
120                      )),
121                  ));
122              }
123          };
124  
125          let content_type = match req.headers().get_one(CONTENT_TYPE_HEADER_NAME) {
126              Some(h) => h.to_string(),
127              None => {
128                  error!("Missing content_type header");
129                  return request::Outcome::Failure((
130                      Status::BadRequest,
131                      PutManifestResponse::BadRequest(format!(
132                          "Missing header {CONTENT_TYPE_HEADER_NAME}"
133                      )),
134                  ));
135              }
136          };
137  
138          request::Outcome::Success(Self {
139              content_length,
140              content_type,
141          })
142      }
143  }
144  
145  #[put("/v2/<name>/manifests/<reference>", data = "<manifest_data>")]
146  pub async fn put_manifest<'a>(
147      name: &str,
148      reference: &str,
149      headers: ManifestHeaders,
150      manifest_data: Vec<u8>,
151      config: &State<Config>,
152      db_pool: &State<Pool<DB>>,
153  ) -> PutManifestResponse<'a> {
154      match services::upload_manifest_service::upload_manifest(
155          name.to_string(),
156          reference.to_string(),
157          headers.content_length,
158          headers.content_type,
159          manifest_data,
160          config,
161          db_pool,
162      )
163      .await
164      {
165          Ok((manifest_id, digest)) => PutManifestResponse::Success(PutManifestResponseData {
166              response: "Upload manifest successful",
167              location: Header::new(LOCATION_HEADER_NAME, format!("/{manifest_id}")),
168              docker_content_digest: Header::new(DOCKER_CONTENT_DIGEST_HEADER_NAME, digest),
169          }),
170          Err(e) => {
171              error!("Failed to upload manifest {e:?}");
172              PutManifestResponse::Failure("Failed to upload manifest")
173          }
174      }
175  }