ast: implement SensitiveStmtNode for BRIEStmt (pingcap#842)

Author: kennytm

ast/misc.go

@@ -16,6 +16,7 @@ package ast
 import (
 	"bytes"
 	"fmt"
+	"net/url"
 	"strconv"
 	"strings"
 
@@ -2436,6 +2437,38 @@ func (n *BRIEStmt) Restore(ctx *format.RestoreCtx) error {
 	return nil
 }
 
+// SecureText implements SensitiveStmtNode
+func (n *BRIEStmt) SecureText() string {
+	// FIXME: this solution is not scalable, and duplicates some logic from BR.
+	redactedStorage := n.Storage
+	u, err := url.Parse(n.Storage)
+	if err == nil {
+		if u.Scheme == "s3" {
+			query := u.Query()
+			for key := range query {
+				switch strings.ToLower(strings.ReplaceAll(key, "_", "-")) {
+				case "access-key", "secret-access-key":
+					query[key] = []string{"xxxxxx"}
+				}
+			}
+			u.RawQuery = query.Encode()
+			redactedStorage = u.String()
+		}
+	}
+
+	redactedStmt := &BRIEStmt{
+		Kind:    n.Kind,
+		Schemas: n.Schemas,
+		Tables:  n.Tables,
+		Storage: redactedStorage,
+		Options: n.Options,
+	}
+
+	var sb strings.Builder
+	_ = redactedStmt.Restore(format.NewRestoreCtx(format.DefaultRestoreFlags, &sb))
+	return sb.String()
+}
+
 // Ident is the table identifier composed of schema name and table name.
 type Ident struct {
 	Schema model.CIStr

ast/misc_test.go

@@ -16,6 +16,7 @@ package ast_test
 import (
 	. "github.com/pingcap/check"
 	"github.com/pingcap/parser"
+	"github.com/pingcap/parser/ast"
 	. "github.com/pingcap/parser/ast"
 	"github.com/pingcap/parser/auth"
 	"github.com/pingcap/parser/mysql"
@@ -277,3 +278,40 @@ func (ts *testMiscSuite) TestChangeStmtRestore(c *C) {
 	}
 	RunNodeRestoreTest(c, testCases, "%s", extractNodeFunc)
 }
+
+func (ts *testMiscSuite) TestBRIESecureText(c *C) {
+	testCases := []struct {
+		input   string
+		secured string
+	}{
+		{
+			input:   "restore database * from 'local:///tmp/br01' snapshot = 23333",
+			secured: `^\QRESTORE DATABASE * FROM 'local:///tmp/br01' SNAPSHOT = 23333\E$`,
+		},
+		{
+			input:   "backup database * to 's3://bucket/prefix?region=us-west-2'",
+			secured: `^\QBACKUP DATABASE * TO 's3://bucket/prefix?region=us-west-2'\E$`,
+		},
+		{
+			// we need to use regexp to match to avoid the random ordering since a map was used.
+			// unfortunately Go's regexp doesn't support lookahead assertion, so the test case below
+			// has false positives.
+			input:   "backup database * to 's3://bucket/prefix?access-key=abcdefghi&secret-access-key=123&force-path-style=true'",
+			secured: `^\QBACKUP DATABASE * TO 's3://bucket/prefix?\E((access-key=xxxxxx|force-path-style=true|secret-access-key=xxxxxx)(&|'$)){3}`,
+		},
+		{
+			input:   "backup database * to 'gcs://bucket/prefix?access-key=irrelevant&credentials-file=/home/user/secrets.txt'",
+			secured: `^\QBACKUP DATABASE * TO 'gcs://bucket/prefix?\E((access-key=irrelevant|credentials-file=/home/user/secrets\.txt)(&|'$)){2}`,
+		},
+	}
+
+	parser := parser.New()
+	for _, tc := range testCases {
+		comment := Commentf("input = %s", tc.input)
+		node, err := parser.ParseOneStmt(tc.input, "", "")
+		c.Assert(err, IsNil, comment)
+		n, ok := node.(ast.SensitiveStmtNode)
+		c.Assert(ok, IsTrue, comment)
+		c.Assert(n.SecureText(), Matches, tc.secured, comment)
+	}
+}