From 774a84300536cba8b1c6e3e1f3f2d21e3bfe7cab Mon Sep 17 00:00:00 2001
From: Paul-Christian Volkmer <code@pcvolkmer.de>
Date: Mon, 26 Dec 2022 19:05:50 +0100
Subject: [PATCH] Add check for plain TCP connections

---
 README.adoc |  8 +++++++-
 src/main.rs | 33 ++++++++++++++++++++++++++++++++-
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/README.adoc b/README.adoc
index faecc51..8397da3 100644
--- a/README.adoc
+++ b/README.adoc
@@ -24,6 +24,11 @@ url = "https://host2.example.com"
 name = "App 1"
 url = "https://app.example.com/actuator/health"
 check_type = "Actuator"
+
+[[checks]]
+name = "App 2"
+url = "tcp://app.example.com:12345"
+check_type = "Tcp"
 ----
 
 Each host or application to be checked consists of `name` and `url`.
@@ -33,8 +38,9 @@ You can optionally specify `check_type`:
 * `Html`: Default value, checks if a request is successful and returns HTTP OK - 200.
 * `Actuator`: Like `Html`, but checks if _Actuator_ shows that the application is up and running.
   ** `Up`: Actuator response indicates, application is up and running.
-  ** `Warn`: Got response but it is not Actuator health `UP` response or redirect.
+  ** `Warn`: Got response, but it is not Actuator health `UP` response or redirect.
   ** `Down`: No response.
+* `Tcp`: Checks if TCP connection to given host and port can be established
 
 To use more than one configuration, pass the config file location as first argument to the application.
 
diff --git a/src/main.rs b/src/main.rs
index 2ec7a5e..4076499 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,10 +1,13 @@
+use reqwest::Url;
 use serde::Deserialize;
 use serde_json::json;
 use std::env;
 use std::fmt::{Display, Formatter, Result};
 use std::fs;
 use std::process;
+use std::str::FromStr;
 use std::time::Duration;
+use tokio::net::TcpStream;
 use tokio::task;
 
 #[derive(Deserialize)]
@@ -29,10 +32,11 @@ struct CheckConfig {
     click_cmd: Option<String>,
 }
 
-#[derive(Deserialize)]
+#[derive(Deserialize, PartialEq)]
 enum CheckType {
     Http,
     Actuator,
+    Tcp,
 }
 
 #[derive(Deserialize)]
@@ -87,6 +91,33 @@ enum CheckState {
 }
 
 async fn check_host(check_config: &CheckConfig) -> CheckResult {
+    if check_config.check_type == Some(CheckType::Tcp) {
+        if let Ok(url) = Url::from_str(check_config.url.as_str()) {
+            if url.scheme() == "tcp" {
+                if url.host_str().is_some() && url.port().is_some() {
+                    let state = match TcpStream::connect(format!(
+                        "{}:{}",
+                        url.host_str().unwrap(),
+                        url.port().unwrap()
+                    ))
+                    .await
+                    {
+                        Ok(_) => CheckState::Up,
+                        _ => CheckState::Down,
+                    };
+                    return CheckResult {
+                        name: check_config.name.to_string(),
+                        state,
+                    };
+                }
+            }
+        }
+        return CheckResult {
+            name: check_config.name.to_string(),
+            state: CheckState::Down,
+        };
+    }
+
     let state = match reqwest::get(check_config.url.as_str()).await {
         Ok(r) => {
             if r.status().is_success() {