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  }