bugabinga.net

package logs

import (
	"bytes"
	"io"
	"sort"
	"sync"
)

const masked = "[MASKED]"

// MaxLogBytes bounds live log output, including the truncation marker.
const MaxLogBytes = 16 << 20

const truncationMarker = "\n[luci] log output truncated\n"

type Sanitizer struct {
	limiter  *limitWriter
	mu       sync.Mutex
	writer   io.Writer
	secrets  [][]byte
	maxLen   int
	pending  []byte
	ansiMode byte
	closed   bool
}

func NewSanitizer(writer io.Writer, secrets [][]byte) *Sanitizer {
	return newSanitizer(writer, secrets, nil)
}

// NewBoundedSanitizer strips ANSI sequences and masks secrets while retaining
// at most MaxLogBytes.  Further output is discarded after one visible marker.
func NewBoundedSanitizer(writer io.Writer, secrets [][]byte) *Sanitizer {
	limiter := &limitWriter{writer: writer, limit: MaxLogBytes}
	return newSanitizer(limiter, secrets, limiter)
}

func newSanitizer(writer io.Writer, secrets [][]byte, limiter *limitWriter) *Sanitizer {
	unique := map[string]struct{}{}
	for _, secret := range secrets {
		if len(secret) > 0 {
			unique[string(secret)] = struct{}{}
		}
	}
	values := make([][]byte, 0, len(unique))
	for secret := range unique {
		values = append(values, []byte(secret))
	}
	sort.Slice(values, func(i, j int) bool { return len(values[i]) > len(values[j]) })
	maxLen := 1
	if len(values) > 0 {
		maxLen = len(values[0])
	}
	return &Sanitizer{writer: writer, secrets: values, maxLen: maxLen, limiter: limiter}
}

// Truncated reports whether NewBoundedSanitizer discarded output.
func (s *Sanitizer) Truncated() bool {
	return s.limiter != nil && s.limiter.Truncated()
}

func (s *Sanitizer) Write(data []byte) (int, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return 0, io.ErrClosedPipe
	}
	for _, value := range data {
		s.consumeANSI(value)
	}
	if err := s.flush(false); err != nil {
		return 0, err
	}
	return len(data), nil
}

func (s *Sanitizer) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil
	}
	s.closed = true
	return s.flush(true)
}

func (s *Sanitizer) consumeANSI(value byte) {
	switch s.ansiMode {
	case 0:
		if value == 0x1b {
			s.ansiMode = 1
			return
		}
		s.pending = append(s.pending, value)
	case 1:
		switch value {
		case '[':
			s.ansiMode = 2
		case ']', 'P', 'X', '^', '_':
			s.ansiMode = 3
		default:
			if value >= 0x20 && value <= 0x2f {
				s.ansiMode = 4
			} else {
				s.ansiMode = 0
			}
		}
	case 2:
		if value >= 0x40 && value <= 0x7e {
			s.ansiMode = 0
		}
	case 3:
		if value == 0x07 {
			s.ansiMode = 0
		} else if value == 0x1b {
			s.ansiMode = 5
		}
	case 4:
		if value >= 0x30 && value <= 0x7e {
			s.ansiMode = 0
		}
	case 5:
		if value == '\\' {
			s.ansiMode = 0
		} else if value != 0x1b {
			s.ansiMode = 3
		}
	}
}

func (s *Sanitizer) flush(all bool) error {
	limit := len(s.pending)
	if !all {
		limit -= s.maxLen - 1
		if limit < 0 {
			limit = 0
		}
	}
	position, plainStart := 0, 0
	for position < limit {
		match := s.secretAt(position)
		if match == 0 {
			position++
			continue
		}
		if plainStart < position {
			if _, err := s.writer.Write(s.pending[plainStart:position]); err != nil {
				return err
			}
		}
		if _, err := io.WriteString(s.writer, masked); err != nil {
			return err
		}
		position += match
		plainStart = position
	}
	if plainStart < position {
		if _, err := s.writer.Write(s.pending[plainStart:position]); err != nil {
			return err
		}
	}
	s.pending = append(s.pending[:0], s.pending[position:]...)
	return nil
}

func (s *Sanitizer) secretAt(position int) int {
	for _, secret := range s.secrets {
		if bytes.HasPrefix(s.pending[position:], secret) {
			return len(secret)
		}
	}
	return 0
}

type limitWriter struct {
	mu        sync.Mutex
	writer    io.Writer
	limit     int
	written   int
	truncated bool
}

func (w *limitWriter) Write(data []byte) (int, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.truncated {
		return len(data), nil
	}
	available := w.limit - len(truncationMarker) - w.written
	if available >= len(data) {
		n, err := w.writer.Write(data)
		w.written += n
		return n, err
	}
	if available > 0 {
		n, err := w.writer.Write(data[:available])
		w.written += n
		if err != nil {
			return n, err
		}
		if n != available {
			return n, io.ErrShortWrite
		}
	}
	n, err := io.WriteString(w.writer, truncationMarker)
	w.written += n
	if err != nil {
		return 0, err
	}
	if n != len(truncationMarker) {
		return 0, io.ErrShortWrite
	}
	w.truncated = true
	return len(data), nil
}

func (w *limitWriter) Truncated() bool {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.truncated
}