From bd02c603c02c4742098bfda1b88cdceb5044e1df Mon Sep 17 00:00:00 2001
From: Yujie Xia <xiayjchn@gmail.com>
Date: Mon, 11 Apr 2022 15:38:05 +0800
Subject: [PATCH 1/7] lightning: reload cert for new connection

---
 br/pkg/lightning/common/security.go | 39 ++++++++++++++++-------------
 1 file changed, 22 insertions(+), 17 deletions(-)

diff --git a/br/pkg/lightning/common/security.go b/br/pkg/lightning/common/security.go
index 9db53f78a5115..7f55c07603919 100644
--- a/br/pkg/lightning/common/security.go
+++ b/br/pkg/lightning/common/security.go
@@ -31,7 +31,6 @@ import (
 	"google.golang.org/grpc/credentials"
 )
 
-// TLS
 type TLS struct {
 	caPath   string
 	certPath string
@@ -50,16 +49,6 @@ func ToTLSConfig(caPath, certPath, keyPath string) (*tls.Config, error) {
 		return nil, nil
 	}
 
-	// Load the client certificates from disk
-	var certificates []tls.Certificate
-	if len(certPath) != 0 && len(keyPath) != 0 {
-		cert, err := tls.LoadX509KeyPair(certPath, keyPath)
-		if err != nil {
-			return nil, errors.Annotate(err, "could not load client key pair")
-		}
-		certificates = []tls.Certificate{cert}
-	}
-
 	// Create a certificate pool from CA
 	certPool := x509.NewCertPool()
 	ca, err := os.ReadFile(caPath)
@@ -72,12 +61,28 @@ func ToTLSConfig(caPath, certPath, keyPath string) (*tls.Config, error) {
 		return nil, errors.New("failed to append ca certs")
 	}
 
-	return &tls.Config{
-		Certificates: certificates,
-		RootCAs:      certPool,
-		NextProtos:   []string{"h2", "http/1.1"}, // specify `h2` to let Go use HTTP/2.
-		MinVersion:   tls.VersionTLS12,
-	}, nil
+	tlsConfig := &tls.Config{
+		RootCAs:    certPool,
+		NextProtos: []string{"h2", "http/1.1"}, // specify `h2` to let Go use HTTP/2.
+		MinVersion: tls.VersionTLS12,
+	}
+
+	if len(certPath) != 0 && len(keyPath) != 0 {
+		loadCert := func() (*tls.Certificate, error) {
+			cert, err := tls.LoadX509KeyPair(certPath, keyPath)
+			if err != nil {
+				return nil, errors.Annotate(err, "could not load client key pair")
+			}
+			return &cert, nil
+		}
+		tlsConfig.GetClientCertificate = func(*tls.CertificateRequestInfo) (*tls.Certificate, error) {
+			return loadCert()
+		}
+		tlsConfig.GetCertificate = func(info *tls.ClientHelloInfo) (*tls.Certificate, error) {
+			return loadCert()
+		}
+	}
+	return tlsConfig, nil
 }
 
 // NewTLS constructs a new HTTP client with TLS configured with the CA,

From 1d8398973ca11a2d1dc115e961362a1896c4d43a Mon Sep 17 00:00:00 2001
From: Yujie Xia <xiayjchn@gmail.com>
Date: Wed, 13 Apr 2022 15:29:57 +0800
Subject: [PATCH 2/7] fix test

---
 br/pkg/lightning/common/security_test.go | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/br/pkg/lightning/common/security_test.go b/br/pkg/lightning/common/security_test.go
