Go · 2543 bytes Raw Blame History
1 // SPDX-License-Identifier: AGPL-3.0-or-later
2
3 // Package scrub masks configured secret values from runner log output.
4 package scrub
5
6 import (
7 "sort"
8 "strings"
9 )
10
11 const Mask = "***"
12
13 type Scrubber struct {
14 values []string
15 replacer *strings.Replacer
16 tail string
17 replacements uint64
18 }
19
20 func New(values []string) *Scrubber {
21 values = normalize(values)
22 if len(values) == 0 {
23 return &Scrubber{}
24 }
25 pairs := make([]string, 0, len(values)*2)
26 for _, v := range values {
27 pairs = append(pairs, v, Mask)
28 }
29 return &Scrubber{values: values, replacer: strings.NewReplacer(pairs...)}
30 }
31
32 func (s *Scrubber) Scrub(chunk []byte) []byte {
33 if s == nil || s.replacer == nil {
34 return append([]byte(nil), chunk...)
35 }
36 combined := s.tail + string(chunk)
37 keep := s.pendingSuffixLen(combined)
38 if keep == len(combined) {
39 s.tail = combined
40 return nil
41 }
42 emit := combined[:len(combined)-keep]
43 s.tail = combined[len(combined)-keep:]
44 s.replacements += countReplacements(emit, s.values)
45 return []byte(s.replacer.Replace(emit))
46 }
47
48 func (s *Scrubber) Flush() []byte {
49 if s == nil || s.tail == "" {
50 return nil
51 }
52 tail := s.tail
53 s.tail = ""
54 if s.replacer == nil {
55 return []byte(tail)
56 }
57 s.replacements += countReplacements(tail, s.values)
58 return []byte(s.replacer.Replace(tail))
59 }
60
61 func (s *Scrubber) Replacements() uint64 {
62 if s == nil {
63 return 0
64 }
65 return s.replacements
66 }
67
68 func normalize(values []string) []string {
69 seen := map[string]struct{}{}
70 out := make([]string, 0, len(values))
71 for _, v := range values {
72 if v == "" {
73 continue
74 }
75 if _, ok := seen[v]; ok {
76 continue
77 }
78 seen[v] = struct{}{}
79 out = append(out, v)
80 }
81 sort.Slice(out, func(i, j int) bool {
82 return len(out[i]) > len(out[j])
83 })
84 return out
85 }
86
87 func countReplacements(input string, values []string) uint64 {
88 var count uint64
89 rest := input
90 for rest != "" {
91 bestAt := -1
92 best := ""
93 for _, value := range values {
94 at := strings.Index(rest, value)
95 if at < 0 {
96 continue
97 }
98 if bestAt == -1 || at < bestAt || (at == bestAt && len(value) > len(best)) {
99 bestAt = at
100 best = value
101 }
102 }
103 if bestAt == -1 {
104 return count
105 }
106 count++
107 rest = rest[bestAt+len(best):]
108 }
109 return count
110 }
111
112 func (s *Scrubber) pendingSuffixLen(combined string) int {
113 keep := 0
114 for _, secret := range s.values {
115 max := len(secret) - 1
116 if max > len(combined) {
117 max = len(combined)
118 }
119 for n := max; n > keep; n-- {
120 if strings.HasSuffix(combined, secret[:n]) {
121 keep = n
122 break
123 }
124 }
125 }
126 return keep
127 }
128