From a088f6224ecb935150c90c5d2c82c32e156138ee Mon Sep 17 00:00:00 2001 From: LittleSheep Date: Sun, 14 Jan 2024 18:39:22 +0800 Subject: [PATCH] :sparkles: Basic Auth --- Cargo.toml | 16 +++++++-- Settings.toml | 1 + src/auth.rs | 50 ++++++++++++++++++++++++++ src/main.rs | 15 +++++--- src/proxies/config.rs | 7 ++-- src/proxies/metrics.rs | 7 +++- src/sideload/mod.rs | 16 +++++++++ src/sideload/overview.rs | 76 ++++++++++++++++++---------------------- src/sideload/regions.rs | 25 +++++++++++++ 9 files changed, 160 insertions(+), 53 deletions(-) create mode 100644 src/auth.rs create mode 100644 src/sideload/regions.rs diff --git a/Cargo.toml b/Cargo.toml index 95fae7b..3bd45fa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,15 +13,25 @@ hyper-util = { version = "0.1.2", features = ["full"] } lazy_static = "1.4.0" mime = "0.3.17" percent-encoding = "2.3.1" -poem = { version = "2.0.0", features = ["tokio-metrics", "websocket", "static-files", "reqwest"] } -poem-openapi = { version = "4.0.0", features = ["swagger-ui"] } +poem = { version = "2.0.0", features = [ + "tokio-metrics", + "websocket", + "static-files", + "reqwest", +] } +poem-openapi = { version = "4.0.0" } queryst = "3.0.0" rand = "0.8.5" regex = "1.10.2" reqwest = { git = "https://github.com/seanmonstar/reqwest.git", branch = "hyper-v1", version = "0.11.23" } serde = "1.0.195" serde_json = "1.0.111" -tokio = { version = "1.35.1", features = ["rt-multi-thread", "macros", "time", "full"] } +tokio = { version = "1.35.1", features = [ + "rt-multi-thread", + "macros", + "time", + "full", +] } tokio-tungstenite = "0.21.0" toml = "0.8.8" tracing = "0.1.40" diff --git a/Settings.toml b/Settings.toml index d1fc034..b31488d 100644 --- a/Settings.toml +++ b/Settings.toml @@ -1,4 +1,5 @@ regions = "./regions" +secret = "aEXcED5xJ3" [listen] proxies = "0.0.0.0:80" diff --git a/src/auth.rs b/src/auth.rs new file mode 100644 index 0000000..79b28c4 --- /dev/null +++ b/src/auth.rs @@ -0,0 +1,50 @@ +use http::StatusCode; +use poem::{ + web::headers::{self, authorization::Basic, HeaderMapExt}, + Endpoint, Error, Middleware, Request, Response, Result, +}; + +pub struct BasicAuth { + pub username: String, + pub password: String, +} + +impl Middleware for BasicAuth { + type Output = BasicAuthEndpoint; + + fn transform(&self, ep: E) -> Self::Output { + BasicAuthEndpoint { + ep, + username: self.username.clone(), + password: self.password.clone(), + } + } +} + +pub struct BasicAuthEndpoint { + ep: E, + username: String, + password: String, +} + +#[poem::async_trait] +impl Endpoint for BasicAuthEndpoint { + type Output = E::Output; + + async fn call(&self, req: Request) -> Result { + if let Some(auth) = req.headers().typed_get::>() { + if auth.0.username() == self.username && auth.0.password() == self.password { + return self.ep.call(req).await; + } + } + Err(Error::from_response( + Response::builder() + .header( + "WWW-Authenticate", + "Basic realm=\"RoadSig\", charset=\"UTF-8\"", + ) + .status(StatusCode::UNAUTHORIZED) + .finish(), + )) + } +} diff --git a/src/main.rs b/src/main.rs index 27195b2..b5859be 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,4 @@ +pub mod auth; mod config; mod proxies; mod sideload; @@ -5,7 +6,7 @@ mod sideload; use std::collections::VecDeque; use lazy_static::lazy_static; -use poem::{listener::TcpListener, Route, Server}; +use poem::{listener::TcpListener, EndpointExt, Route, Server}; use poem_openapi::OpenApiService; use proxies::RoadInstance; use tokio::sync::Mutex; @@ -64,7 +65,6 @@ async fn main() -> Result<(), std::io::Error> { // Sideload let sideload = OpenApiService::new(sideload::SideloadApi, "Sideload API", "1.0") .server("http://localhost:3000/cgi"); - let sideload_ui = sideload.swagger_ui(); let sideload_server = Server::new(TcpListener::bind( config::C @@ -74,9 +74,14 @@ async fn main() -> Result<(), std::io::Error> { .unwrap_or("0.0.0.0:81".to_string()), )) .run( - Route::new() - .nest("/cgi", sideload) - .nest("/swagger", sideload_ui), + Route::new().nest("/cgi", sideload).with(auth::BasicAuth { + username: "RoadSign".to_string(), + password: config::C + .read() + .await + .get_string("secret") + .unwrap_or("password".to_string()), + }), ); tokio::try_join!(proxies_server, sideload_server)?; diff --git a/src/proxies/config.rs b/src/proxies/config.rs index c6c1835..ae4f2d5 100644 --- a/src/proxies/config.rs +++ b/src/proxies/config.rs @@ -1,18 +1,19 @@ use std::collections::HashMap; +use poem_openapi::Object; use queryst::parse; use serde::{Deserialize, Serialize}; use serde_json::json; use super::responder::StaticResponderConfig; -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Object, Clone, Serialize, Deserialize)] pub struct Region { pub id: String, pub locations: Vec, } -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Object, Clone, Serialize, Deserialize)] pub struct Location { pub id: String, pub hosts: Vec, @@ -23,7 +24,7 @@ pub struct Location { pub destinations: Vec, } -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Object, Clone, Serialize, Deserialize)] pub struct Destination { pub id: String, pub uri: String, diff --git a/src/proxies/metrics.rs b/src/proxies/metrics.rs index dad7275..4ecd101 100644 --- a/src/proxies/metrics.rs +++ b/src/proxies/metrics.rs @@ -47,6 +47,8 @@ pub struct RoadMetrics { pub recent_errors: VecDeque, } +const MAX_TRACE_COUNT: usize = 10; + impl RoadMetrics { pub fn get_success_rate(&self) -> f64 { if self.requests_count > 0 { @@ -60,6 +62,9 @@ impl RoadMetrics { self.requests_count += 1; self.recent_successes .push_back(RoadTrace::from_structs(reg, loc, end)); + if self.recent_successes.len() > MAX_TRACE_COUNT { + self.recent_successes.pop_front(); + } } pub fn add_faliure_request( @@ -73,7 +78,7 @@ impl RoadMetrics { self.failures_count += 1; self.recent_errors .push_back(RoadTrace::from_structs_with_error(reg, loc, end, err)); - if self.recent_errors.len() > 10 { + if self.recent_errors.len() > MAX_TRACE_COUNT { self.recent_errors.pop_front(); } } diff --git a/src/sideload/mod.rs b/src/sideload/mod.rs index 95a4645..e69231e 100644 --- a/src/sideload/mod.rs +++ b/src/sideload/mod.rs @@ -1,3 +1,19 @@ +use poem_openapi::OpenApi; + pub mod overview; +pub mod regions; pub struct SideloadApi; + +#[OpenApi] +impl SideloadApi { + #[oai(path = "/", method = "get")] + async fn index(&self) -> overview::OverviewResponse { + overview::index().await + } + + #[oai(path = "/regions", method = "get")] + async fn regions_index(&self) -> regions::RegionResponse { + regions::index().await + } +} diff --git a/src/sideload/overview.rs b/src/sideload/overview.rs index ba2bc4c..8c3ac6b 100644 --- a/src/sideload/overview.rs +++ b/src/sideload/overview.rs @@ -1,4 +1,4 @@ -use poem_openapi::{payload::Json, ApiResponse, Object, OpenApi}; +use poem_openapi::{payload::Json, ApiResponse, Object}; use crate::{ proxies::{ @@ -8,17 +8,15 @@ use crate::{ ROAD, }; -use super::SideloadApi; - #[derive(ApiResponse)] -enum OverviewResponse { +pub enum OverviewResponse { /// Return the overview data. #[oai(status = 200)] Ok(Json), } #[derive(Debug, Object, Clone, PartialEq)] -struct OverviewData { +pub struct OverviewData { /// Loaded regions count #[oai(read_only)] regions: usize, @@ -42,40 +40,36 @@ struct OverviewData { recent_errors: Vec, } -#[OpenApi] -impl SideloadApi { - #[oai(path = "/", method = "get")] - async fn index(&self) -> OverviewResponse { - let locked_app = ROAD.lock().await; - let regions = locked_app.regions.clone(); - let locations = regions - .iter() - .flat_map(|item| item.locations.clone()) - .collect::>(); - let destinations = locations - .iter() - .flat_map(|item| item.destinations.clone()) - .collect::>(); - OverviewResponse::Ok(Json(OverviewData { - regions: regions.len(), - locations: locations.len(), - destinations: destinations.len(), - requests_count: locked_app.metrics.requests_count, - successes_count: locked_app.metrics.requests_count - locked_app.metrics.failures_count, - faliures_count: locked_app.metrics.failures_count, - success_rate: locked_app.metrics.get_success_rate(), - recent_successes: locked_app - .metrics - .recent_successes - .clone() - .into_iter() - .collect::>(), - recent_errors: locked_app - .metrics - .recent_errors - .clone() - .into_iter() - .collect::>(), - })) - } +pub async fn index() -> OverviewResponse { + let locked_app = ROAD.lock().await; + let regions = locked_app.regions.clone(); + let locations = regions + .iter() + .flat_map(|item| item.locations.clone()) + .collect::>(); + let destinations = locations + .iter() + .flat_map(|item| item.destinations.clone()) + .collect::>(); + OverviewResponse::Ok(Json(OverviewData { + regions: regions.len(), + locations: locations.len(), + destinations: destinations.len(), + requests_count: locked_app.metrics.requests_count, + successes_count: locked_app.metrics.requests_count - locked_app.metrics.failures_count, + faliures_count: locked_app.metrics.failures_count, + success_rate: locked_app.metrics.get_success_rate(), + recent_successes: locked_app + .metrics + .recent_successes + .clone() + .into_iter() + .collect::>(), + recent_errors: locked_app + .metrics + .recent_errors + .clone() + .into_iter() + .collect::>(), + })) } diff --git a/src/sideload/regions.rs b/src/sideload/regions.rs new file mode 100644 index 0000000..827d603 --- /dev/null +++ b/src/sideload/regions.rs @@ -0,0 +1,25 @@ +use poem_openapi::{payload::Json, ApiResponse}; + +use crate::{proxies::config::Region, ROAD}; + +#[derive(ApiResponse)] +pub enum RegionResponse { + /// Return the region data. + #[oai(status = 200)] + Ok(Json), + /// Return the list of region data. + #[oai(status = 200)] + OkMany(Json>), + /// Return the region data after created. + #[oai(status = 201)] + Created(Json), + /// Return was not found. + #[oai(status = 404)] + NotFound, +} + +pub async fn index() -> RegionResponse { + let locked_app = ROAD.lock().await; + + RegionResponse::OkMany(Json(locked_app.regions.clone())) +}