index a68b8e1f4437a..3359935b73bf0 100644
--- a/br/pkg/lightning/common/security_test.go
+++ b/br/pkg/lightning/common/security_test.go
@@ -81,15 +81,29 @@ func TestInvalidTLS(t *testing.T) {
 	_, err = common.NewTLS(caPath, "", "", "localhost")
 	require.Regexp(t, "failed to append ca certs", err.Error())
 
+	err = os.WriteFile(caPath, []byte(`-----BEGIN CERTIFICATE-----
+MIIBITCBxwIUf04/Hucshr7AynmgF8JeuFUEf9EwCgYIKoZIzj0EAwIwEzERMA8G
+A1UEAwwIYnJfdGVzdHMwHhcNMjIwNDEzMDcyNDQxWhcNMjIwNDE1MDcyNDQxWjAT
+MREwDwYDVQQDDAhicl90ZXN0czBZMBMGByqGSM49AgEGCCqGSM49AwEHA0IABL+X
+wczUg0AbaFFaCI+FAk3K9vbB9JeIORgGKS+F1TKip5tvm96g7S5lq8SgY38SXVc3
+0yS3YqWZqnRjWi+sLwIwCgYIKoZIzj0EAwIDSQAwRgIhAJcpSwsUhqkM08LK1gYC
+ze4ZnCkwJdP2VdpI3WZsoI7zAiEAjP8X1c0iFwYxdAbQAveX+9msVrzyUpZOohi4
+RtgQTNI=
+-----END CERTIFICATE-----
+`), 0o644)
+	require.NoError(t, err)
+
 	certPath := filepath.Join(tempDir, "test.pem")
 	keyPath := filepath.Join(tempDir, "test.key")
-	_, err = common.NewTLS(caPath, certPath, keyPath, "localhost")
+	tls, err := common.NewTLS(caPath, certPath, keyPath, "localhost")
+	_, err = tls.TLSConfig().GetCertificate(nil)
 	require.Regexp(t, "could not load client key pair: open.*", err.Error())
 
 	err = os.WriteFile(certPath, []byte("invalid cert content"), 0o644)
 	require.NoError(t, err)
 	err = os.WriteFile(keyPath, []byte("invalid key content"), 0o600)
 	require.NoError(t, err)
-	_, err = common.NewTLS(caPath, certPath, keyPath, "localhost")
+	tls, err = common.NewTLS(caPath, certPath, keyPath, "localhost")
+	_, err = tls.TLSConfig().GetCertificate(nil)
 	require.Regexp(t, "could not load client key pair: tls.*", err.Error())
 }

From 09f7c66e50c59117176535d1ba0f7b9bf276f5ab Mon Sep 17 00:00:00 2001
From: Yujie Xia <xiayjchn@gmail.com>
Date: Fri, 15 Apr 2022 15:36:11 +0800
Subject: [PATCH 3/7] add test lightning_reload_cert

---
 br/pkg/lightning/lightning.go                 | 81 +++++++++++++++++++
 br/tests/lightning_reload_cert/config.toml    |  0
 .../data/test-schema-create.sql               |  1 +
 .../data/test.t-schema.sql                    |  1 +
 .../lightning_reload_cert/data/test.t.sql     |  1 +
 br/tests/lightning_reload_cert/run.sh         | 26 ++++++
 6 files changed, 110 insertions(+)
 create mode 100644 br/tests/lightning_reload_cert/config.toml
 create mode 100644 br/tests/lightning_reload_cert/data/test-schema-create.sql
 create mode 100644 br/tests/lightning_reload_cert/data/test.t-schema.sql
 create mode 100644 br/tests/lightning_reload_cert/data/test.t.sql
 create mode 100644 br/tests/lightning_reload_cert/run.sh

diff --git a/br/pkg/lightning/lightning.go b/br/pkg/lightning/lightning.go
index 4ad8622c8d5fb..d6c4e6eb54a6b 100644
--- a/br/pkg/lightning/lightning.go
+++ b/br/pkg/lightning/lightning.go
@@ -17,7 +17,11 @@ package lightning
 import (
 	"compress/gzip"
 	"context"
+	"crypto/ecdsa"
+	"crypto/rand"
+	"crypto/x509"
 	"encoding/json"
+	"encoding/pem"
 	"fmt"
 	"io"
 	"net"
@@ -317,6 +321,29 @@ func (l *Lightning) run(taskCtx context.Context, taskCfg *config.Config, o *opti
 		failpoint.Return(nil)
 	})
 
+	failpoint.Inject("SetCertExpiredSoon", func(val failpoint.Value) {
+		rootKeyPath := val.(string)
+		rootCaPath := taskCfg.Security.CAPath
+		keyPath := taskCfg.Security.KeyPath
+		certPath := taskCfg.Security.CertPath
+		certBytes, err := os.ReadFile(certPath)
+		if err != nil {
+			panic(err)
+		}
+		if err := os.WriteFile(certPath+".old", certBytes, 0o644); err != nil {
+			panic(err)
+		}
+		if err := updateCertExpiry(rootKeyPath, rootCaPath, keyPath, certPath, time.Second*10); err != nil {
+			panic(err)
+		}
+		// Must restore the original cert before the new cert is expired.
+		time.AfterFunc(time.Second*5, func() {
+			if err := os.Rename(certPath+".old", certPath); err != nil {
+				panic(err)
+			}
+		})
+	})
+
 	if err := taskCfg.TiDB.Security.RegisterMySQL(); err != nil {
 		return common.ErrInvalidTLSConfig.Wrap(err)
 	}
