playground: auto tiproxy sign certs (#2372)

* playground: auto tiproxy sign certs Signed-off-by: xhe <[email protected]> * fix lint Signed-off-by: xhe <[email protected]> * fix lint Signed-off-by: xhe <[email protected]> --------- Signed-off-by: xhe <[email protected]>
pingcap · Mar 25, 2024 · 1ab96db · 1ab96db
1 parent f5e6dbf
commit 1ab96db
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 7 deletions.
diff --git a/components/playground/instance/tidb.go b/components/playground/instance/tidb.go
@@ -29,12 +29,13 @@ type TiDBInstance struct {
 	instance
 	pds []*PDInstance
 	Process
-	enableBinlog bool
-	isDisaggMode bool
+	tiproxyCertDir string
+	enableBinlog   bool
+	isDisaggMode   bool
 }
 
 // NewTiDBInstance return a TiDBInstance
-func NewTiDBInstance(binPath string, dir, host, configPath string, id, port int, pds []*PDInstance, enableBinlog bool, isDisaggMode bool) *TiDBInstance {
+func NewTiDBInstance(binPath string, dir, host, configPath string, id, port int, pds []*PDInstance, tiproxyCertDir string, enableBinlog bool, isDisaggMode bool) *TiDBInstance {
 	if port <= 0 {
 		port = 4000
 	}
@@ -48,9 +49,10 @@ func NewTiDBInstance(binPath string, dir, host, configPath string, id, port int,
 			StatusPort: utils.MustGetFreePort("0.0.0.0", 10080),
 			ConfigPath: configPath,
 		},
-		pds:          pds,
-		enableBinlog: enableBinlog,
-		isDisaggMode: isDisaggMode,
+		tiproxyCertDir: tiproxyCertDir,
+		pds:            pds,
+		enableBinlog:   enableBinlog,
+		isDisaggMode:   isDisaggMode,
 	}
 }
 

diff --git a/components/playground/instance/tidb_config.go b/components/playground/instance/tidb_config.go
@@ -13,6 +13,11 @@
 
 package instance
 
+import (
+	"os"
+	"path/filepath"
+)
+
 func (inst *TiDBInstance) getConfig() map[string]any {
 	config := make(map[string]any)
 	config["security.auto-tls"] = true
@@ -22,5 +27,14 @@ func (inst *TiDBInstance) getConfig() map[string]any {
 		config["disaggregated-tiflash"] = true
 	}
 
+	tiproxyCrtPath := filepath.Join(inst.tiproxyCertDir, "tiproxy.crt")
+	tiproxyKeyPath := filepath.Join(inst.tiproxyCertDir, "tiproxy.key")
+	_, err1 := os.Stat(tiproxyCrtPath)
+	_, err2 := os.Stat(tiproxyKeyPath)
+	if err1 == nil && err2 == nil {
+		config["security.session-token-signing-cert"] = tiproxyCrtPath
+		config["security.session-token-signing-key"] = tiproxyKeyPath
+	}
+
 	return config
 }
diff --git a/components/playground/instance/tiproxy.go b/components/playground/instance/tiproxy.go
@@ -15,13 +15,15 @@ package instance
 
 import (
 	"context"
+	"encoding/pem"
 	"fmt"
 	"os"
 	"path/filepath"
 	"strings"
 
 	"github.com/BurntSushi/toml"
 	"github.com/pingcap/tiup/pkg/cluster/spec"
+	"github.com/pingcap/tiup/pkg/crypto"
 	tiupexec "github.com/pingcap/tiup/pkg/exec"
 	"github.com/pingcap/tiup/pkg/utils"
 )
@@ -35,6 +37,37 @@ type TiProxy struct {
 
 var _ Instance = &TiProxy{}
 
+// GenTiProxySessionCerts will create a self-signed certs for TiProxy session migration. NOTE that this cert is directly used by TiDB.
+func GenTiProxySessionCerts(dir string) error {
+	if _, err := os.Stat(filepath.Join(dir, "tiproxy.crt")); err == nil {
+		return nil
+	}
+
+	ca, err := crypto.NewCA("tiproxy")
+	if err != nil {
+		return err
+	}
+	privKey, err := crypto.NewKeyPair(crypto.KeyTypeRSA, crypto.KeySchemeRSASSAPSSSHA256)
+	if err != nil {
+		return err
+	}
+	csr, err := privKey.CSR("tiproxy", "tiproxy", nil, nil)
+	if err != nil {
+		return err
+	}
+	cert, err := ca.Sign(csr)
+	if err != nil {
+		return err
+	}
+	if err := utils.SaveFileWithBackup(filepath.Join(dir, "tiproxy.key"), privKey.Pem(), ""); err != nil {
+		return err
+	}
+	return utils.SaveFileWithBackup(filepath.Join(dir, "tiproxy.crt"), pem.EncodeToMemory(&pem.Block{
+		Type:  "CERTIFICATE",
+		Bytes: cert,
+	}), "")
+}
+
 // NewTiProxy create a TiProxy instance.
 func NewTiProxy(binPath string, dir, host, configPath string, id int, port int, pds []*PDInstance) *TiProxy {
 	if port <= 0 {

diff --git a/components/playground/playground.go b/components/playground/playground.go
@@ -740,7 +740,7 @@ func (p *Playground) addInstance(componentID string, pdRole instance.PDRole, tif
 			p.rms = append(p.rms, inst)
 		}
 	case spec.ComponentTiDB:
-		inst := instance.NewTiDBInstance(cfg.BinPath, dir, host, cfg.ConfigPath, id, cfg.Port, p.pds, p.enableBinlog(), p.bootOptions.Mode == "tidb-disagg")
+		inst := instance.NewTiDBInstance(cfg.BinPath, dir, host, cfg.ConfigPath, id, cfg.Port, p.pds, dataDir, p.enableBinlog(), p.bootOptions.Mode == "tidb-disagg")
 		ins = inst
 		p.tidbs = append(p.tidbs, inst)
 	case spec.ComponentTiKV:
@@ -752,6 +752,9 @@ func (p *Playground) addInstance(componentID string, pdRole instance.PDRole, tif
 		ins = inst
 		p.tiflashs = append(p.tiflashs, inst)
 	case spec.ComponentTiProxy:
+		if err := instance.GenTiProxySessionCerts(dataDir); err != nil {
+			return nil, err
+		}
 		inst := instance.NewTiProxy(cfg.BinPath, dir, host, cfg.ConfigPath, id, cfg.Port, p.pds)
 		ins = inst
 		p.tiproxys = append(p.tiproxys, inst)