From e84581ee2c0cdcd21add369019bbc3c9d8c9c2bc Mon Sep 17 00:00:00 2001
From: MacRat <m@crat.jp>
Date: Sat, 17 Dec 2022 21:00:59 +0900
Subject: [PATCH] feat(scheme): set ping's privilege mode automatically

---
 README.md                       |  4 ---
 internal/scheme/ping.go         | 43 ++++++++-------------------------
 internal/scheme/ping_test.go    | 41 -------------------------------
 internal/scheme/ping_unix.go    | 22 +++++++++++++++++
 internal/scheme/ping_windows.go | 15 ++++++++++++
 5 files changed, 47 insertions(+), 78 deletions(-)
 create mode 100644 internal/scheme/ping_unix.go
 create mode 100644 internal/scheme/ping_windows.go

diff --git a/README.md b/README.md
index 2b843ed..e800e7e 100644
--- a/README.md
+++ b/README.md
@@ -180,10 +180,6 @@ Send ICMP echo request (a.k.a. ping command) and check if the target is connecte
 Ayd sends 3 packets in 1 second and expects all packets to return.
 These parameter can changed by `AYD_PING_PACKETS` and `AYD_PING_PERIOD` environment variable.
 
-In Linux or MacOS, Ayd use non-privileged ICMP in default. So, you can use ping even if rootless.
-But this way is not work on some platforms for instance docker container.
-Please set `yes` to `AYD_PRIVILEGED` environment variable to use privileged ICMP.
-
 You can specify IPv4 or IPv6 with `ping4:` or `ping6:` scheme.
 
 Ping will timeout in 30 seconds after sent all packets and report as failure.
diff --git a/internal/scheme/ping.go b/internal/scheme/ping.go
index 5978085..4d9401d 100644
--- a/internal/scheme/ping.go
+++ b/internal/scheme/ping.go
@@ -98,39 +98,13 @@ func newAutoPinger() *autoPingerStruct {
 	}
 }
 
-func getAydPrivilegedEnv() bool {
-	switch strings.ToLower(os.Getenv("AYD_PRIVILEGED")) {
-	case "", "0", "no", "false":
-		return false
-	}
-	return true
-}
-
-func makePingers() (v4, v6 *pinger.Pinger) {
-	v4 = pinger.NewIPv4()
-	v6 = pinger.NewIPv6()
-
-	if getAydPrivilegedEnv() {
-		v4.SetPrivileged(true)
-		v6.SetPrivileged(true)
-	}
-
-	return v4, v6
-}
-
 func (p *autoPingerStruct) start() (teardown func(), err error) {
-	p.v4, p.v6 = makePingers()
+	p.v4 = pinger.NewIPv4()
+	p.v6 = pinger.NewIPv6()
 
 	ctx, stop := context.WithCancel(context.Background())
 
-	if err := p.v4.Start(ctx); err != nil {
-		stop()
-		p.v4 = nil
-		p.v6 = nil
-		return nil, err
-	}
-
-	if err := p.v6.Start(ctx); err != nil {
+	if err := p.startPingers(ctx); err != nil {
 		stop()
 		p.v4 = nil
 		p.v6 = nil
@@ -212,11 +186,14 @@ var (
 
 // checkPingPermission tries to prepare pinger for check if it has permission.
 func checkPingPermission() error {
-	ctx, stop := context.WithCancel(context.Background())
-	defer stop()
+	stop, err := autoPinger.start()
 
-	p, _ := makePingers()
-	return p.Start(ctx)
+	if err != nil {
+		return err
+	}
+
+	stop()
+	return nil
 }
 
 // PingProbe is a Prober implementation for SNMP echo request aka ping.
diff --git a/internal/scheme/ping_test.go b/internal/scheme/ping_test.go
index cc5198f..7cfa578 100644
--- a/internal/scheme/ping_test.go
+++ b/internal/scheme/ping_test.go
@@ -2,8 +2,6 @@ package scheme_test
 
 import (
 	"context"
-	"errors"
-	"os"
 	"runtime"
 	"strings"
 	"testing"
@@ -102,45 +100,6 @@ func TestPingProbe_Probe(t *testing.T) {
 	})
 }
 
-func TestPingProbe(t *testing.T) {
-	if runtime.GOOS == "windows" {
-		t.Skip("unprivileged mode is not supported on windows")
-	}
-
-	privileged := os.Getenv("AYD_PRIVILEGED")
-	t.Cleanup(func() {
-		os.Setenv("AYD_PRIVILEGED", privileged)
-	})
-
-	tests := []struct {
-		Env  string
-		Fail bool
-	}{
-		{"1", true},
-		{"0", false},
-		{"yes", true},
-		{"no", false},
-		{"true", true},
-		{"false", false},
-		{"TRUE", true},
-		{"False", false},
-	}
-
-	for _, tt := range tests {
-		t.Run("AYD_PRIVILEGED="+tt.Env, func(t *testing.T) {
-			os.Setenv("AYD_PRIVILEGED", tt.Env)
-			_, err := scheme.NewPingProbe(&api.URL{Scheme: "ping", Opaque: "localhost"})
-
-			if tt.Fail && !errors.Is(err, scheme.ErrFailedToPreparePing) {
-				t.Errorf("expected permission error but got %v", err)
-			}
-			if !tt.Fail && err != nil {
-				t.Errorf("expected nil but got error: %s", err)
-			}
-		})
-	}
-}
-
 func BenchmarkPingProbe(b *testing.B) {
 	p := testutil.NewProber(b, "ping:localhost")
 
diff --git a/internal/scheme/ping_unix.go b/internal/scheme/ping_unix.go
new file mode 100644
index 0000000..5b247ae
--- /dev/null
+++ b/internal/scheme/ping_unix.go
@@ -0,0 +1,22 @@
+//go:build linux || darwin
+// +build linux darwin
+
+package scheme
+
+import (
+	"context"
+)
+
+func (p *autoPingerStruct) startPingers(ctx context.Context) error {
+	if err := p.v4.Start(ctx); err == nil {
+		return p.v6.Start(ctx)
+	}
+
+	p.v4.SetPrivileged(true)
+	p.v4.SetPrivileged(false)
+
+	if err := p.v4.Start(ctx); err != nil {
+		return err
+	}
+	return p.v6.Start(ctx)
+}
diff --git a/internal/scheme/ping_windows.go b/internal/scheme/ping_windows.go
new file mode 100644
index 0000000..539c26d
--- /dev/null
+++ b/internal/scheme/ping_windows.go
@@ -0,0 +1,15 @@
+//go:build windows
+// +build windows
+
+package scheme
+
+import (
+	"context"
+)
+
+func (p *autoPingerStruct) startPingers(ctx context.Context) error {
+	if err := p.v4.Start(ctx); err != nil {
+		return err
+	}
+	return p.v6.Start(ctx)
+}