endstat/src/main.rs

extern crate actix_web;
extern crate env_logger;
extern crate reqwest;
extern crate ron;
extern crate serde;
#[macro_use]
extern crate tera;

mod utils;

use self::utils::EpochTimestamp;
use actix_web::{
    error::ErrorInternalServerError,
    middleware::Logger,
    web::{resource, Data},
    App, Error as WebError, HttpResponse, HttpServer, Result as WebResult,
};
use reqwest::{Client, Url, UrlError};
use ron::de::from_str;
use serde::{Deserialize, Serialize};
use std::{
    error::Error,
    fs::read_to_string,
    sync::{Arc, Mutex, MutexGuard},
};
use tera::{Context, Tera};

#[derive(Deserialize, Serialize, Debug, Clone)]
struct EndpointConf {
    label: Option<String>,
    endpoint: Option<String>,
    port: Option<u16>,
    code: Option<u16>,
    body: Option<String>,
}

#[derive(Deserialize, Serialize, Debug, Clone)]
struct WebsiteConf {
    label: String,
    base: String,
    endpoints: Vec<EndpointConf>,
}

#[derive(Deserialize, Serialize, Debug, Clone)]
struct Config {
    refresh_time: u64,
    bind_address: String,
    websites: Vec<WebsiteConf>,
}

#[derive(Clone, Serialize)]
pub struct Status {
    status: u8,
    location: String,
    domain: String,
    endpoint: String,
    error: Option<String>,
}

#[derive(Serialize)]
pub struct FetchResults {
    last_update: EpochTimestamp,
    refresh_time: u64,
    config: Config,
    statuses: Vec<Status>,
}

type StatusState = Arc<Mutex<FetchResults>>;

fn index(tmpl: Data<Tera>, state: Data<StatusState>) -> WebResult<HttpResponse, WebError> {
    let state = update_state(state.lock().unwrap());
    let mut ctx = Context::new();
    ctx.insert("results", &*state);
    let s = tmpl
        .render("index.html", &tera::Context::new())
        .map_err(|_| ErrorInternalServerError("Template error"))?;
    Ok(HttpResponse::Ok().content_type("text/html").body(s))
}

fn json_endpoint(state: Data<StatusState>) -> HttpResponse {
    let state = update_state(state.lock().unwrap());
    HttpResponse::Ok().json(&state.statuses)
}

fn update_state(mut state: MutexGuard<FetchResults>) -> MutexGuard<FetchResults> {
    if EpochTimestamp::now() - state.last_update >= state.refresh_time {
        state.last_update = EpochTimestamp::now();
        state.statuses = update_status(&state.config);
    }

    state
}

fn main() -> Result<(), Box<Error>> {
    let config = from_str::<Config>(&read_to_string("./endstat_conf.ron")?)?;
    let bind_addr = config.bind_address.clone();
    std::env::set_var("RUST_LOG", "actix_web=info");
    env_logger::init();

    HttpServer::new(move || {
        let state = Arc::from(Mutex::from(FetchResults {
            last_update: EpochTimestamp::now(),
            refresh_time: config.refresh_time.clone(),
            config: config.clone(),
            statuses: update_status(&config),
        }));
        let tera = compile_templates!(concat!(env!("CARGO_MANIFEST_DIR"), "/templates/**/*"));

        App::new()
            .data(state)
            .data(tera)
            .wrap(Logger::default())
            .service(resource("/").to(index))
            .service(resource("/api").to(json_endpoint))
    })
    .bind(&bind_addr)?
    .run()?;

    Ok(())
}

fn update_status(config: &Config) -> Vec<Status> {
    let client = Client::new();
    let mut results: Vec<Status> = vec![];

    for website_conf in &config.websites {
        for endpoint in &website_conf.endpoints {
            let (label, path, port, code, body) = get_endpoint_info(endpoint.clone());
            let url = get_url(&website_conf.base, &path, port).expect("reading config");
            if let Ok(mut res) = client.get(&url).send() {
                let res_body = res.text().expect("could not get body of request");
                let does_code_match = res.status() == code;
                let does_body_match = body.is_empty() || res_body == body;

                results.push(if !does_code_match {
                    Status {
                        status: 1,
                        location: url,
                        domain: website_conf.label.clone(),
                        endpoint: label,
                        error: Some(format!(
                            "Status code mismatch! {} != {}",
                            res.status().as_u16(),
                            code
                        )),
                    }
                } else if !does_body_match {
                    Status {
                        status: 2,
                        location: url,
                        domain: website_conf.label.clone(),
                        endpoint: label,
                        error: Some(format!("Body mismatch! {} != {}", res_body, body)),
                    }
                } else {
                    Status {
                        status: 0,
                        location: url,
                        domain: website_conf.label.clone(),
                        endpoint: label,
                        error: None,
                    }
                });
            }
        }
    }

    results
}

fn get_url(base: &String, path: &String, port: Option<u16>) -> Result<String, UrlError> {
    let mut url = Url::parse(base)?.join(path)?;
    if let Err(e) = url.set_port(port) {
        println!("{:?}", e);
    }
    Ok(url.into_string())
}

fn get_endpoint_info(endpoint: EndpointConf) -> (String, String, Option<u16>, u16, String) {
    let path = endpoint.endpoint.unwrap_or_default();
    let label = endpoint.label.unwrap_or_else(|| path.clone());
    let code = endpoint.code.unwrap_or_else(|| 200);
    let body = endpoint.body.unwrap_or_default();

    (label, path, endpoint.port, code, body)
}