From da76e349539c15cd7a6a98f83a8597fad636b9cf Mon Sep 17 00:00:00 2001
From: Morgan Tocker <tocker@gmail.com>
Date: Wed, 27 Oct 2021 00:38:49 -0600
Subject: [PATCH] *: add support for --initialize-secure and
 --initialize-insecure (#28487)

---
 config/config.go     |  7 ++++---
 session/bootstrap.go | 17 ++++++++++++++---
 tidb-server/main.go  | 29 +++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/config/config.go b/config/config.go
index d8be4a87bf6a3..6167b33e11593 100644
--- a/config/config.go
+++ b/config/config.go
@@ -374,9 +374,10 @@ type Security struct {
 	// EnableSEM prevents SUPER users from having full access.
 	EnableSEM bool `toml:"enable-sem" json:"enable-sem"`
 	// Allow automatic TLS certificate generation
-	AutoTLS       bool   `toml:"auto-tls" json:"auto-tls"`
-	MinTLSVersion string `toml:"tls-version" json:"tls-version"`
-	RSAKeySize    int    `toml:"rsa-key-size" json:"rsa-key-size"`
+	AutoTLS         bool   `toml:"auto-tls" json:"auto-tls"`
+	MinTLSVersion   string `toml:"tls-version" json:"tls-version"`
+	RSAKeySize      int    `toml:"rsa-key-size" json:"rsa-key-size"`
+	SecureBootstrap bool   `toml:"secure-bootstrap" json:"secure-bootstrap"`
 }
 
 // The ErrConfigValidationFailed error is used so that external callers can do a type assertion
diff --git a/session/bootstrap.go b/session/bootstrap.go
index 9b8156fb76627..6f16db51a0ad6 100644
--- a/session/bootstrap.go
+++ b/session/bootstrap.go
@@ -23,6 +23,7 @@ import (
 	"encoding/hex"
 	"flag"
 	"fmt"
+	osuser "os/user"
 	"runtime/debug"
 	"strconv"
 	"strings"
@@ -1700,10 +1701,20 @@ func doDDLWorks(s Session) {
 // TODO: sanitize.
 func doDMLWorks(s Session) {
 	mustExecute(s, "BEGIN")
-
-	// Insert a default user with empty password.
-	mustExecute(s, `INSERT HIGH_PRIORITY INTO mysql.user VALUES
+	if config.GetGlobalConfig().Security.SecureBootstrap {
+		// If secure bootstrap is enabled, we create a root@localhost account which can login with auth_socket.
+		// i.e. mysql -S /tmp/tidb.sock -uroot
+		// The auth_socket plugin will validate that the user matches $USER.
+		u, err := osuser.Current()
+		if err != nil {
+			logutil.BgLogger().Fatal("failed to read current user. unable to secure bootstrap.", zap.Error(err))
+		}
+		mustExecute(s, `INSERT HIGH_PRIORITY INTO mysql.user VALUES
+		("localhost", "root", %?, "auth_socket", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "N", "Y", "Y", "Y", "Y", "Y", "Y", "Y")`, u.Username)
+	} else {
+		mustExecute(s, `INSERT HIGH_PRIORITY INTO mysql.user VALUES
 		("%", "root", "", "mysql_native_password", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "Y", "N", "Y", "Y", "Y", "Y", "Y", "Y", "Y")`)
+	}
 
 	// Init global system variables table.
 	values := make([]string, 0, len(variable.GetSysVars()))
diff --git a/tidb-server/main.go b/tidb-server/main.go
index db40c13d43e39..3191e48035939 100644
--- a/tidb-server/main.go
+++ b/tidb-server/main.go
@@ -110,6 +110,9 @@ const (
 	nmProxyProtocolNetworks      = "proxy-protocol-networks"
 	nmProxyProtocolHeaderTimeout = "proxy-protocol-header-timeout"
 	nmAffinityCPU                = "affinity-cpus"
+
+	nmInitializeSecure   = "initialize-secure"
+	nmInitializeInsecure = "initialize-insecure"
 )
 
 var (
@@ -152,6 +155,10 @@ var (
 	// PROXY Protocol
 	proxyProtocolNetworks      = flag.String(nmProxyProtocolNetworks, "", "proxy protocol networks allowed IP or *, empty mean disable proxy protocol support")
 	proxyProtocolHeaderTimeout = flag.Uint(nmProxyProtocolHeaderTimeout, 5, "proxy protocol header read timeout, unit is second.")
+
+	// Security
+	initializeSecure   = flagBoolean(nmInitializeSecure, false, "bootstrap tidb-server in secure mode")
+	initializeInsecure = flagBoolean(nmInitializeInsecure, true, "bootstrap tidb-server in insecure mode")
 )
 
 func main() {
@@ -505,6 +512,28 @@ func overrideConfig(cfg *config.Config) {
 	if actualFlags[nmProxyProtocolHeaderTimeout] {
 		cfg.ProxyProtocol.HeaderTimeout = *proxyProtocolHeaderTimeout
 	}
+
+	// Sanity check: can't specify both options
+	if actualFlags[nmInitializeSecure] && actualFlags[nmInitializeInsecure] {
+		err = fmt.Errorf("the options --initialize-insecure and --initialize-secure are mutually exclusive")
+		terror.MustNil(err)
+	}
+	// The option --initialize-secure=true ensures that a secure bootstrap is used.
+	if actualFlags[nmInitializeSecure] {
+		cfg.Security.SecureBootstrap = *initializeSecure
+	}
+	// The option --initialize-insecure=true/false was used.
+	// Store the inverted value of this to the secure bootstrap cfg item
+	if actualFlags[nmInitializeInsecure] {
+		cfg.Security.SecureBootstrap = !*initializeInsecure
+	}
+	// Secure bootstrap initializes with Socket authentication
+	// which is not supported on windows. Only the insecure bootstrap
+	// method is supported.
+	if runtime.GOOS == "windows" && cfg.Security.SecureBootstrap {
+		err = fmt.Errorf("the option --initialize-secure is not supported on Windows")
+		terror.MustNil(err)
+	}
 }
 
 func setGlobalVars() {