github.com/nilium/gitlab-runner@v12.5.0+incompatible/helpers/trace/buffer.go (about) 1 package trace 2 3 import ( 4 "bufio" 5 "bytes" 6 "fmt" 7 "io" 8 "io/ioutil" 9 "os" 10 "sync" 11 12 "github.com/markelog/trie" 13 14 "gitlab.com/gitlab-org/gitlab-runner/helpers" 15 ) 16 17 const maskedText = "[MASKED]" 18 const defaultBytesLimit = 4 * 1024 * 1024 // 4MB 19 20 type Buffer struct { 21 writer io.WriteCloser 22 lock sync.RWMutex 23 logFile *os.File 24 logSize int 25 logWriter *bufio.Writer 26 advanceBuffer bytes.Buffer 27 bytesLimit int 28 finish chan struct{} 29 30 maskTree *trie.Trie 31 } 32 33 func (b *Buffer) SetMasked(values []string) { 34 if len(values) == 0 { 35 b.maskTree = nil 36 return 37 } 38 39 maskTree := trie.New() 40 for _, value := range values { 41 maskTree.Add(value, nil) 42 } 43 b.maskTree = maskTree 44 } 45 46 func (b *Buffer) SetLimit(size int) { 47 b.bytesLimit = size 48 } 49 50 func (b *Buffer) Size() int { 51 return b.logSize 52 } 53 54 func (b *Buffer) Reader(offset, n int) (io.ReadSeeker, error) { 55 b.lock.Lock() 56 defer b.lock.Unlock() 57 58 err := b.logWriter.Flush() 59 if err != nil { 60 return nil, err 61 } 62 63 return io.NewSectionReader(b.logFile, int64(offset), int64(n)), nil 64 } 65 66 func (b *Buffer) Bytes(offset, n int) ([]byte, error) { 67 reader, err := b.Reader(offset, n) 68 if err != nil { 69 return nil, err 70 } 71 72 return ioutil.ReadAll(reader) 73 } 74 75 func (b *Buffer) Write(data []byte) (n int, err error) { 76 return b.writer.Write(data) 77 } 78 79 func (b *Buffer) Finish() { 80 // wait for trace to finish 81 b.writer.Close() 82 <-b.finish 83 } 84 85 func (b *Buffer) Close() { 86 _ = b.logFile.Close() 87 _ = os.Remove(b.logFile.Name()) 88 } 89 90 func (b *Buffer) advanceAllUnsafe() error { 91 n, err := b.advanceBuffer.WriteTo(b.logWriter) 92 b.logSize += int(n) 93 return err 94 } 95 96 func (b *Buffer) advanceAll() { 97 b.lock.Lock() 98 defer b.lock.Unlock() 99 100 b.advanceAllUnsafe() 101 } 102 103 // advanceLogUnsafe is assumed to be run every character 104 func (b *Buffer) advanceLogUnsafe() error { 105 // advance all if no masking is enabled 106 if b.maskTree == nil { 107 return b.advanceAllUnsafe() 108 } 109 110 rest := b.advanceBuffer.String() 111 results := b.maskTree.Search(rest) 112 if len(results) == 0 { 113 // we can advance as no match was found 114 return b.advanceAllUnsafe() 115 } 116 117 // full match was found 118 if len(results) == 1 && results[0].Key == rest { 119 b.advanceBuffer.Reset() 120 b.advanceBuffer.WriteString(maskedText) 121 return b.advanceAllUnsafe() 122 } 123 124 // partial match, wait for more characters 125 return nil 126 } 127 128 func (b *Buffer) limitExceededMessage() string { 129 return fmt.Sprintf("\n%sJob's log exceeded limit of %v bytes.%s\n", helpers.ANSI_BOLD_RED, b.bytesLimit, helpers.ANSI_RESET) 130 } 131 132 func (b *Buffer) writeRune(r rune) error { 133 b.lock.Lock() 134 defer b.lock.Unlock() 135 136 // over trace limit 137 if b.logSize > b.bytesLimit { 138 return io.EOF 139 } 140 141 if _, err := b.advanceBuffer.WriteRune(r); err != nil { 142 return err 143 } 144 145 if err := b.advanceLogUnsafe(); err != nil { 146 return err 147 } 148 149 // under trace limit 150 if b.logSize <= b.bytesLimit { 151 return nil 152 } 153 154 b.advanceBuffer.Reset() 155 b.advanceBuffer.WriteString(b.limitExceededMessage()) 156 return b.advanceAllUnsafe() 157 } 158 159 func (b *Buffer) process(pipe *io.PipeReader) { 160 defer pipe.Close() 161 162 reader := bufio.NewReader(pipe) 163 164 for { 165 r, s, err := reader.ReadRune() 166 if s <= 0 { 167 break 168 } else if err == nil { 169 b.writeRune(r) 170 } else { 171 // ignore invalid characters 172 continue 173 } 174 } 175 176 b.advanceAll() 177 close(b.finish) 178 } 179 180 func New() (*Buffer, error) { 181 logFile, err := ioutil.TempFile("", "trace") 182 if err != nil { 183 return nil, err 184 } 185 186 reader, writer := io.Pipe() 187 buffer := &Buffer{ 188 writer: writer, 189 bytesLimit: defaultBytesLimit, 190 finish: make(chan struct{}), 191 logFile: logFile, 192 logWriter: bufio.NewWriter(logFile), 193 } 194 go buffer.process(reader) 195 return buffer, nil 196 }