@@ -905,3 +932,57 @@ func SwitchMode(ctx context.Context, cfg *config.Config, tls *common.TLS, mode s
 		},
 	)
 }
+
+func updateCertExpiry(rootKeyPath, rootCaPath, keyPath, certPath string, expiry time.Duration) error {
+	rootKey, err := parsePrivateKey(rootKeyPath)
+	if err != nil {
+		return err
+	}
+	rootCaPem, err := os.ReadFile(rootCaPath)
+	if err != nil {
+		return err
+	}
+	rootCaDer, _ := pem.Decode(rootCaPem)
+	rootCa, err := x509.ParseCertificate(rootCaDer.Bytes)
+	if err != nil {
+		return err
+	}
+	key, err := parsePrivateKey(keyPath)
+	if err != nil {
+		return err
+	}
+	certPem, err := os.ReadFile(certPath)
+	if err != nil {
+		panic(err)
+	}
+	certDer, _ := pem.Decode(certPem)
+	cert, err := x509.ParseCertificate(certDer.Bytes)
+	if err != nil {
+		return err
+	}
+	cert.NotBefore = time.Now()
+	cert.NotAfter = time.Now().Add(expiry)
+	derBytes, err := x509.CreateCertificate(rand.Reader, cert, rootCa, &key.PublicKey, rootKey)
+	if err != nil {
+		return err
+	}
+	return os.WriteFile(certPath, pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes}), 0o644)
+}
+
+func parsePrivateKey(keyPath string) (*ecdsa.PrivateKey, error) {
+	keyPemBlock, err := os.ReadFile(keyPath)
+	if err != nil {
+		return nil, err
+	}
+	var keyDERBlock *pem.Block
+	for {
+		keyDERBlock, keyPemBlock = pem.Decode(keyPemBlock)
+		if keyDERBlock == nil {
+			return nil, errors.New("failed to find PEM block with type ending in \"PRIVATE KEY\"")
+		}
+		if keyDERBlock.Type == "PRIVATE KEY" || strings.HasSuffix(keyDERBlock.Type, " PRIVATE KEY") {
+			break
+		}
+	}
+	return x509.ParseECPrivateKey(keyDERBlock.Bytes)
+}
diff --git a/br/tests/lightning_reload_cert/config.toml b/br/tests/lightning_reload_cert/config.toml
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/br/tests/lightning_reload_cert/data/test-schema-create.sql b/br/tests/lightning_reload_cert/data/test-schema-create.sql
new file mode 100644
index 0000000000000..14379bd68472a
--- /dev/null
+++ b/br/tests/lightning_reload_cert/data/test-schema-create.sql
@@ -0,0 +1 @@
+CREATE DATABASE test;
diff --git a/br/tests/lightning_reload_cert/data/test.t-schema.sql b/br/tests/lightning_reload_cert/data/test.t-schema.sql
new file mode 100644
index 0000000000000..57a1b65732950
--- /dev/null
+++ b/br/tests/lightning_reload_cert/data/test.t-schema.sql
@@ -0,0 +1 @@
+CREATE TABLE t(a INT PRIMARY KEY, b int);
diff --git a/br/tests/lightning_reload_cert/data/test.t.sql b/br/tests/lightning_reload_cert/data/test.t.sql
new file mode 100644
index 0000000000000..30e06b42e169b
--- /dev/null
+++ b/br/tests/lightning_reload_cert/data/test.t.sql
@@ -0,0 +1 @@
+INSERT INTO t VALUES (1,1);
diff --git a/br/tests/lightning_reload_cert/run.sh b/br/tests/lightning_reload_cert/run.sh
new file mode 100644
index 0000000000000..b2f801b96c654
--- /dev/null
+++ b/br/tests/lightning_reload_cert/run.sh
@@ -0,0 +1,26 @@
+#!/bin/bash
+#
+# Copyright 2022 PingCAP, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -eux
+
+cp "$TEST_DIR/certs/lightning.pem" "$TEST_DIR/certs/lightning-valid.pem"
+trap 'mv "$TEST_DIR/certs/lightning-valid.pem" "$TEST_DIR/certs/lightning.pem"' EXIT
+
+# shellcheck disable=SC2089
+export GO_FAILPOINTS="github.com/pingcap/tidb/br/pkg/lightning/SetCertExpiredSoon=return(\"$TEST_DIR/certs/ca.key\")"
+export GO_FAILPOINTS="${GO_FAILPOINTS};github.com/pingcap/tidb/br/pkg/lightning/restore/SlowDownWriteRows=sleep(15000)"
+
+run_lightning --backend='local'

