From ef049e517171af67afeaa90cd29e578f0ec5b751 Mon Sep 17 00:00:00 2001
From: Joakim Hulthe <joakim@hulthe.net>
Date: Mon, 1 Feb 2021 16:17:03 +0100
Subject: [PATCH] Add more detailed Http mode config

---
 example.config.ron |  5 ++-
 src/health/mod.rs  | 84 ++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 78 insertions(+), 11 deletions(-)

diff --git a/example.config.ron b/example.config.ron
index cc2d454..cbcdfc0 100644
--- a/example.config.ron
+++ b/example.config.ron
@@ -3,7 +3,10 @@
     services: {
         "example": (
             url: "https://example.org",
-            mode: Check200,
+            mode: (
+                up_status_codes: [Status2XX],
+                down_status_codes: [Status(404)],
+            ),
         ),
     },
 )
diff --git a/src/health/mod.rs b/src/health/mod.rs
index 08526cd..492c895 100644
--- a/src/health/mod.rs
+++ b/src/health/mod.rs
@@ -1,20 +1,21 @@
 use chrono::{DateTime, Utc};
 use futures::lock::Mutex;
-use log::info;
+use log::{error, info};
+use reqwest::{Response, StatusCode};
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 
 pub type ServiceId = String;
 
-#[derive(Debug)]
+#[derive(Debug, Copy, Clone)]
 pub enum HealthStatus {
     /// The service is up and running
     Up,
 
-    /// The service gave an error when performing the health check
+    /// The service gave a valid error when performing the health check
     Errored,
 
-    /// The service did not respond to the health check
+    /// The service seems to be completely offline
     Down,
 }
 
@@ -25,9 +26,26 @@ pub struct ServiceConfig {
 }
 
 #[derive(Serialize, Deserialize)]
-pub enum HttpHealthCheckMode {
-    /// Expect the server to return a 200 status code. The body is ignored.
-    Check200,
+pub enum PartialStatusCode {
+    Any,
+    Status1XX,
+    Status2XX,
+    Status3XX,
+    Status4XX,
+    Status5XX,
+    Status(u16),
+}
+
+/// If the response status code matches matches any in `ok_status_codes`, status will be Up
+/// else if it matches any in `down_status_codes` status will be Down, else status will be Errored
+#[derive(Serialize, Deserialize)]
+pub struct HttpHealthCheckMode {
+    /// If the response status code matches one in this list, the status will be Up
+    pub up_status_codes: Vec<PartialStatusCode>,
+
+    /// If the response status code matches one in this list, the status will be Down
+    #[serde(default)]
+    pub down_status_codes: Vec<PartialStatusCode>,
 }
 
 pub struct HealthState {
@@ -62,15 +80,61 @@ impl HealthState {
         *self.last_update.lock().await = Some(Utc::now());
 
         for (id, status) in &self.health {
-            let url = &self.config.services[id].url;
+            let service_config = &self.config.services[id];
+            let url = &service_config.url;
             info!("service [{}] querying {}", id, url);
             let new_status = match reqwest::get(url).await {
                 Err(_) => HealthStatus::Down,
-                Ok(response) if response.status().is_success() => HealthStatus::Up,
-                Ok(_response) => HealthStatus::Errored,
+                Ok(response) => service_config.mode.status_of(&response),
             };
             info!("service [{}] new status {:?}", id, new_status);
             *status.lock().await = Some(new_status);
         }
     }
 }
+
+impl HttpHealthCheckMode {
+    fn status_of(&self, response: &Response) -> HealthStatus {
+        let response_status = response.status();
+
+        fn validate_status(
+            expected: &PartialStatusCode,
+            actual: StatusCode,
+            if_valid: HealthStatus,
+        ) -> Option<HealthStatus> {
+            let if_valid = Some(if_valid);
+            match expected {
+                PartialStatusCode::Any => if_valid,
+                &PartialStatusCode::Status(code) => StatusCode::from_u16(code)
+                    .map_err(|e| error!("invalid status code: {}", e))
+                    .ok()
+                    .filter(|status| status == &actual)
+                    .and_then(|_| if_valid),
+                PartialStatusCode::Status5XX if actual.is_server_error() => if_valid,
+                PartialStatusCode::Status4XX if actual.is_server_error() => if_valid,
+                PartialStatusCode::Status3XX if actual.is_redirection() => if_valid,
+                PartialStatusCode::Status2XX if actual.is_success() => if_valid,
+                PartialStatusCode::Status1XX if actual.is_informational() => if_valid,
+                _ => None,
+            }
+        }
+
+        // Check if response status matches expected status codes for Up
+        let check_up = self
+            .up_status_codes
+            .iter()
+            .flat_map(|up_code| validate_status(&up_code, response_status, HealthStatus::Up));
+
+        // Check if response status matches expected status codes for Down
+        let check_down = self
+            .down_status_codes
+            .iter()
+            .flat_map(|down_code| validate_status(&down_code, response_status, HealthStatus::Down));
+
+        // Compute status, defaulting to Errored if neigher Up nor Down matched
+        check_up
+            .chain(check_down)
+            .next()
+            .unwrap_or(HealthStatus::Errored)
+    }
+}