From da2fd6de275066c0c5873c86eb9c3a024bea41e5 Mon Sep 17 00:00:00 2001
From: okJiang <819421878@qq.com>
Date: Mon, 3 Jun 2024 16:00:40 +0800
Subject: [PATCH] fix https client

---
 tools/pd-ctl/pdctl/command/cluster_command.go |  8 +-
 tools/pd-ctl/pdctl/command/health_command.go  |  8 +-
 tools/pd-ctl/tests/health/health_test.go      | 84 +++++++++++++++++++
 3 files changed, 92 insertions(+), 8 deletions(-)

diff --git a/tools/pd-ctl/pdctl/command/cluster_command.go b/tools/pd-ctl/pdctl/command/cluster_command.go
index 0b3cda8a8677..cd9612571073 100644
--- a/tools/pd-ctl/pdctl/command/cluster_command.go
+++ b/tools/pd-ctl/pdctl/command/cluster_command.go
@@ -19,10 +19,10 @@ import "github.com/spf13/cobra"
 // NewClusterCommand return a cluster subcommand of rootCmd
 func NewClusterCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:               "cluster",
-		Short:             "show the cluster information",
-		PersistentPreRunE: requirePDClient,
-		Run:               showClusterCommandFunc,
+		Use:     "cluster",
+		Short:   "show the cluster information",
+		PreRunE: requirePDClient,
+		Run:     showClusterCommandFunc,
 	}
 	cmd.AddCommand(NewClusterStatusCommand())
 	return cmd
diff --git a/tools/pd-ctl/pdctl/command/health_command.go b/tools/pd-ctl/pdctl/command/health_command.go
index a10ee1183970..54aaa525184c 100644
--- a/tools/pd-ctl/pdctl/command/health_command.go
+++ b/tools/pd-ctl/pdctl/command/health_command.go
@@ -21,10 +21,10 @@ import (
 // NewHealthCommand return a health subcommand of rootCmd
 func NewHealthCommand() *cobra.Command {
 	m := &cobra.Command{
-		Use:               "health",
-		Short:             "show all node's health information of the PD cluster",
-		PersistentPreRunE: requirePDClient,
-		Run:               showHealthCommandFunc,
+		Use:     "health",
+		Short:   "show all node's health information of the PD cluster",
+		PreRunE: requirePDClient,
+		Run:     showHealthCommandFunc,
 	}
 	return m
 }
diff --git a/tools/pd-ctl/tests/health/health_test.go b/tools/pd-ctl/tests/health/health_test.go
index 9150a56c91b2..f1d3c7cfbf11 100644
--- a/tools/pd-ctl/tests/health/health_test.go
+++ b/tools/pd-ctl/tests/health/health_test.go
@@ -17,14 +17,21 @@ package health_test
 import (
 	"context"
 	"encoding/json"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"github.com/tikv/pd/pkg/utils/grpcutil"
 	"github.com/tikv/pd/server/api"
 	"github.com/tikv/pd/server/cluster"
+	"github.com/tikv/pd/server/config"
 	pdTests "github.com/tikv/pd/tests"
 	ctl "github.com/tikv/pd/tools/pd-ctl/pdctl"
 	"github.com/tikv/pd/tools/pd-ctl/tests"
+	"go.etcd.io/etcd/pkg/transport"
 )
 
 func TestHealth(t *testing.T) {
@@ -68,3 +75,80 @@ func TestHealth(t *testing.T) {
 	re.NoError(json.Unmarshal(output, &h))
 	re.Equal(healths, h)
 }
+
+func TestHealthTLS(t *testing.T) {
+	re := require.New(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	certPath := "../cert"
+	certScript := "../cert_opt.sh"
+	// generate certs
+	if err := os.Mkdir(certPath, 0755); err != nil {
+		t.Fatal(err)
+	}
+	if err := exec.Command(certScript, "generate", certPath).Run(); err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		if err := exec.Command(certScript, "cleanup", certPath).Run(); err != nil {
+			t.Fatal(err)
+		}
+		if err := os.RemoveAll(certPath); err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	tlsInfo := transport.TLSInfo{
+		KeyFile:       filepath.Join(certPath, "pd-server-key.pem"),
+		CertFile:      filepath.Join(certPath, "pd-server.pem"),
+		TrustedCAFile: filepath.Join(certPath, "ca.pem"),
+	}
+	tc, err := pdTests.NewTestCluster(ctx, 1, func(conf *config.Config, _ string) {
+		conf.Security.TLSConfig = grpcutil.TLSConfig{
+			KeyPath:  tlsInfo.KeyFile,
+			CertPath: tlsInfo.CertFile,
+			CAPath:   tlsInfo.TrustedCAFile,
+		}
+		conf.AdvertiseClientUrls = strings.ReplaceAll(conf.AdvertiseClientUrls, "http", "https")
+		conf.ClientUrls = strings.ReplaceAll(conf.ClientUrls, "http", "https")
+		conf.AdvertisePeerUrls = strings.ReplaceAll(conf.AdvertisePeerUrls, "http", "https")
+		conf.PeerUrls = strings.ReplaceAll(conf.PeerUrls, "http", "https")
+		conf.InitialCluster = strings.ReplaceAll(conf.InitialCluster, "http", "https")
+	})
+	re.NoError(err)
+	defer tc.Destroy()
+	err = tc.RunInitialServers()
+	re.NoError(err)
+	tc.WaitLeader()
+	cmd := ctl.GetRootCmd()
+
+	client := tc.GetEtcdClient()
+	members, err := cluster.GetMembers(client)
+	re.NoError(err)
+	healthMembers := cluster.CheckHealth(tc.GetHTTPClient(), members)
+	healths := []api.Health{}
+	for _, member := range members {
+		h := api.Health{
+			Name:       member.Name,
+			MemberID:   member.MemberId,
+			ClientUrls: member.ClientUrls,
+			Health:     false,
+		}
+		if _, ok := healthMembers[member.GetMemberId()]; ok {
+			h.Health = true
+		}
+		healths = append(healths, h)
+	}
+
+	pdAddr := tc.GetConfig().GetClientURL()
+	pdAddr = strings.ReplaceAll(pdAddr, "http", "https")
+	args := []string{"-u", pdAddr, "health",
+		"--cacert=../cert/ca.pem",
+		"--cert=../cert/client.pem",
+		"--key=../cert/client-key.pem"}
+	output, err := tests.ExecuteCommand(cmd, args...)
+	re.NoError(err)
+	h := make([]api.Health, len(healths))
+	re.NoError(json.Unmarshal(output, &h))
+	re.Equal(healths, h)
+}