use axum::extract::{Path, Request, State};
use axum::http::StatusCode;
use axum::middleware::Next;
use axum::response::Response;
use kellnr_appstate::{CrateIoStorageState, CratesIoPrefetchSenderState, SettingsState};
use kellnr_common::cratesio_downloader::{CLIENT, download_crate};
use kellnr_common::cratesio_prefetch_msg::{CratesioPrefetchMsg, DownloadData};
use kellnr_common::original_name::OriginalName;
use kellnr_common::version::Version;
use kellnr_error::api_error::ApiResult;
use reqwest::Url;
use tracing::{error, trace, warn};
use crate::registry_error::RegistryError;
use crate::search_params::SearchParams;
pub async fn cratesio_enabled(
State(settings): SettingsState,
request: Request,
next: Next,
) -> Result<Response, StatusCode> {
if settings.proxy.enabled {
Ok(next.run(request).await)
} else {
Err(StatusCode::NOT_FOUND)
}
}
pub async fn search(params: SearchParams) -> ApiResult<String> {
let url = Url::parse(&format!(
"https://crates.io/api/v1/crates?q={}&per_page={}",
params.q, params.per_page.0
))
.map_err(RegistryError::UrlParseError)?;
let response = CLIENT
.get(url)
.send()
.await
.map_err(RegistryError::RequestError)?;
let body = response.text().await.map_err(RegistryError::RequestError)?;
Ok(body)
}
pub async fn download(
Path((name, version)): Path<(OriginalName, Version)>,
State(crate_storage): CrateIoStorageState,
State(sender): CratesIoPrefetchSenderState,
State(settings): SettingsState,
) -> Result<Vec<u8>, StatusCode> {
trace!("Downloading crate: {name} ({version})");
if let Some(file) = crate_storage.get(&name, &version).await {
let msg = DownloadData {
name: name.into(),
version,
};
if let Err(e) = sender.send(CratesioPrefetchMsg::IncDownloadCnt(msg)) {
warn!("Failed to send IncDownloadCnt message: {e}");
}
Ok(file)
} else {
let crate_data = download_crate(&name, &version, &settings.proxy.url).await?;
let _save = crate_storage
.put(&name, &version, crate_data.clone())
.await
.map_err(|e| {
error!("Failed to save crate to disk: {e}");
StatusCode::UNPROCESSABLE_ENTITY
})?;
crate_storage
.get(&name, &version)
.await
.ok_or(StatusCode::NOT_FOUND)
}
}
#[cfg(test)]
mod tests {
use std::path::PathBuf;
use std::sync::Arc;
use axum::body::Body;
use axum::http::Request;
use axum::routing::get;
use axum::{Router, middleware};
use http_body_util::BodyExt;
use kellnr_appstate::AppStateData;
use kellnr_common::util::generate_rand_string;
use kellnr_db::mock::MockDb;
use kellnr_settings::Settings;
use kellnr_storage::cached_crate_storage::DynStorage;
use kellnr_storage::cratesio_crate_storage::CratesIoCrateStorage;
use kellnr_storage::fs_storage::FSStorage;
use tower::ServiceExt;
use super::*;
#[tokio::test]
async fn download_not_existing_package() {
let settings = get_settings();
let kellnr = TestKellnr::new(settings);
let r = kellnr
.client
.clone()
.oneshot(
Request::get("/api/v1/cratesio/does_not_exist/0.1.0/download")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(r.status(), StatusCode::NOT_FOUND);
}
#[tokio::test]
async fn download_invalid_package_name() {
let settings = get_settings();
let kellnr = TestKellnr::new(settings);
let r = kellnr
.client
.clone()
.oneshot(
Request::get("/api/v1/cratesio/-invalid_name/0.1.0/download")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(r.status(), StatusCode::NOT_FOUND);
}
#[tokio::test]
async fn download_not_existing_version() {
let settings = get_settings();
let kellnr = TestKellnr::new(settings);
let r = kellnr
.client
.clone()
.oneshot(
Request::get("/api/v1/cratesio/test-lib/99.1.0/download")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(r.status(), StatusCode::NOT_FOUND);
}
#[tokio::test]
async fn download_invalid_package_version() {
let settings = get_settings();
let kellnr = TestKellnr::new(settings);
let r = kellnr
.client
.clone()
.oneshot(
Request::get("/api/v1/cratesio/invalid_version/0.a.0/download")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(r.status(), StatusCode::NOT_FOUND);
}
#[tokio::test]
async fn download_valid_package() {
let settings = get_settings();
let kellnr = TestKellnr::new(settings);
let r = kellnr
.client
.clone()
.oneshot(
Request::get("/api/v1/cratesio/adler/1.0.2/download")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(r.status(), StatusCode::OK);
let body = r.into_body().collect().await.unwrap().to_bytes();
assert_eq!(12778, body.len());
}
#[tokio::test]
async fn cratesio_disabled_returns_404() {
let mut settings = get_settings();
settings.proxy.enabled = false;
let kellnr = TestKellnr::new(settings);
let r = kellnr
.client
.clone()
.oneshot(
Request::get("/api/v1/cratesio/adler/1.0.2/download")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(r.status(), StatusCode::NOT_FOUND);
}
struct TestKellnr {
path: PathBuf,
client: Router,
}
fn get_settings() -> Settings {
Settings {
registry: kellnr_settings::Registry {
data_dir: "/tmp/".to_string() + &generate_rand_string(10),
session_age_seconds: 10,
..kellnr_settings::Registry::default()
},
proxy: kellnr_settings::Proxy {
enabled: true,
..kellnr_settings::Proxy::default()
},
..Settings::default()
}
}
impl TestKellnr {
fn new(settings: Settings) -> Self {
std::fs::create_dir_all(&settings.registry.data_dir).unwrap();
TestKellnr {
path: PathBuf::from(&settings.registry.data_dir),
client: app(settings),
}
}
}
impl Drop for TestKellnr {
fn drop(&mut self) {
rm_rf::remove(&self.path).expect("Cannot remove TestKellnr");
}
}
fn app(settings: Settings) -> Router {
let storage = Box::new(FSStorage::new(&settings.crates_io_path()).unwrap()) as DynStorage;
let cs = CratesIoCrateStorage::new(&settings, storage);
let mut db = MockDb::new();
db.expect_increase_cached_download_counter()
.returning(|_, _| Ok(()));
let state = AppStateData {
settings: settings.into(),
cratesio_storage: cs.into(),
db: Arc::<MockDb>::new(db),
..kellnr_appstate::test_state()
};
let routes = Router::new()
.route("/", get(search))
.route("/{package}/{version}/download", get(download))
.route_layer(middleware::from_fn_with_state(
state.clone(),
cratesio_enabled,
));
Router::new()
.nest("/api/v1/cratesio", routes)
.with_state(state)
}
}