This commit is contained in:
2023-07-25 20:37:28 +02:00
parent a8ca9c0d2c
commit 680275bd00
2 changed files with 71 additions and 84 deletions

View File

@@ -3,58 +3,9 @@ pub mod views;
pub mod services; pub mod services;
use std::net::SocketAddr; use std::net::SocketAddr;
use axum::{Router, routing::get, middleware::{self, Next}, http::{self, Request}, response::Response};
use axum_prometheus::PrometheusMetricLayer;
use config::CONFIG;
use reqwest::StatusCode;
use views::{download, get_filename};
use tracing::info; use tracing::info;
use tower_http::trace::{TraceLayer, self};
use tracing::Level;
use crate::views::get_router;
async fn auth<B>(req: Request<B>, next: Next<B>) -> Result<Response, StatusCode> {
let auth_header = req.headers()
.get(http::header::AUTHORIZATION)
.and_then(|header| header.to_str().ok());
let auth_header = if let Some(auth_header) = auth_header {
auth_header
} else {
return Err(StatusCode::UNAUTHORIZED);
};
if auth_header != CONFIG.api_key {
return Err(StatusCode::UNAUTHORIZED);
}
Ok(next.run(req).await)
}
async fn get_router() -> Router {
let (prometheus_layer, metric_handle) = PrometheusMetricLayer::pair();
let app_router = Router::new()
.route("/download/:source_id/:remote_id/:file_type", get(download))
.route("/filename/:book_id/:file_type", get(get_filename))
.layer(middleware::from_fn(auth))
.layer(prometheus_layer);
let metric_router = Router::new()
.route("/metrics", get(|| async move { metric_handle.render() }));
Router::new()
.nest("/", app_router)
.nest("/", metric_router)
.layer(
TraceLayer::new_for_http()
.make_span_with(trace::DefaultMakeSpan::new()
.level(Level::INFO))
.on_response(trace::DefaultOnResponse::new()
.level(Level::INFO)),
)
}
#[tokio::main] #[tokio::main]

View File

@@ -1,40 +1,36 @@
use axum::{ use axum::{
body::StreamBody, body::StreamBody,
extract::Path, extract::Path,
http::{header, HeaderMap, StatusCode, header::AUTHORIZATION}, http::{header, StatusCode},
response::{IntoResponse, AppendHeaders}, response::{AppendHeaders, IntoResponse},
}; };
use axum::{
http::{self, Request},
middleware::{self, Next},
response::Response,
routing::get,
Router,
};
use axum_prometheus::PrometheusMetricLayer;
use base64::{engine::general_purpose, Engine}; use base64::{engine::general_purpose, Engine};
use tokio_util::io::ReaderStream; use tokio_util::io::ReaderStream;
use tower_http::trace::{self, TraceLayer};
use tracing::Level;
use crate::{config, services::{book_library::get_book, filename_getter::get_filename_by_book, downloader::book_download}}; use crate::{
config::CONFIG,
fn check_authorization(headers: HeaderMap) -> Result<(), (StatusCode, String)> { services::{
let config_api_key = config::CONFIG.api_key.clone(); book_library::get_book, downloader::book_download, filename_getter::get_filename_by_book,
},
let api_key = match headers.get(AUTHORIZATION) {
Some(v) => v,
None => return Err((StatusCode::FORBIDDEN, "No api-key!".to_string())),
}; };
if config_api_key != api_key.to_str().unwrap() {
return Err((StatusCode::FORBIDDEN, "Wrong api-key!".to_string()))
}
Ok(())
}
pub async fn download( pub async fn download(
Path((source_id, remote_id, file_type)): Path<(u32, u32, String)>, Path((source_id, remote_id, file_type)): Path<(u32, u32, String)>,
headers: HeaderMap
) -> impl IntoResponse { ) -> impl IntoResponse {
check_authorization(headers)?;
let download_result = match book_download(source_id, remote_id, file_type.as_str()).await { let download_result = match book_download(source_id, remote_id, file_type.as_str()).await {
Ok(v) => v, Ok(v) => v,
Err(_) => { Err(_) => return Err((StatusCode::NO_CONTENT, "Can't download!".to_string())),
return Err((StatusCode::NO_CONTENT, "Can't download!".to_string()))
},
}; };
let data = match download_result { let data = match download_result {
@@ -52,21 +48,20 @@ pub async fn download(
let encoder = general_purpose::STANDARD; let encoder = general_purpose::STANDARD;
let headers = AppendHeaders([ let headers = AppendHeaders([
(header::CONTENT_DISPOSITION, format!("attachment; filename={filename_ascii}")), (
(header::HeaderName::from_static("x-filename-b64"), encoder.encode(filename)) header::CONTENT_DISPOSITION,
format!("attachment; filename={filename_ascii}"),
),
(
header::HeaderName::from_static("x-filename-b64"),
encoder.encode(filename),
),
]); ]);
Ok((headers, body)) Ok((headers, body))
} }
pub async fn get_filename( pub async fn get_filename(Path((book_id, file_type)): Path<(u32, String)>) -> (StatusCode, String) {
Path((book_id, file_type)): Path<(u32, String)>,
headers: HeaderMap
) -> (StatusCode, String){
if let Err(v) = check_authorization(headers) {
return v;
}
let filename = match get_book(book_id).await { let filename = match get_book(book_id).await {
Ok(book) => get_filename_by_book(&book, file_type.as_str(), false, false), Ok(book) => get_filename_by_book(&book, file_type.as_str(), false, false),
Err(_) => return (StatusCode::BAD_REQUEST, "Book not found!".to_string()), Err(_) => return (StatusCode::BAD_REQUEST, "Book not found!".to_string()),
@@ -74,3 +69,44 @@ pub async fn get_filename(
(StatusCode::OK, filename) (StatusCode::OK, filename)
} }
async fn auth<B>(req: Request<B>, next: Next<B>) -> Result<Response, StatusCode> {
let auth_header = req
.headers()
.get(http::header::AUTHORIZATION)
.and_then(|header| header.to_str().ok());
let auth_header = if let Some(auth_header) = auth_header {
auth_header
} else {
return Err(StatusCode::UNAUTHORIZED);
};
if auth_header != CONFIG.api_key {
return Err(StatusCode::UNAUTHORIZED);
}
Ok(next.run(req).await)
}
pub async fn get_router() -> Router {
let (prometheus_layer, metric_handle) = PrometheusMetricLayer::pair();
let app_router = Router::new()
.route("/download/:source_id/:remote_id/:file_type", get(download))
.route("/filename/:book_id/:file_type", get(get_filename))
.layer(middleware::from_fn(auth))
.layer(prometheus_layer);
let metric_router =
Router::new().route("/metrics", get(|| async move { metric_handle.render() }));
Router::new()
.nest("/", app_router)
.nest("/", metric_router)
.layer(
TraceLayer::new_for_http()
.make_span_with(trace::DefaultMakeSpan::new().level(Level::INFO))
.on_response(trace::DefaultOnResponse::new().level(Level::INFO)),
)
}