From 157dc7aa8f1ef8628c81ae55d68dae1c58db06c6 Mon Sep 17 00:00:00 2001
From: Guo Jix <729324352@qq.com>
Date: Fri, 23 Aug 2024 03:15:44 +0000
Subject: [PATCH] use excludeLocalhost flag

---
 config.go         |  2 +-
 endpoints.go      |  4 ++--
 endpoints_test.go | 48 +++++++++++++++++++++++++++++++++++++++++++----
 main.go           |  2 +-
 4 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/config.go b/config.go
index 6615b4f..62c0768 100644
--- a/config.go
+++ b/config.go
@@ -10,7 +10,7 @@ import (
 type globalConfig struct {
 	endpoints           []string
 	useClusterEndpoints bool
-	remote              bool
+	excludeLocalhost    bool
 
 	dialTimeout      time.Duration
 	commandTimeout   time.Duration
diff --git a/endpoints.go b/endpoints.go
index c82c8e6..ff66175 100644
--- a/endpoints.go
+++ b/endpoints.go
@@ -40,7 +40,7 @@ func endpoints(gcfg globalConfig) ([]string, error) {
 	return endpointsFromCluster(gcfg)
 }
 
-func isLocalEp(ep string) bool {
+func IsLocalEp(ep string) bool {
 	return strings.Contains(ep, "127.0.0.1") || strings.Contains(ep, "localhost")
 }
 
@@ -56,7 +56,7 @@ func endpointsFromCluster(gcfg globalConfig) ([]string, error) {
 		if !m.GetIsLearner() {
 			for _, ep := range m.ClientURLs {
 				// filter local endpoint when set remote
-				if !gcfg.remote || !isLocalEp(ep) {
+				if !gcfg.excludeLocalhost || !IsLocalEp(ep) {
 					eps = append(eps, ep)
 				}
 			}
diff --git a/endpoints_test.go b/endpoints_test.go
index 577069a..2255d74 100644
--- a/endpoints_test.go
+++ b/endpoints_test.go
@@ -111,7 +111,6 @@ func TestEndpointDedup(t *testing.T) {
 	}
 }
 
-
 func TestEndpointRemote(t *testing.T) {
 	oldCreateClient := createClient
 	t.Cleanup(func() {
@@ -127,7 +126,7 @@ func TestEndpointRemote(t *testing.T) {
 		name               string
 		returnedMemberList *clientv3.MemberListResponse
 		expectedEndpoints  []string
-		remote             bool
+		excludeLocalhost   bool
 	}{
 		{
 			"normal",
@@ -164,7 +163,7 @@ func TestEndpointRemote(t *testing.T) {
 	for _, testcase := range testcases {
 		t.Run(testcase.name, func(t *testing.T) {
 			fakeClient.memberListResp = testcase.returnedMemberList
-			ep, err := endpointsFromCluster(globalConfig{endpoints: []string{"https://localhost:2379"}, remote: testcase.remote})
+			ep, err := endpointsFromCluster(globalConfig{endpoints: []string{"https://localhost:2379"}, excludeLocalhost: testcase.excludeLocalhost})
 			if err != nil {
 				t.Error(err)
 			}
@@ -176,7 +175,6 @@ func TestEndpointRemote(t *testing.T) {
 	}
 }
 
-
 type fakeClientURLClient struct {
 	*clientv3.Client
 	memberListResp *clientv3.MemberListResponse
@@ -190,3 +188,45 @@ func (f fakeClientURLClient) MemberList(ctx context.Context, opts ...clientv3.Op
 func (fakeClientURLClient) Close() error {
 	return nil
 }
+
+func TestIsLocalEp(t *testing.T) {
+	testcases := []struct {
+		name   string
+		ep     string
+		desire bool
+	}{
+		{
+			"exclude",
+			"127.0.0.1:2379",
+			true,
+		},
+		{
+			"include",
+			"10.7.7.7:2379",
+			false,
+		},
+		{
+			"excludehttp",
+			"http://127.0.0.1:2379",
+			true,
+		},
+		{
+			"includehttps",
+			"http://10.7.7.7:2379",
+			false,
+		},
+		{
+			"includehttpsdomain",
+			"https://abc-0.ns1-etcd.ns1.svc.cluster.local.:2379",
+			false,
+		},
+	}
+
+	for _, testcase := range testcases {
+		t.Run(testcase.name, func(t *testing.T) {
+			if IsLocalEp(testcase.ep) != testcase.desire {
+				t.Errorf("expected %v, got %v", testcase.desire, IsLocalEp(testcase.ep))
+			}
+		})
+	}
+}
diff --git a/main.go b/main.go
index f79bd33..c0fd918 100644
--- a/main.go
+++ b/main.go
@@ -23,7 +23,7 @@ func newDefragCommand() *cobra.Command {
 
 	defragCmd.Flags().StringSliceVar(&globalCfg.endpoints, "endpoints", []string{"127.0.0.1:2379"}, "comma separated etcd endpoints")
 	defragCmd.Flags().BoolVar(&globalCfg.useClusterEndpoints, "cluster", false, "use all endpoints from the cluster member list")
-	defragCmd.Flags().BoolVar(&globalCfg.remote, "remote", true, "whether execute on remote machine")
+	defragCmd.Flags().BoolVar(&globalCfg.excludeLocalhost, "remote", true, "whether to exclude localhost endpoints")
 
 	defragCmd.Flags().DurationVar(&globalCfg.dialTimeout, "dial-timeout", 2*time.Second, "dial timeout for client connections")
 	defragCmd.Flags().DurationVar(&globalCfg.commandTimeout, "command-timeout", 30*time.Second, "command timeout (excluding dial timeout)")