From 3d1b4e566d95ea2d57f1ffc2db95e43078c329ff Mon Sep 17 00:00:00 2001
From: Nicolas Thumann <me@n-thumann.de>
Date: Wed, 7 Feb 2024 03:15:51 +0100
Subject: [PATCH] feat(client): Add network config for ICMP endpoint client
 (#661)

* feat(client): Add network to config

* feat(client): Use network client config for pinger

* feat(client): Add client network configuration and demo to README

* feat(client): Add tests for pinger using network config

* feat(client): Drop integration tests

* feat(client): Add comment to tests

* feat(client): Add tests

* Update README.md

---------

Co-authored-by: TwiN <twin@linux.com>
---
 README.md             |  1 +
 client/client.go      |  8 +++-----
 client/client_test.go | 26 ++++++++++++++++++++++++++
 client/config.go      |  4 ++++
 4 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index beff1f26..6c028a54 100644
--- a/README.md
+++ b/README.md
@@ -362,6 +362,7 @@ the client used to send the request.
 | `client.oauth2.scopes[]`               | A list of `scopes` which should be used for the `Client credentials flow`.  | required `[""]` |
 | `client.identity-aware-proxy`          | Google Identity-Aware-Proxy client configuration.                           | `{}`            |
 | `client.identity-aware-proxy.audience` | The Identity-Aware-Proxy audience. (client-id of the IAP oauth2 credential) | required `""`   |
+| `client.network`                       | The network to use for ICMP endpoint client (`ip`, `ip4` or `ip6`).         | `"ip"`          |
 
 > 📝 Some of these parameters are ignored based on the type of endpoint. For instance, there's no certificate involved
 in ICMP requests (ping), therefore, setting `client.insecure` to `true` for an endpoint of that type will not do anything.
diff --git a/client/client.go b/client/client.go
index 6d1d5a49..7a00a767 100644
--- a/client/client.go
+++ b/client/client.go
@@ -235,10 +235,7 @@ func ExecuteSSHCommand(sshClient *ssh.Client, body string, config *Config) (bool
 //
 // Note that this function takes at least 100ms, even if the address is 127.0.0.1
 func Ping(address string, config *Config) (bool, time.Duration) {
-	pinger, err := ping.NewPinger(address)
-	if err != nil {
-		return false, 0
-	}
+	pinger := ping.New(address)
 	pinger.Count = 1
 	pinger.Timeout = config.Timeout
 	// Set the pinger's privileged mode to true for every GOOS except darwin
@@ -247,7 +244,8 @@ func Ping(address string, config *Config) (bool, time.Duration) {
 	// Note that for this to work on Linux, Gatus must run with sudo privileges.
 	// See https://github.com/prometheus-community/pro-bing#linux
 	pinger.SetPrivileged(runtime.GOOS != "darwin")
-	err = pinger.Run()
+	pinger.SetNetwork(config.Network)
+	err := pinger.Run()
 	if err != nil {
 		return false, 0
 	}
diff --git a/client/client_test.go b/client/client_test.go
index 3e91dda8..3cb117a9 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -83,6 +83,32 @@ func TestPing(t *testing.T) {
 			t.Error("Round-trip time returned on failure should've been 0")
 		}
 	}
+	// Can't perform integration tests (e.g. pinging public targets by single-stacked hostname) here,
+	// because ICMP is blocked in the network of GitHub-hosted runners.
+	if success, rtt := Ping("127.0.0.1", &Config{Timeout: 500 * time.Millisecond, Network: "ip"}); !success {
+		t.Error("expected true")
+		if rtt == 0 {
+			t.Error("Round-trip time returned on failure should've been 0")
+		}
+	}
+	if success, rtt := Ping("::1", &Config{Timeout: 500 * time.Millisecond, Network: "ip"}); !success {
+		t.Error("expected true")
+		if rtt == 0 {
+			t.Error("Round-trip time returned on failure should've been 0")
+		}
+	}
+	if success, rtt := Ping("::1", &Config{Timeout: 500 * time.Millisecond, Network: "ip4"}); success {
+		t.Error("expected false, because the IP isn't an IPv4 address")
+		if rtt != 0 {
+			t.Error("Round-trip time returned on failure should've been 0")
+		}
+	}
+	if success, rtt := Ping("127.0.0.1", &Config{Timeout: 500 * time.Millisecond, Network: "ip6"}); success {
+		t.Error("expected false, because the IP isn't an IPv6 address")
+		if rtt != 0 {
+			t.Error("Round-trip time returned on failure should've been 0")
+		}
+	}
 }
 
 func TestCanPerformStartTLS(t *testing.T) {
diff --git a/client/config.go b/client/config.go
index 4f0ad781..aa514e5a 100644
--- a/client/config.go
+++ b/client/config.go
@@ -30,6 +30,7 @@ var (
 		Insecure:       false,
 		IgnoreRedirect: false,
 		Timeout:        defaultTimeout,
+		Network:        "ip",
 	}
 )
 
@@ -64,6 +65,9 @@ type Config struct {
 	IAPConfig *IAPConfig `yaml:"identity-aware-proxy,omitempty"`
 
 	httpClient *http.Client
+
+	// Network (ip, ip4 or ip6) for the ICMP client
+	Network string `yaml:"network"`
 }
 
 // DNSResolverConfig is the parsed configuration from the DNSResolver config string.