From 2e62a7646116d1b2eb5dedcafe63ad02b5459c46 Mon Sep 17 00:00:00 2001
From: Xavier Vello <xavier@posthog.com>
Date: Mon, 6 May 2024 13:40:12 +0200
Subject: [PATCH] add tests

---
 hook-worker/src/dns.rs    | 50 +++++++++++++++++++++++++++++++++++++++
 hook-worker/src/error.rs  |  5 ++--
 hook-worker/src/worker.rs |  4 ++--
 3 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/hook-worker/src/dns.rs b/hook-worker/src/dns.rs
index ca444ec..54c574b 100644
--- a/hook-worker/src/dns.rs
+++ b/hook-worker/src/dns.rs
@@ -88,3 +88,53 @@ impl Resolve for PublicIPv4Resolver {
         Box::pin(future_result)
     }
 }
+
+mod tests {
+    use super::*;
+    use std::str::FromStr;
+
+    #[tokio::test]
+    async fn it_resolves_google_com() {
+        let resolver: PublicIPv4Resolver = PublicIPv4Resolver {};
+        let addrs = resolver
+            .resolve(Name::from_str("google.com").unwrap())
+            .await
+            .expect("lookup has failed");
+        assert!(addrs.count() > 0, "empty address list")
+    }
+
+    #[tokio::test]
+    async fn it_denies_ipv6_google_com() {
+        let resolver: PublicIPv4Resolver = PublicIPv4Resolver {};
+        match resolver
+            .resolve(Name::from_str("ipv6.google.com").unwrap())
+            .await
+        {
+            Ok(_) => panic!("should have failed"),
+            Err(err) => assert!(err.downcast_ref::<NoPublicIPError>().is_some()),
+        }
+    }
+
+    #[tokio::test]
+    async fn it_denies_localhost() {
+        let resolver: PublicIPv4Resolver = PublicIPv4Resolver {};
+        match resolver.resolve(Name::from_str("localhost").unwrap()).await {
+            Ok(_) => panic!("should have failed"),
+            Err(err) => assert!(err.is::<NoPublicIPError>()),
+        }
+    }
+
+    #[tokio::test]
+    async fn it_propagates_unknown_domain() {
+        let resolver: PublicIPv4Resolver = PublicIPv4Resolver {};
+        match resolver
+            .resolve(Name::from_str("invalid.domain.unknown").unwrap())
+            .await
+        {
+            Ok(_) => panic!("should have failed"),
+            Err(err) => assert!(err
+                .to_string()
+                .contains("failed to lookup address information")),
+        }
+    }
+}
diff --git a/hook-worker/src/error.rs b/hook-worker/src/error.rs
index 207049e..68a077e 100644
--- a/hook-worker/src/error.rs
+++ b/hook-worker/src/error.rs
@@ -67,7 +67,7 @@ impl fmt::Display for WebhookRequestError {
                     None => "No response from the server".to_string(),
                 };
                 if is_error_source::<NoPublicIPError>(error) {
-                    writeln!(f, "{}: {}", error ,NoPublicIPError)?;
+                    writeln!(f, "{}: {}", error, NoPublicIPError)?;
                 } else {
                     writeln!(f, "{}", error)?;
                 }
@@ -140,8 +140,9 @@ pub enum WorkerError {
 }
 
 /// Check the error and it's sources (recursively) to return true if an error of the given type is found.
+/// TODO: use Error::sources() when stable
 pub fn is_error_source<T: Error + 'static>(err: &(dyn std::error::Error + 'static)) -> bool {
-    if err.downcast_ref::<T>().is_some() {
+    if err.is::<NoPublicIPError>() {
         return true;
     }
     match err.source() {
diff --git a/hook-worker/src/worker.rs b/hook-worker/src/worker.rs
index 7fc3d0e..7e2693e 100644
--- a/hook-worker/src/worker.rs
+++ b/hook-worker/src/worker.rs
@@ -577,8 +577,8 @@ mod tests {
             webhook_job_parameters.clone(),
             webhook_job_metadata,
         )
-            .await
-            .expect("failed to enqueue job");
+        .await
+        .expect("failed to enqueue job");
         let worker = WebhookWorker::new(
             &worker_id,
             &queue,