Refactored badly grown code to make service detection more robust

Fixes falsely detecting unknown tls service on badly behaving ports

leech/src/modules/os_detection/mod.rs

@@ -20,6 +20,7 @@ use crate::modules::os_detection::syn_scan::find_open_and_closed_port;
 use crate::modules::os_detection::tcp_fingerprint::fingerprint_tcp;
 use crate::modules::os_detection::tcp_fingerprint::TcpFingerprint;
 use crate::modules::service_detection::tcp::OneShotTcpSettings;
+use crate::modules::service_detection::tcp::ProbeTcpResult;
 
 pub mod errors;
 mod fingerprint_db;
@@ -681,17 +682,15 @@ async fn os_detect_ssh(
         return Ok(OperatingSystemInfo::default());
     };
 
-    let Ok(result) = result else {
+    let ProbeTcpResult::Ok(data) = result else {
         // TOOD: might want to return differently if the error is a specific one, but right now it's a dynamic error
         // without proper error information for us to match on.
         return Ok(OperatingSystemInfo::default());
     };
 
-    if let Some(data) = result {
-        if data.starts_with(b"SSH-") {
-            if let Some(end) = data.iter().find_position(|&&c| c == b'\r' || c == b'\n') {
-                return Ok(os_detect_ssh_header(&data[0..end.0]));
-            }
+    if data.starts_with(b"SSH-") {
+        if let Some(end) = data.iter().find_position(|&&c| c == b'\r' || c == b'\n') {
+            return Ok(os_detect_ssh_header(&data[0..end.0]));
         }
     }
 

leech/src/modules/service_detection/error.rs

@@ -26,19 +26,6 @@ impl<E: Error> Deref for Extended<E> {
     }
 }
 
-/// Syntax extension which allows adding a context string to any error
-pub trait ErrorExt: Error + Sized {
-    fn context(self, msg: impl Into<Cow<'static, str>>) -> Extended<Self>;
-}
-impl<E: Error + Sized> ErrorExt for E {
-    fn context(self, msg: impl Into<Cow<'static, str>>) -> Extended<Self> {
-        Extended {
-            inner: self,
-            context: msg.into(),
-        }
-    }
-}
-
 /// Syntax extension which allows adding a context string to a result's error
 pub trait ResultExt {
     /// The `Result`'s `Ok` type

leech/src/modules/service_detection/generated.rs

@@ -8,7 +8,6 @@ use regex::bytes::Regex;
 use tokio::net::UdpSocket;
 
 use crate::modules::service_detection::tcp::OneShotTcpSettings;
-use crate::modules::service_detection::DynResult;
 
 include!(concat!(env!("OUT_DIR"), "/generated_probes.rs"));
 
@@ -23,10 +22,9 @@ pub struct AllProbes {
     pub rust_udp_probes: [Vec<RustProbe<UdpFn>>; 3],
 }
 
-pub type TcpFn = for<'a> fn(&'a OneShotTcpSettings) -> BoxFuture<'a, DynResult<Match>>;
-pub type TlsFn =
-    for<'a> fn(&'a OneShotTcpSettings, Option<&'static str>) -> BoxFuture<'a, DynResult<Match>>;
-pub type UdpFn = for<'a> fn(&'a mut UdpSocket) -> BoxFuture<'a, DynResult<Match>>;
+pub type TcpFn = for<'a> fn(&'a OneShotTcpSettings) -> BoxFuture<'a, Match>;
+pub type TlsFn = for<'a> fn(&'a OneShotTcpSettings, Option<&'static str>) -> BoxFuture<'a, Match>;
+pub type UdpFn = for<'a> fn(&'a mut UdpSocket) -> BoxFuture<'a, Match>;
 
 /// A probe implemented in rust
 pub struct RustProbe<F> {

leech/src/modules/service_detection/probe_impls/postgres.rs

@@ -9,35 +9,36 @@ use log::trace;
 
 use crate::modules::service_detection::generated::Match;
 use crate::modules::service_detection::tcp::OneShotTcpSettings;
-use crate::modules::service_detection::DynResult;
+use crate::modules::service_detection::tcp::ProbeTcpResult;
+use crate::modules::service_detection::tcp::ProbeTlsResult;
 use crate::utils::DebuggableBytes;
 
-pub async fn probe_tcp(settings: &OneShotTcpSettings) -> DynResult<Match> {
-    let Some(data) = settings.probe_tcp(&create_startup_message()).await? else {
-        return Ok(Match::No);
+pub async fn probe_tcp(settings: &OneShotTcpSettings) -> Match {
+    let ProbeTcpResult::Ok(data) = settings.probe_tcp(&create_startup_message()).await else {
+        return Match::No;
     };
     trace!(target: "postgres", "Got data: {data:x?}");
-    Ok(if parse_response(data).is_some() {
+    if parse_response(data).is_some() {
         Match::Exact
     } else {
         Match::No
-    })
+    }
 }
 
 // NB: postgres' protocol provides no method to enforce ssl.
 // Instead admins could block the auth step when the connection is unencrypted.
 // Therefore, no special treatment for ssl should be required.
 // https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-SSL
-pub async fn probe_tls(settings: &OneShotTcpSettings, alpn: Option<&str>) -> DynResult<Match> {
-    let Ok(Some(data)) = settings.probe_tls(&create_startup_message(), alpn).await? else {
-        return Ok(Match::No);
+pub async fn probe_tls(settings: &OneShotTcpSettings, alpn: Option<&str>) -> Match {
+    let ProbeTlsResult::Ok(data) = settings.probe_tls(&create_startup_message(), alpn).await else {
+        return Match::No;
     };
     trace!(target: "postgres", "Got data: {data:x?}");
-    Ok(if parse_response(data).is_some() {
+    if parse_response(data).is_some() {
         Match::Exact
     } else {
         Match::No
-    })
+    }
 }
 
 fn create_startup_message() -> Vec<u8> {