From e5004d84ec9500eb9ee7c18ea5f481129f93ca58 Mon Sep 17 00:00:00 2001
From: Ti Chi Robot <ti-community-prow-bot@tidb.io>
Date: Wed, 18 Oct 2023 13:22:59 +0800
Subject: [PATCH] lightning: fix lightning TLS.WithHost to access the correct
 host (#45749) (#47734)

close pingcap/tidb#45747
---
 br/pkg/lightning/common/BUILD.bazel      |  2 +-
 br/pkg/lightning/common/security.go      |  8 +++++
 br/pkg/lightning/common/security_test.go | 43 ++++++++++++++++++++++++
 br/tests/run.sh                          |  1 +
 4 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/br/pkg/lightning/common/BUILD.bazel b/br/pkg/lightning/common/BUILD.bazel
index 15cc00048753d..44b370a799cde 100644
--- a/br/pkg/lightning/common/BUILD.bazel
+++ b/br/pkg/lightning/common/BUILD.bazel
@@ -102,7 +102,7 @@ go_test(
     ],
     embed = [":common"],
     flaky = True,
-    shard_count = 19,
+    shard_count = 20,
     deps = [
         "//br/pkg/errors",
         "//br/pkg/lightning/log",
diff --git a/br/pkg/lightning/common/security.go b/br/pkg/lightning/common/security.go
index 46a5ffb821ffc..9932b10c8ed9b 100644
--- a/br/pkg/lightning/common/security.go
+++ b/br/pkg/lightning/common/security.go
@@ -20,6 +20,7 @@ import (
 	"net"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/br/pkg/httputil"
@@ -88,8 +89,15 @@ func NewTLSFromMockServer(server *httptest.Server) *TLS {
 	}
 }
 
+// GetMockTLSUrl returns tls's host for mock test
+func GetMockTLSUrl(tls *TLS) string {
+	return tls.url
+}
+
 // WithHost creates a new TLS instance with the host replaced.
 func (tc *TLS) WithHost(host string) *TLS {
+	host = strings.TrimPrefix(host, "http://")
+	host = strings.TrimPrefix(host, "https://")
 	var url string
 	if tc.inner != nil {
 		url = "https://" + host
diff --git a/br/pkg/lightning/common/security_test.go b/br/pkg/lightning/common/security_test.go
index e34ef3622500c..4ba9825efc883 100644
--- a/br/pkg/lightning/common/security_test.go
+++ b/br/pkg/lightning/common/security_test.go
@@ -70,6 +70,49 @@ func TestGetJSONSecure(t *testing.T) {
 	require.Equal(t, "/dddd", result.Path)
 }
 
+func TestWithHost(t *testing.T) {
+	mockTLSServer := httptest.NewTLSServer(http.HandlerFunc(respondPathHandler))
+	defer mockTLSServer.Close()
+	mockServer := httptest.NewServer(http.HandlerFunc(respondPathHandler))
+	defer mockServer.Close()
+
+	testCases := []struct {
+		expected string
+		host     string
+		secure   bool
+	}{
+		{
+			"https://127.0.0.1:2379",
+			"http://127.0.0.1:2379",
+			true,
+		},
+		{
+			"http://127.0.0.1:2379",
+			"https://127.0.0.1:2379",
+			false,
+		},
+		{
+			"http://127.0.0.1:2379/pd/api/v1/stores",
+			"127.0.0.1:2379/pd/api/v1/stores",
+			false,
+		},
+		{
+			"https://127.0.0.1:2379",
+			"127.0.0.1:2379",
+			true,
+		},
+	}
+
+	for _, testCase := range testCases {
+		server := mockServer
+		if testCase.secure {
+			server = mockTLSServer
+		}
+		tls := common.NewTLSFromMockServer(server)
+		require.Equal(t, testCase.expected, common.GetMockTLSUrl(tls.WithHost(testCase.host)))
+	}
+}
+
 func TestInvalidTLS(t *testing.T) {
 	tempDir := t.TempDir()
 	caPath := filepath.Join(tempDir, "ca.pem")
diff --git a/br/tests/run.sh b/br/tests/run.sh
index d8b471b00c932..7e934131193d5 100755
--- a/br/tests/run.sh
+++ b/br/tests/run.sh
@@ -17,6 +17,7 @@
 set -eu
 export PATH="tests/_utils:bin:$PATH"
 export TEST_DIR=/tmp/backup_restore_test
+export COV_DIR="/tmp/group_cover"
 
 # Create COV_DIR if not exists
 if [ -d "$COV_DIR" ]; then