endstat/src/updater.rs

use crate::{
    config::*,
    handlers::{EndpointStatus, StatusGroup},
    State,
};
use chrono::Utc;
use reqwest::{Client, RedirectPolicy, Url, UrlError};
use ring::{
    digest::{digest, SHA256},
    test::from_hex,
};
use std::time::Duration;

pub fn update_state(state: State) {
    let new_statuses = { Some(update_status(&state.read().unwrap().config)) };
    let mut write_state = state.try_write().expect("Could not unlock");
    write_state.timestamp = Utc::now();
    write_state.timestamp_str = Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
    write_state.groups = new_statuses.unwrap();
}

pub fn update_status(config: &Config) -> Vec<StatusGroup> {
    let mut results: Vec<StatusGroup> = Vec::with_capacity(config.websites.len());

    for website_conf in &config.websites {
        let mut group = Vec::with_capacity(website_conf.endpoints.len());
        for endpoint in &website_conf.endpoints {
            let mut client_builder = Client::builder().timeout(Some(Duration::from_secs(5)));

            if let Some(false) = endpoint.follow_redirects {
                client_builder = client_builder.redirect(RedirectPolicy::none());
            }

            let client = client_builder.build().unwrap();
            group.push(get_result(website_conf, &client, endpoint));
        }

        results.push(StatusGroup {
            label: website_conf.label.clone(),
            endpoints: group,
        });
    }

    results
}

fn get_result(
    website_conf: &WebsiteConfig,
    client: &Client,
    endpoint: &EndpointConfig,
) -> EndpointStatus {
    let (label, path, port, code, body) = get_endpoint_info(endpoint.clone());
    let url = get_url(&website_conf.base, &path, port).expect("reading config");

    let ping_start = Utc::now();
    let ping_result = client.get(&url).send();
    let rtt = Utc::now() - ping_start;
    let rtt_string = if rtt.num_seconds() > 0 {
        format!("{}s", rtt.num_milliseconds() as f64 / 1000.)
    } else {
        format!("{}ms", rtt.num_milliseconds())
    };

    match ping_result {
        Ok(mut res) => {
            let res_body = res.text().expect("could not get body of request");
            let mut error = None;

            if res.status() != code {
                error = append_err_msg(
                    error,
                    format!(
                        "Status code mismatch: {} != {}.",
                        res.status().as_u16(),
                        code
                    ),
                );
            }

            if let Some(expected_hash) = &endpoint.body_hash {
                let expected = from_hex(expected_hash).unwrap();
                let actual = digest(&SHA256, String::from(res_body).as_bytes());
                if &expected != &actual.as_ref() {
                    error = append_err_msg(error, String::from("Body hash mismatch."));
                }
            } else if !body.is_empty() && res_body != body {
                error = append_err_msg(
                    error,
                    format!("Body mismatch: {} != {}.", res_body.len(), body.len()),
                );
            }

            if rtt.num_seconds() >= 2 {
                error = append_err_msg(error, format!("RTT too long: {}", rtt_string));
            }

            if error.is_some() {
                EndpointStatus::warn(url, label, rtt_string, error)
            } else {
                EndpointStatus::ok(url, label, rtt_string)
            }
        }
        Err(e) => {
            if let Some(true) = endpoint.should_err {
                EndpointStatus::ok(url, label, rtt_string)
            } else {
                EndpointStatus::error(url, label, Some(format!("{}", e)))
            }
        }
    }
}

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

fn append_err_msg(optional: Option<String>, to_append: String) -> Option<String> {
    // Turns out Rust doesn't recognize that only one of the FnOnce occurs, so
    // as a result both closures want ownership of to_append
    // Some(optional.map_or_else(|| to_append, |msg| format!("{} {}", msg, to_append)))

    Some(match optional {
        Some(e) => format!("{} {}", e, to_append),
        None => to_append,
    })
}