From ecced7bc940e2beb050cb270ed64baf6f24c9eaf Mon Sep 17 00:00:00 2001 From: Zhaofeng Li Date: Fri, 19 May 2023 02:48:52 -0600 Subject: [PATCH] Refactor --- nix-actions-cache/src/api.rs | 112 ++++++++++++++++++++++++++++++++++ nix-actions-cache/src/main.rs | 108 +------------------------------- 2 files changed, 115 insertions(+), 105 deletions(-) create mode 100644 nix-actions-cache/src/api.rs diff --git a/nix-actions-cache/src/api.rs b/nix-actions-cache/src/api.rs new file mode 100644 index 0000000..1e94a35 --- /dev/null +++ b/nix-actions-cache/src/api.rs @@ -0,0 +1,112 @@ +//! Binary Cache API. + +use std::io; + +use axum::{ + extract::{BodyStream, Extension, Path}, + response::Redirect, + routing::{get, put}, + Router, +}; +use tokio_stream::StreamExt; +use tokio_util::io::StreamReader; + +use super::State; +use crate::error::{Error, Result}; + +pub fn get_router() -> Router { + Router::new() + .route("/nix-cache-info", get(get_nix_cache_info)) + // .narinfo + .route("/:path", get(get_narinfo)) + .route("/:path", put(put_narinfo)) + // .nar + .route("/nar/:path", get(get_nar)) + .route("/nar/:path", put(put_nar)) +} + +async fn get_nix_cache_info() -> &'static str { + // TODO: Make StoreDir configurable + r#"WantMassQuery: 1 +StoreDir: /nix/store +Priority: 41 +"# +} + +async fn get_narinfo( + Extension(state): Extension, + Path(path): Path, +) -> Result { + let components: Vec<&str> = path.splitn(2, '.').collect(); + + if components.len() != 2 { + return Err(Error::NotFound); + } + + if components[1] != "narinfo" { + return Err(Error::NotFound); + } + + let store_path_hash = components[0].to_string(); + let key = format!("{}.narinfo", store_path_hash); + + if let Some(url) = state.api.get_file_url(&[&key]).await? { + return Ok(Redirect::temporary(&url)); + } + + if let Some(upstream) = &state.upstream { + Ok(Redirect::temporary(&format!("{}/{}", upstream, path))) + } else { + Err(Error::NotFound) + } +} +async fn put_narinfo( + Extension(state): Extension, + Path(path): Path, + body: BodyStream, +) -> Result<()> { + let components: Vec<&str> = path.splitn(2, '.').collect(); + + if components.len() != 2 { + return Err(Error::BadRequest); + } + + if components[1] != "narinfo" { + return Err(Error::BadRequest); + } + + let store_path_hash = components[0].to_string(); + let key = format!("{}.narinfo", store_path_hash); + let allocation = state.api.allocate_file_with_random_suffix(&key).await?; + let stream = StreamReader::new( + body.map(|r| r.map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))), + ); + state.api.upload_file(allocation, stream).await?; + + Ok(()) +} + +async fn get_nar(Extension(state): Extension, Path(path): Path) -> Result { + if let Some(url) = state.api.get_file_url(&[&path]).await? { + return Ok(Redirect::temporary(&url)); + } + + if let Some(upstream) = &state.upstream { + Ok(Redirect::temporary(&format!("{}/nar/{}", upstream, path))) + } else { + Err(Error::NotFound) + } +} +async fn put_nar( + Extension(state): Extension, + Path(path): Path, + body: BodyStream, +) -> Result<()> { + let allocation = state.api.allocate_file_with_random_suffix(&path).await?; + let stream = StreamReader::new( + body.map(|r| r.map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))), + ); + state.api.upload_file(allocation, stream).await?; + + Ok(()) +} diff --git a/nix-actions-cache/src/main.rs b/nix-actions-cache/src/main.rs index 1b11706..c7ea6f1 100644 --- a/nix-actions-cache/src/main.rs +++ b/nix-actions-cache/src/main.rs @@ -13,26 +13,18 @@ deny(unused_imports, unused_mut, unused_variables,) )] +mod api; mod error; -use std::io; use std::net::SocketAddr; use std::path::PathBuf; use std::sync::Arc; -use axum::{ - extract::{BodyStream, Extension, Path}, - response::Redirect, - routing::{get, put}, - Router, -}; +use axum::{extract::Extension, routing::get, Router}; use clap::Parser; use tokio::fs; -use tokio_stream::StreamExt; -use tokio_util::io::StreamReader; use tracing_subscriber::EnvFilter; -use error::{Error, Result}; use gha_cache::{Api, Credentials}; type State = Arc; @@ -105,15 +97,7 @@ async fn main() { upstream: args.upstream, }); - let app = Router::new() - .route("/", get(root)) - .route("/nix-cache-info", get(get_nix_cache_info)) - // .narinfo - .route("/:path", get(get_narinfo)) - .route("/:path", put(put_narinfo)) - // .nar - .route("/nar/:path", get(get_nar)) - .route("/nar/:path", put(put_nar)); + let app = Router::new().route("/", get(root)).merge(api::get_router()); #[cfg(debug_assertions)] let app = app @@ -153,89 +137,3 @@ async fn dump_api_stats( async fn root() -> &'static str { "cache the world 🚀" } - -async fn get_nix_cache_info() -> &'static str { - // TODO: Make StoreDir configurable - r#"WantMassQuery: 1 -StoreDir: /nix/store -Priority: 41 -"# -} - -async fn get_narinfo( - Extension(state): Extension, - Path(path): Path, -) -> Result { - let components: Vec<&str> = path.splitn(2, '.').collect(); - - if components.len() != 2 { - return Err(Error::NotFound); - } - - if components[1] != "narinfo" { - return Err(Error::NotFound); - } - - let store_path_hash = components[0].to_string(); - let key = format!("{}.narinfo", store_path_hash); - - if let Some(url) = state.api.get_file_url(&[&key]).await? { - return Ok(Redirect::temporary(&url)); - } - - if let Some(upstream) = &state.upstream { - Ok(Redirect::temporary(&format!("{}/{}", upstream, path))) - } else { - Err(Error::NotFound) - } -} -async fn put_narinfo( - Extension(state): Extension, - Path(path): Path, - body: BodyStream, -) -> Result<()> { - let components: Vec<&str> = path.splitn(2, '.').collect(); - - if components.len() != 2 { - return Err(Error::BadRequest); - } - - if components[1] != "narinfo" { - return Err(Error::BadRequest); - } - - let store_path_hash = components[0].to_string(); - let key = format!("{}.narinfo", store_path_hash); - let allocation = state.api.allocate_file_with_random_suffix(&key).await?; - let stream = StreamReader::new( - body.map(|r| r.map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))), - ); - state.api.upload_file(allocation, stream).await?; - - Ok(()) -} - -async fn get_nar(Extension(state): Extension, Path(path): Path) -> Result { - if let Some(url) = state.api.get_file_url(&[&path]).await? { - return Ok(Redirect::temporary(&url)); - } - - if let Some(upstream) = &state.upstream { - Ok(Redirect::temporary(&format!("{}/nar/{}", upstream, path))) - } else { - Err(Error::NotFound) - } -} -async fn put_nar( - Extension(state): Extension, - Path(path): Path, - body: BodyStream, -) -> Result<()> { - let allocation = state.api.allocate_file_with_random_suffix(&path).await?; - let stream = StreamReader::new( - body.map(|r| r.map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))), - ); - state.api.upload_file(allocation, stream).await?; - - Ok(()) -}