From 2e3739106aa0ab7ef048c2d99e90b25772a47964 Mon Sep 17 00:00:00 2001
From: Yujie Xia <xiayjchn@gmail.com>
Date: Sun, 17 Apr 2022 15:45:20 +0800
Subject: [PATCH 4/7] more strict perm

---
 br/pkg/lightning/lightning.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/br/pkg/lightning/lightning.go b/br/pkg/lightning/lightning.go
index d6c4e6eb54a6b..871928a740aac 100644
--- a/br/pkg/lightning/lightning.go
+++ b/br/pkg/lightning/lightning.go
@@ -330,7 +330,7 @@ func (l *Lightning) run(taskCtx context.Context, taskCfg *config.Config, o *opti
 		if err != nil {
 			panic(err)
 		}
-		if err := os.WriteFile(certPath+".old", certBytes, 0o644); err != nil {
+		if err := os.WriteFile(certPath+".old", certBytes, 0o600); err != nil {
 			panic(err)
 		}
 		if err := updateCertExpiry(rootKeyPath, rootCaPath, keyPath, certPath, time.Second*10); err != nil {
@@ -966,7 +966,7 @@ func updateCertExpiry(rootKeyPath, rootCaPath, keyPath, certPath string, expiry
 	if err != nil {
 		return err
 	}
-	return os.WriteFile(certPath, pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes}), 0o644)
+	return os.WriteFile(certPath, pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes}), 0o600)
 }
 
 func parsePrivateKey(keyPath string) (*ecdsa.PrivateKey, error) {

From 82272b2d0d5a98842ce343f7cb60f4c5b4358ef8 Mon Sep 17 00:00:00 2001
From: Yujie Xia <xiayjchn@gmail.com>
Date: Tue, 19 Apr 2022 16:53:49 +0800
Subject: [PATCH 5/7] add negative test

---
 br/pkg/lightning/lightning.go              | 13 -----------
 br/tests/lightning_reload_cert/config.toml |  2 ++
 br/tests/lightning_reload_cert/run.sh      | 27 ++++++++++++++++++++++
 3 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/br/pkg/lightning/lightning.go b/br/pkg/lightning/lightning.go
index 871928a740aac..b857e93f032b8 100644
--- a/br/pkg/lightning/lightning.go
+++ b/br/pkg/lightning/lightning.go
@@ -326,22 +326,9 @@ func (l *Lightning) run(taskCtx context.Context, taskCfg *config.Config, o *opti
 		rootCaPath := taskCfg.Security.CAPath
 		keyPath := taskCfg.Security.KeyPath
 		certPath := taskCfg.Security.CertPath
-		certBytes, err := os.ReadFile(certPath)
-		if err != nil {
-			panic(err)
-		}
-		if err := os.WriteFile(certPath+".old", certBytes, 0o600); err != nil {
-			panic(err)
-		}
 		if err := updateCertExpiry(rootKeyPath, rootCaPath, keyPath, certPath, time.Second*10); err != nil {
 			panic(err)
 		}
-		// Must restore the original cert before the new cert is expired.
-		time.AfterFunc(time.Second*5, func() {
-			if err := os.Rename(certPath+".old", certPath); err != nil {
-				panic(err)
-			}
-		})
 	})
 
 	if err := taskCfg.TiDB.Security.RegisterMySQL(); err != nil {
diff --git a/br/tests/lightning_reload_cert/config.toml b/br/tests/lightning_reload_cert/config.toml
index e69de29bb2d1d..850796f5cdb8a 100644
--- a/br/tests/lightning_reload_cert/config.toml
+++ b/br/tests/lightning_reload_cert/config.toml
@@ -0,0 +1,2 @@
+[checkpoint]
+enable = false
diff --git a/br/tests/lightning_reload_cert/run.sh b/br/tests/lightning_reload_cert/run.sh
index b2f801b96c654..a08a0b6c0dc9d 100644
--- a/br/tests/lightning_reload_cert/run.sh
+++ b/br/tests/lightning_reload_cert/run.sh
@@ -23,4 +23,31 @@ trap 'mv "$TEST_DIR/certs/lightning-valid.pem" "$TEST_DIR/certs/lightning.pem"'
 export GO_FAILPOINTS="github.com/pingcap/tidb/br/pkg/lightning/SetCertExpiredSoon=return(\"$TEST_DIR/certs/ca.key\")"
 export GO_FAILPOINTS="${GO_FAILPOINTS};github.com/pingcap/tidb/br/pkg/lightning/restore/SlowDownWriteRows=sleep(15000)"
 
+# 1. After 10s, the certificate will be expired and import should report connection error.
+run_lightning --backend='local' &
+shpid="$!"
+sleep 15
+ok=0
+for _ in {0..60}; do
+  if grep -Fq "connection closed before server preface received" "$TEST_DIR"/lightning.log; then
+    ok=1
+    break
+  fi
+  sleep 1
+done
+# Lightning process is wrapped by a shell process, use pstree to extract it out.
+pid=$(pstree -pT "$shpid" | grep -Eo "tidb-lightning\.\([0-9]*\)" | grep -Eo "[0-9]*")
+if [ -n "$pid" ]; then
+  kill -9 "$pid" &>/dev/null || true
+fi
+if [ "$ok" = "0" ]; then
+  echo "lightning should report connection error due to certificate expired, but not error is reported"
+  exit 1
+fi
+# Do some cleanup.
+cp "$TEST_DIR/certs/lightning-valid.pem" "$TEST_DIR/certs/lightning.pem"
+rm -rf "$TEST_DIR/lightning_reload_cert.sorted" "$TEST_DIR"/lightning.log
+
+# 2. Replace the certificate with a valid certificate before it is expired. Lightning should import successfully.
+sleep 10 && cp "$TEST_DIR/certs/lightning-valid.pem" "$TEST_DIR/certs/lightning.pem" &
 run_lightning --backend='local'

From 6a3ec5c15dab53bcf4331caf21422243f7cd9001 Mon Sep 17 00:00:00 2001
From: Yujie Xia <xiayjchn@gmail.com>
Date: Tue, 19 Apr 2022 18:11:07 +0800
Subject: [PATCH 6/7] pstree: invalid option

---
 br/tests/lightning_reload_cert/run.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/br/tests/lightning_reload_cert/run.sh b/br/tests/lightning_reload_cert/run.sh
index a08a0b6c0dc9d..265cf1bcb55eb 100644
--- a/br/tests/lightning_reload_cert/run.sh
+++ b/br/tests/lightning_reload_cert/run.sh
@@ -36,7 +36,7 @@ for _ in {0..60}; do
   sleep 1
 done
 # Lightning process is wrapped by a shell process, use pstree to extract it out.
-pid=$(pstree -pT "$shpid" | grep -Eo "tidb-lightning\.\([0-9]*\)" | grep -Eo "[0-9]*")
+pid=$(pstree -p "$shpid" | grep -Eo "tidb-lightning\.\([0-9]*\)" | grep -Eo "[0-9]*")
 if [ -n "$pid" ]; then
   kill -9 "$pid" &>/dev/null || true
 fi

From 789f5897cbae983d534ad0bec5f898e09ab67d0a Mon Sep 17 00:00:00 2001
From: Yujie Xia <xiayjchn@gmail.com>
Date: Wed, 20 Apr 2022 10:31:45 +0800
Subject: [PATCH 7/7] Update br/tests/lightning_reload_cert/run.sh

---
 br/tests/lightning_reload_cert/run.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/br/tests/lightning_reload_cert/run.sh b/br/tests/lightning_reload_cert/run.sh
index 265cf1bcb55eb..e06ef8d7fbf51 100644
--- a/br/tests/lightning_reload_cert/run.sh
+++ b/br/tests/lightning_reload_cert/run.sh
@@ -41,7 +41,7 @@ if [ -n "$pid" ]; then
   kill -9 "$pid" &>/dev/null || true
 fi
 if [ "$ok" = "0" ]; then
-  echo "lightning should report connection error due to certificate expired, but not error is reported"
+  echo "lightning should report connection error due to certificate expired, but no error is reported"
   exit 1
 fi
 # Do some cleanup.