Skip to content

Commit d6de964

Browse files
xheboxxhebox
authored
cmd: add new subcommand for redacting logs (#51759)
close #51758
1 parent 1718d7b commit d6de964

File tree

5 files changed

+164
-0
lines changed

5 files changed

+164
-0
lines changed

‎cmd/tidb-server/BUILD.bazel

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ go_library(
4848
"//pkg/util/memory",
4949
"//pkg/util/metricsutil",
5050
"//pkg/util/printer",
51+
"//pkg/util/redact",
5152
"//pkg/util/sem",
5253
"//pkg/util/signal",
5354
"//pkg/util/stmtsummary/v2:stmtsummary",

‎cmd/tidb-server/main.go

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ import (
7272
"github.com/pingcap/tidb/pkg/util/memory"
7373
"github.com/pingcap/tidb/pkg/util/metricsutil"
7474
"github.com/pingcap/tidb/pkg/util/printer"
75+
"github.com/pingcap/tidb/pkg/util/redact"
7576
"github.com/pingcap/tidb/pkg/util/sem"
7677
"github.com/pingcap/tidb/pkg/util/signal"
7778
stmtsummaryv2 "github.com/pingcap/tidb/pkg/util/stmtsummary/v2"
@@ -122,6 +123,8 @@ const (
122123
nmRepairList = "repair-list"
123124
nmTempDir = "temp-dir"
124125

126+
nmRedact = "redact"
127+
125128
nmProxyProtocolNetworks = "proxy-protocol-networks"
126129
nmProxyProtocolHeaderTimeout = "proxy-protocol-header-timeout"
127130
nmProxyProtocolFallbackable = "proxy-protocol-fallbackable"
@@ -173,6 +176,9 @@ var (
173176
metricsAddr *string
174177
metricsInterval *uint
175178

179+
// subcommand collect-log
180+
redactFlag *bool
181+
176182
// PROXY Protocol
177183
proxyProtocolNetworks *string
178184
proxyProtocolHeaderTimeout *uint
@@ -227,6 +233,9 @@ func initFlagSet() *flag.FlagSet {
227233
metricsAddr = fset.String(nmMetricsAddr, "", "prometheus pushgateway address, leaves it empty will disable prometheus push.")
228234
metricsInterval = fset.Uint(nmMetricsInterval, 15, "prometheus client push interval in second, set \"0\" to disable prometheus push.")
229235

236+
// subcommand collect-log
237+
redactFlag = flagBoolean(fset, nmRedact, false, "remove sensitive words from marked tidb logs, if `./tidb-server --redact=xxx collect-log <input> <output>` subcommand is used")
238+
230239
// PROXY Protocol
231240
proxyProtocolNetworks = fset.String(nmProxyProtocolNetworks, "", "proxy protocol networks allowed IP or *, empty mean disable proxy protocol support")
232241
proxyProtocolHeaderTimeout = fset.Uint(nmProxyProtocolHeaderTimeout, 5, "proxy protocol header read timeout, unit is second. (Deprecated: as proxy protocol using lazy mode, header read timeout no longer used)")
@@ -253,6 +262,16 @@ func initFlagSet() *flag.FlagSet {
253262

254263
func main() {
255264
fset := initFlagSet()
265+
if args := fset.Args(); len(args) != 0 {
266+
if args[0] == "collect-logs" && len(args) > 1 {
267+
output := "-"
268+
if len(args) > 2 {
269+
output = args[2]
270+
}
271+
terror.MustNil(redact.DeRedactFile(*redactFlag, args[1], output))
272+
return
273+
}
274+
}
256275
config.InitializeConfig(*configPath, *configCheck, *configStrict, overrideConfig, fset)
257276
if *version {
258277
setVersions()

‎pkg/util/redact/BUILD.bazel

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ go_library(
55
srcs = ["redact.go"],
66
importpath = "github.com/pingcap/tidb/pkg/util/redact",
77
visibility = ["//visibility:public"],
8+
deps = ["@com_github_pingcap_errors//:errors"],
89
)
910

1011
go_test(

‎pkg/util/redact/redact.go

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,15 @@
1515
package redact
1616

1717
import (
18+
"bufio"
19+
"bytes"
1820
"fmt"
21+
"io"
22+
"os"
23+
"path/filepath"
1924
"strings"
25+
26+
"github.com/pingcap/errors"
2027
)
2128

2229
var (
@@ -60,3 +67,108 @@ func (s redactStringer) String() string {
6067
func Stringer(mode string, input fmt.Stringer) redactStringer {
6168
return redactStringer{mode, input}
6269
}
70+
71+
// DeRedactFile will deredact the input file, either removing marked contents, or remove the marker. It works line by line.
72+
func DeRedactFile(remove bool, input string, output string) error {
73+
ifile, err := os.Open(filepath.Clean(input))
74+
if err != nil {
75+
return errors.WithStack(err)
76+
}
77+
defer ifile.Close()
78+
79+
var ofile io.Writer
80+
if output == "-" {
81+
ofile = os.Stdout
82+
} else {
83+
//nolint: gosec
84+
file, err := os.OpenFile(filepath.Clean(output), os.O_TRUNC|os.O_CREATE|os.O_WRONLY, 0644)
85+
if err != nil {
86+
return errors.WithStack(err)
87+
}
88+
defer file.Close()
89+
ofile = file
90+
}
91+
92+
return DeRedact(remove, ifile, ofile, "\n")
93+
}
94+
95+
// DeRedact is similar to DeRedactFile, but act on reader/writer, it works line by line.
96+
func DeRedact(remove bool, input io.Reader, output io.Writer, sep string) error {
97+
sc := bufio.NewScanner(input)
98+
out := bufio.NewWriter(output)
99+
defer out.Flush()
100+
buf := bytes.NewBuffer(nil)
101+
s := bufio.NewReader(nil)
102+
103+
for sc.Scan() {
104+
s.Reset(strings.NewReader(sc.Text()))
105+
start := false
106+
for {
107+
ch, _, err := s.ReadRune()
108+
if err == io.EOF {
109+
break
110+
}
111+
if err != nil {
112+
return errors.WithStack(err)
113+
}
114+
if ch == '‹' {
115+
if start {
116+
// must be '<'
117+
pch, _, err := s.ReadRune()
118+
if err != nil {
119+
return errors.WithStack(err)
120+
}
121+
if pch == ch {
122+
_, _ = buf.WriteRune(ch)
123+
} else {
124+
_, _ = buf.WriteRune(ch)
125+
_, _ = buf.WriteRune(pch)
126+
}
127+
} else {
128+
start = true
129+
buf.Reset()
130+
}
131+
} else if ch == '›' {
132+
if start {
133+
// peek the next
134+
pch, _, err := s.ReadRune()
135+
if err != nil && err != io.EOF {
136+
return errors.WithStack(err)
137+
}
138+
if pch == ch {
139+
_, _ = buf.WriteRune(ch)
140+
} else {
141+
start = false
142+
if err != io.EOF {
143+
// unpeek it
144+
if err := s.UnreadRune(); err != nil {
145+
return errors.WithStack(err)
146+
}
147+
}
148+
if remove {
149+
_ = out.WriteByte('?')
150+
} else {
151+
_, err = io.Copy(out, buf)
152+
if err != nil {
153+
return errors.WithStack(err)
154+
}
155+
}
156+
}
157+
} else {
158+
_, _ = out.WriteRune(ch)
159+
}
160+
} else if start {
161+
_, _ = buf.WriteRune(ch)
162+
} else {
163+
_, _ = out.WriteRune(ch)
164+
}
165+
}
166+
if start {
167+
_, _ = out.WriteRune('‹')
168+
_, _ = out.WriteString(buf.String())
169+
}
170+
_, _ = out.WriteString(sep)
171+
}
172+
173+
return nil
174+
}

‎pkg/util/redact/redact_test.go

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
package redact
1616

1717
import (
18+
"bytes"
19+
"strings"
1820
"testing"
1921

2022
"github.com/stretchr/testify/require"
@@ -44,3 +46,32 @@ func TestRedact(t *testing.T) {
4446
require.Equal(t, c.output, Stringer(c.mode, &testStringer{c.input}).String())
4547
}
4648
}
49+
50+
func TestDeRedact(t *testing.T) {
51+
for _, c := range []struct {
52+
remove bool
53+
input string
54+
output string
55+
}{
56+
{true, "‹fxcv›ggg", "?ggg"},
57+
{false, "‹fxcv›ggg", "fxcvggg"},
58+
{true, "fxcv", "fxcv"},
59+
{false, "fxcv", "fxcv"},
60+
{true, "‹fxcv›ggg‹fxcv›eee", "?ggg?eee"},
61+
{false, "‹fxcv›ggg‹fxcv›eee", "fxcvgggfxcveee"},
62+
{true, "‹›", "?"},
63+
{false, "‹›", ""},
64+
{true, "gg‹ee", "gg‹ee"},
65+
{false, "gg‹ee", "gg‹ee"},
66+
{true, "gg›ee", "gg›ee"},
67+
{false, "gg›ee", "gg›ee"},
68+
{true, "gg‹ee‹ee", "gg‹ee‹ee"},
69+
{false, "gg‹ee‹gg", "gg‹ee‹gg"},
70+
{true, "gg›ee›gg", "gg›ee›gg"},
71+
{false, "gg›ee›ee", "gg›ee›ee"},
72+
} {
73+
w := bytes.NewBuffer(nil)
74+
require.NoError(t, DeRedact(c.remove, strings.NewReader(c.input), w, ""))
75+
require.Equal(t, c.output, w.String())
76+
}
77+
}

0 commit comments

Comments
 (0)