From fd992d7e3c03f37fbcafe9d3f26c72a2ead3b2a7 Mon Sep 17 00:00:00 2001 From: Toby Vincent Date: Thu, 26 Sep 2024 17:31:16 -0500 Subject: feat!: impl full api --- src/api.rs | 73 ++++++++++++++++++ src/api/services.rs | 40 ++++++++++ src/error.rs | 22 ++++++ src/lib.rs | 82 ++++++++------------ src/main.rs | 93 ++++++++++++++-------- src/service.rs | 203 ++++++++++++++++++------------------------------- src/service/http.rs | 49 ++++++++++++ src/service/systemd.rs | 33 ++++++++ src/service/tcp.rs | 28 +++++++ 9 files changed, 409 insertions(+), 214 deletions(-) create mode 100644 src/api.rs create mode 100644 src/api/services.rs create mode 100644 src/service/http.rs create mode 100644 src/service/systemd.rs create mode 100644 src/service/tcp.rs (limited to 'src') diff --git a/src/api.rs b/src/api.rs new file mode 100644 index 0000000..e6a91ba --- /dev/null +++ b/src/api.rs @@ -0,0 +1,73 @@ +use std::collections::HashMap; + +use axum::{extract::State, response::IntoResponse, Json}; +use serde::{Deserialize, Serialize}; + +use crate::{service::Services, Check, Error, Status}; + +pub mod services; + +pub fn router() -> axum::Router { + use axum::routing::get; + + axum::Router::new() + .route("/healthcheck", get(healthcheck)) + .merge(services::router()) + .fallback(fallback) +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct Health { + pub status: Status, + pub output: Option, + pub checks: HashMap, +} + +impl From for Health { + fn from(value: T) -> Self { + Health { + status: Status::Fail, + output: Some(value.to_string()), + ..Default::default() + } + } +} + +impl IntoResponse for Health { + fn into_response(self) -> axum::response::Response { + Json(self).into_response() + } +} + +pub async fn healthcheck(State(services): State) -> Health { + let checks = match services.check().await { + Ok(c) => c, + Err(err) => { + return Health { + status: Status::Fail, + output: Some(err.to_string()), + ..Default::default() + } + } + }; + + let (status, output) = match checks + .values() + .filter(|s| !matches!(s.status, Status::Pass)) + .count() + { + 0 => (Status::Pass, None), + 1 => (Status::Fail, Some("1 issue detected".to_string())), + n => (Status::Fail, Some(format!("{n} issues detected"))), + }; + + Health { + status, + output, + checks, + } +} + +pub async fn fallback(uri: axum::http::Uri) -> Error { + Error::RouteNotFound(uri) +} diff --git a/src/api/services.rs b/src/api/services.rs new file mode 100644 index 0000000..59e891f --- /dev/null +++ b/src/api/services.rs @@ -0,0 +1,40 @@ +use std::collections::HashMap; + +use axum::{ + extract::{Path, Query, State}, + Json, Router, +}; +use axum_extra::routing::Resource; +use serde::{Deserialize, Serialize}; + +use crate::{service::Services, Check, Error, Status}; + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct ServiceQuery { + pub name: Option, + pub state: Option, +} + +pub fn router() -> Router { + Resource::named("services").index(index).show(show).into() +} + +pub async fn index( + Query(query): Query, + State(services): State, +) -> Result>, Error> { + services + .check_filtered(|name| (!query.name.as_ref().is_some_and(|s| s != name))) + .await + .map(Json) +} + +pub async fn show( + Path(name): Path, + State(services): State, +) -> Result { + services + .check_one(&name) + .await + .ok_or_else(|| Error::ServiceNotFound(name))? +} diff --git a/src/error.rs b/src/error.rs index ef30b97..109c944 100644 --- a/src/error.rs +++ b/src/error.rs @@ -13,4 +13,26 @@ pub enum Error { #[error("Invalid HTTP method")] Method, + + #[error("Axum error: {0}")] + Axum(#[from] axum::Error), + + #[error("Route not found: {0}")] + RouteNotFound(axum::http::Uri), + + #[error("Service not found: {0}")] + ServiceNotFound(String), +} + +impl axum::response::IntoResponse for Error { + fn into_response(self) -> axum::response::Response { + use axum::http::StatusCode; + + let status = match self { + Self::RouteNotFound(_) | Self::ServiceNotFound(_) => StatusCode::NOT_FOUND, + _ => StatusCode::INTERNAL_SERVER_ERROR, + }; + + (status, self.to_string()).into_response() + } } diff --git a/src/lib.rs b/src/lib.rs index 3437dca..2c9fa91 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,61 +1,39 @@ -use std::ops::Range; +use std::path::PathBuf; -pub use crate::{ - error::{Error, Result}, - service::{Service, Status}, -}; +use serde::{Deserialize, Serialize}; +use service::Services; +use tower_http::services::ServeDir; +pub use crate::error::{Error, Result}; + +pub mod api; pub mod error; pub mod service; -pub fn generate(title: String, mut services: Vec, template: String) -> String { - let client = reqwest::blocking::Client::new(); - - let [up, down, unknown] = std::thread::scope(|s| { - let mut handles = Vec::new(); - for service in services.iter_mut() { - handles.push(s.spawn(|| service.check(client.clone()))); - } +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum Status { + #[default] + Pass, + Fail, + Warn, +} - handles - .into_iter() - .map(|h| h.join().expect("Joining thread")) - .fold([0, 0, 0], |[up, down, unknown], res| match res { - Ok(true) => [up + 1, down, unknown], - Ok(false) => [up, down + 1, unknown], - Err(_) => [up, down, unknown + 1], - }) - }); +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct Check { + pub status: Status, + pub output: Option, +} - template - .match_indices("{{services}}") - .zip(template.match_indices("{{end}}")) - .map(|(start, stop)| { - ( - start.0 + start.1.len()..stop.0, - start.0..stop.0 + stop.1.len(), - ) - }) - .collect::>() - .into_iter() - .fold(template, |mut template, (Range { start, end }, outer)| { - let replace_with = services - .iter() - .map(|service| { - template[start..end] - .replace("{name}", &service.name) - .replace("{title}", &service.kind.to_string()) - .replace("{state}", &service.state.to_string()) - .replace("{level}", &service.state.as_level()) - }) - .collect::(); +impl axum::response::IntoResponse for Check { + fn into_response(self) -> axum::response::Response { + axum::Json(self).into_response() + } +} - template.replace_range(outer, &replace_with); - template - }) - .replace("{title}", &title) - .replace("{status}", if down > 0 { "error" } else { "ok" }) - .replace("{up}", &up.to_string()) - .replace("{down}", &down.to_string()) - .replace("{unknown}", &unknown.to_string()) +pub fn router(root: PathBuf) -> axum::Router { + axum::Router::new() + .nest_service("/", ServeDir::new(root)) + .nest("/api", api::router()) + .layer(tower_http::trace::TraceLayer::new_for_http()) } diff --git a/src/main.rs b/src/main.rs index 2ff9fd3..97ed111 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,42 +1,69 @@ -use std::{fs::File, io::Write, path::PathBuf}; +use std::{fs::File, path::PathBuf}; +use tracing::level_filters::LevelFilter; +use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter}; -use statsrv::Service; +use statsrv::service::Services; + +#[cfg(not(debug_assertions))] +const DEFAULT_CONFIG: &str = "/etc/statsrv.toml"; +#[cfg(debug_assertions)] +const DEFAULT_CONFIG: &str = "./config.toml"; + +#[tokio::main] +async fn main() -> Result<(), Box> { + tracing_subscriber::registry() + .with( + EnvFilter::builder() + .with_default_directive(LevelFilter::INFO.into()) + .from_env_lossy(), + ) + .with(tracing_subscriber::fmt::layer()) + .init(); + + let config = match Config::parse() { + Ok(c) => c, + Err(err) => { + tracing::debug!("Failed to read config file: `{err}`"); + tracing::debug!("Using default config values"); + Default::default() + } + }; + + let router = statsrv::router(config.root).with_state(config.services); + + let listener = tokio::net::TcpListener::bind(config.address).await.unwrap(); + tracing::info!("listening on {}", listener.local_addr().unwrap()); + + axum::serve(listener, router).await.map_err(Into::into) +} #[derive(Debug, Clone, serde::Deserialize)] +#[serde(default)] pub struct Config { - pub title: String, - pub template_path: PathBuf, - pub output_dir: Option, - pub address: Option, - pub services: Vec, + pub root: PathBuf, + pub address: String, + pub services: Services, } -fn main() -> Result<(), main_error::MainError> { - let mut args = std::env::args().skip(1); - - let config_path = args - .next() - .unwrap_or_else(|| "/etc/statsrv.toml".to_string()); - let config_file = File::open(config_path)?; - let config_toml = std::io::read_to_string(config_file)?; - let Config { - title, - template_path: template, - output_dir, - address: _, - services, - } = toml::from_str(&config_toml)?; - - let template_file = File::open(template)?; - let template = std::io::read_to_string(template_file)?; - let status_page = statsrv::generate(title, services, template); - - if let Some(output_dir) = output_dir { - std::fs::create_dir_all(&output_dir)?; - let mut html_writer = File::create(output_dir.join("index.html"))?; - - html_writer.write_all(status_page.as_bytes())?; +impl Config { + fn parse() -> Result> { + let config_path = std::env::args().nth(1).unwrap_or_else(|| { + tracing::debug!("Falling back to default config location"); + DEFAULT_CONFIG.to_string() + }); + + let config_file = File::open(&config_path)?; + let config_toml = std::io::read_to_string(config_file)?; + toml::from_str(&config_toml).map_err(Into::into) } +} - Ok(()) +impl Default for Config { + fn default() -> Self { + Self { + root: PathBuf::from("./"), + address: String::from("127.0.0.1:8080"), + services: Services::new(Default::default()), + } + } } diff --git a/src/service.rs b/src/service.rs index c5eb0d7..677db17 100644 --- a/src/service.rs +++ b/src/service.rs @@ -1,160 +1,105 @@ -use std::{fmt::Display, process::Command}; +use std::{collections::HashMap, fmt::Display}; -use reqwest::blocking::Client; +use futures::{stream::FuturesOrdered, TryStreamExt}; +use http::Http; use serde::Deserialize; +use systemd::Systemd; +use tcp::Tcp; -use crate::Error; +use crate::{Check, Error}; + +pub mod http; +pub mod systemd; +pub mod tcp; #[derive(Debug, Clone, Deserialize)] -pub struct Service { - pub name: String, +pub struct Services { #[serde(flatten)] - pub kind: Kind, - #[serde(skip)] - pub state: State, + inner: HashMap, + #[serde(skip, default = "Services::default_client")] + client: reqwest::Client, } -impl Service { - pub fn check(&mut self, client: Client) -> Result { - self.state = self.kind.get_state(client)?; - Ok(self.state.is_operational()) +impl Services { + pub fn new(services: HashMap) -> Self { + let client = reqwest::Client::new(); + Self { + inner: services, + client, + } } -} - -#[derive(Debug, Clone, Deserialize)] -#[serde(tag = "type", rename_all = "lowercase")] -pub enum Kind { - Tcp { - address: String, - }, - Http { - url: String, - #[serde(default = "Kind::default_method")] - method: String, - #[serde(default = "Kind::default_code")] - status_code: u16, - }, - Systemd { - service: String, - }, -} -impl Kind { - fn default_method() -> String { - "GET".to_string() + fn default_client() -> reqwest::Client { + reqwest::Client::new() } - fn default_code() -> u16 { - 200 + pub async fn check(&self) -> Result, Error> { + let checks = self + .inner + .values() + .map(|service| service.check(self.client.clone())) + .collect::>() + .try_collect::>() + .await?; + + Ok(self + .inner + .keys() + .cloned() + .zip(checks) + .collect::>()) } - pub fn get_state(&self, client: Client) -> Result { - let state = match self { - Kind::Tcp { address } => { - if std::net::TcpStream::connect(address).is_ok() { - State::Operational - } else { - State::Down("Unreachable".to_string()) - } - } - Kind::Http { - method, - url, - status_code, - } => { - match client - .request(method.parse().map_err(|_| Error::Method)?, url) - .send()? - .status() - { - s if s.as_u16() == *status_code => State::Operational, - s => State::Down(s.to_string()), - } - } - Kind::Systemd { service } => { - let output = Command::new("systemctl") - .arg("is-active") - .arg(service) - .output()?; - - if output.status.success() { - State::Operational - } else { - State::Down(String::from_utf8_lossy(&output.stdout).to_string()) - } - } - }; - - Ok(state) + pub async fn check_one(&self, name: &str) -> Option> { + Some(self.inner.get(name)?.check(self.client.clone()).await) } -} -impl Display for Kind { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Kind::Tcp { address } => write!(f, "tcp://{address}"), - Kind::Http { method, url, .. } => write!(f, "{method} {url}"), - Kind::Systemd { service } => write!(f, "{service}"), - } + pub async fn check_filtered

(&self, mut predicate: P) -> Result, Error> + where + P: FnMut(&String) -> bool, + { + let checks = self + .inner + .iter() + .filter_map(|(s, service)| predicate(s).then_some(service)) + .map(|service| service.check(self.client.clone())) + .collect::>() + .try_collect::>() + .await?; + + Ok(self + .inner + .keys() + .cloned() + .zip(checks) + .collect::>()) } } -#[derive(Debug, Clone, Default)] -pub struct Status { - pub info: String, - pub state: State, -} - -#[derive(Debug, Clone, Default)] -pub enum State { - #[default] - Unknown, - Operational, - Down(String), +#[derive(Debug, Clone, Deserialize)] +#[serde(untagged)] +pub enum Service { + Http(Http), + Tcp(Tcp), + Systemd(Systemd), } -impl State { - /// Returns `true` if this is a `Unknown` variant. - pub fn is_unknown(&self) -> bool { - matches!(self, Self::Unknown) - } - - /// Returns `true` if this is a `Operational` variant. - pub fn is_operational(&self) -> bool { - matches!(self, Self::Operational) - } - - /// Returns `true` if this is a `Down` variant. - pub fn is_down(&self) -> bool { - matches!(self, Self::Down(_)) - } - - /// Converts the `State` into an `Option` containing `String` description if the `State` was - /// `Down` and `None` otherwise. - pub fn down_value(self) -> Option { - match self { - State::Unknown => None, - State::Operational => None, - State::Down(s) => Some(s), - } - } - - pub fn as_level(&self) -> String { +impl Service { + pub async fn check(&self, client: reqwest::Client) -> Result { match self { - State::Unknown => "warning", - State::Operational => "ok", - State::Down(_) => "error", + Service::Http(http) => http.check(client).await, + Service::Tcp(tcp) => tcp.check().await, + Service::Systemd(systemd) => systemd.check().await, } - .to_string() } } -impl Display for State { +impl Display for Service { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { - State::Unknown => write!(f, "Unknown"), - State::Operational => write!(f, "Operational"), - State::Down(s) => write!(f, "{s}"), + Service::Http(http) => http.fmt(f), + Service::Tcp(tcp) => tcp.fmt(f), + Service::Systemd(systemd) => systemd.fmt(f), } } } diff --git a/src/service/http.rs b/src/service/http.rs new file mode 100644 index 0000000..15696a1 --- /dev/null +++ b/src/service/http.rs @@ -0,0 +1,49 @@ +use std::fmt::Display; + +use serde::Deserialize; + +use crate::{Check, Error, Status}; + +#[derive(Debug, Clone, Deserialize)] +pub struct Http { + pub url: String, + #[serde(default = "Http::default_method")] + pub method: String, + #[serde(default = "Http::default_code")] + pub status_code: u16, +} + +impl Http { + fn default_method() -> String { + "GET".to_string() + } + + fn default_code() -> u16 { + 200 + } +} + +impl Display for Http { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{} {}", self.method, self.url) + } +} + +impl Http { + pub async fn check(&self, client: reqwest::Client) -> Result { + let status_code = client + .request(self.method.parse().map_err(|_| Error::Method)?, &self.url) + .send() + .await? + .status() + .as_u16(); + + match status_code == self.status_code { + true => Ok(Check::default()), + false => Ok(Check { + status: Status::Fail, + output: Some(format!("Status code: {status_code}")), + }), + } + } +} diff --git a/src/service/systemd.rs b/src/service/systemd.rs new file mode 100644 index 0000000..2e3b74c --- /dev/null +++ b/src/service/systemd.rs @@ -0,0 +1,33 @@ +use std::{fmt::Display, process::Command}; + +use serde::Deserialize; + +use crate::{Check, Error, Status}; + +#[derive(Debug, Clone, Deserialize)] +pub struct Systemd { + pub service: String, +} + +impl Display for Systemd { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}.service", self.service.trim_end_matches(".service")) + } +} + +impl Systemd { + pub async fn check(&self) -> Result { + let output = Command::new("systemctl") + .arg("is-active") + .arg(&self.service) + .output()?; + + let stdout = String::from_utf8_lossy(&output.stdout).to_string(); + Ok((!output.status.success()) + .then(|| Check { + status: Status::Fail, + output: Some(format!("Service state: {}", stdout.trim())), + }) + .unwrap_or_default()) + } +} diff --git a/src/service/tcp.rs b/src/service/tcp.rs new file mode 100644 index 0000000..5f55091 --- /dev/null +++ b/src/service/tcp.rs @@ -0,0 +1,28 @@ +use std::fmt::Display; + +use serde::Deserialize; + +use crate::{Check, Error, Status}; + +#[derive(Debug, Clone, Deserialize)] +pub struct Tcp { + pub address: String, +} + +impl Display for Tcp { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "tcp://{}", self.address) + } +} + +impl Tcp { + pub async fn check(&self) -> Result { + Ok(std::net::TcpStream::connect(&self.address) + .err() + .map(|err| Check { + status: Status::Fail, + output: Some(format!("error: {err}")), + }) + .unwrap_or_default()) + } +} -- cgit v1.2.3-70-g09d2