From b22a62016b52bf708916c5761b25a46b7a046495 Mon Sep 17 00:00:00 2001
From: clerie <git@clerie.de>
Date: Mon, 2 Jun 2025 19:25:50 +0200
Subject: [PATCH] Check server response for some simple payload

---
 Cargo.lock  | 39 +++++++++++++++++++++++++++++++++++++++
 Cargo.toml  |  1 +
 src/main.rs | 42 ++++++++++++++++++++++++++++++++++++++----
 3 files changed, 78 insertions(+), 4 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index c02053b..09aec7d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -17,6 +17,15 @@ version = "2.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "512761e0bb2578dd7380c6baaa0f4ce03e84f95e960231d1dec8bf4d7d6e2627"
 
+[[package]]
+name = "aho-corasick"
+version = "1.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916"
+dependencies = [
+ "memchr",
+]
+
 [[package]]
 name = "axum"
 version = "0.8.4"
@@ -340,6 +349,35 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "regex"
+version = "1.11.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b544ef1b4eac5dc2db33ea63606ae9ffcfac26c1416a2806ae0bf5f56b201191"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-automata",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-automata"
+version = "0.4.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "809e8dc61f6de73b46c85f4c96486310fe304c434cfa43669d7b40f711150908"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-syntax"
+version = "0.8.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c"
+
 [[package]]
 name = "rustc-demangle"
 version = "0.1.24"
@@ -619,6 +657,7 @@ name = "xmpp-blackbox-exporter"
 version = "0.1.0"
 dependencies = [
  "axum",
+ "regex",
  "serde",
  "tokio",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index 0107e0f..bce2175 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -5,5 +5,6 @@ edition = "2021"
 
 [dependencies]
 axum = "0.8.4"
+regex = "1.11.1"
 serde = { version = "1.0.219", features = ["derive"] }
 tokio = { version = "1.45.1", features = ["io-util", "rt-multi-thread"] }
diff --git a/src/main.rs b/src/main.rs
index 3a47fb8..41447f5 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,6 +3,7 @@ use axum::{
     routing::get,
     Router,
 };
+use regex::Regex;
 use serde::Deserialize;
 use tokio::io::{
     AsyncBufReadExt,
@@ -20,6 +21,24 @@ struct ProbeClientToServerQuery {
     port: u16,
 }
 
+struct ProbeFacts {
+    is_xmpp_client: bool,
+    is_xmpp_stream: bool,
+}
+
+impl ProbeFacts {
+    fn new() -> Self {
+        Self {
+            is_xmpp_client: false,
+            is_xmpp_stream: false,
+        }
+    }
+
+    fn probe_success(&self) -> bool {
+        self.is_xmpp_client
+        && self.is_xmpp_stream
+    }
+}
 
 #[tokio::main]
 async fn main() {
@@ -63,7 +82,9 @@ async fn route_index() -> String {
     return String::from("Prometheus exporter for checking XMPP server availability");
 }
 
-async fn probe_client_to_server(domain: &str, hostname: &str, port: u16) -> Result<(), String> {
+async fn probe_client_to_server(domain: &str, hostname: &str, port: u16) -> Result<ProbeFacts, String> {
+    let mut probe_facts = ProbeFacts::new();
+
     let mut stream = TcpStream::connect((hostname, port)).await.unwrap();
 
     let connect_string = format!("<stream:stream to='{}' xmlns:stream='http://etherx.jabber.org/streams' xmlns='jabber:client' xml:lang='en' version='1.0'></stream:stream>", domain);
@@ -77,17 +98,30 @@ async fn probe_client_to_server(domain: &str, hostname: &str, port: u16) -> Resu
 
     println!("{}", response);
 
-    Ok(())
+    let re_match_xmpp_stream = Regex::new(r"http://etherx\.jabber\.org/streams").unwrap();
+    if re_match_xmpp_stream.is_match(&response) {
+        probe_facts.is_xmpp_stream = true;
+    }
+
+    let re_match_xmpp_client = Regex::new(r"jabber:client").unwrap();
+    if re_match_xmpp_client.is_match(&response) {
+        probe_facts.is_xmpp_client = true;
+    }
+
+    Ok(probe_facts)
 }
 
 async fn route_probe_client_to_server(
     Query(query): Query<ProbeClientToServerQuery>,
 ) -> Result<String, String> {
 
-    probe_client_to_server(&query.domain, &query.hostname, query.port).await.unwrap();
+    let probe_facts = probe_client_to_server(&query.domain, &query.hostname, query.port).await.unwrap();
 
     let mut out = String::new();
-    out.push_str(&format!("{} {}:{}", query.domain, query.hostname, query.port));
+
+    out.push_str(&format!("xmpp_probe_is_xmpp_client{{domain=\"{}\", hostname=\"{}\", port=\"{}\"}} {}\n", query.domain, query.hostname, query.port, probe_facts.is_xmpp_client as u8));
+    out.push_str(&format!("xmpp_probe_is_xmpp_stream{{domain=\"{}\", hostname=\"{}\", port=\"{}\"}} {}\n", query.domain, query.hostname, query.port, probe_facts.is_xmpp_stream as u8));
+    out.push_str(&format!("xmpp_probe_success{{domain=\"{}\", hostname=\"{}\", port=\"{}\"}} {}\n", query.domain, query.hostname, query.port, probe_facts.probe_success() as u8));
 
     return Ok(out);
 }