Data scrubbing options for protecting sensitive data https://godoc.org/maze.io/x/scrub
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

167 lines
3.6 KiB

package scrub
import (
"bytes"
"io"
"sync"
"sync/atomic"
"time"
)
// Buffer can be written to and will scrub.
// By default the Buffer scrubs before each Read invocation.
type Buffer struct {
Scrubber Scrubber
mutex sync.RWMutex
b *bytes.Buffer // buffer contents
i int64 // index since last read/scrub
w chan struct{}
scrubAlways bool
scrubSize int
closeIdle chan struct{}
}
// NewBuffer returns a Buffer with the selected Scrubber.
func NewBuffer(scrubber Scrubber) *Buffer {
return &Buffer{
Scrubber: scrubber,
b: new(bytes.Buffer),
w: make(chan struct{}),
}
}
// Reset cancels any remaining flushers and empties the buffer.
func (b *Buffer) Reset() {
b.mutex.Lock()
// Signal any flusher.
if b.closeIdle != nil {
close(b.closeIdle)
b.closeIdle = nil
}
// Signal any blocked reader.
select {
case b.w <- struct{}{}:
default:
}
// Flush the buffer.
b.b.Reset()
b.i = 0
// Back to defaults
b.scrubAlways = false
b.scrubSize = 0
b.mutex.Unlock()
}
// ScrubAfter scrubs the contents of the internal buffer after no Write has happened for timeout.
// This is useful to make sure no unscrubbed secrets remain in memory, or if you want to scrub
// after a burst of writes happen and then pauses, such as when scrubbing output of a terminal
// session.
// If timeout <= 0 then the Buffer will be scrubbed for each write.
func (b *Buffer) ScrubAfter(timeout time.Duration) CancelFunc {
b.mutex.Lock()
if b.scrubAlways = timeout <= 0; b.scrubAlways {
b.mutex.Unlock()
return func() {}
}
// If ScrubAfter is already running, close previous runner.
if b.closeIdle != nil {
close(b.closeIdle)
b.closeIdle = nil
}
b.closeIdle = make(chan struct{})
go func(closed <-chan struct{}, timeout time.Duration) {
ticker := time.NewTicker(timeout)
defer ticker.Stop()
for {
select {
case <-closed:
return
case <-ticker.C:
b.scrub()
}
}
}(b.closeIdle, timeout)
b.mutex.Unlock()
return func() {
b.mutex.Lock()
close(b.closeIdle)
b.closeIdle = nil
b.mutex.Unlock()
}
}
// ScrubSize scrubs the contents of the buffer if it is larger than size.
// If size is equal to or less than 0, the buffer will be scrubbed immediately.
func (b *Buffer) ScrubSize(size int) {
b.mutex.Lock()
b.scrubSize = size
b.scrubAlways = size <= 0
b.mutex.Unlock()
}
// Read data from the buffer.
// If no data has been written yet, this will block until a write occurs.
func (b *Buffer) Read(p []byte) (n int, err error) {
if len(p) == 0 {
return
}
if b.b.Len() == 0 && atomic.LoadInt64(&b.i) == 0 {
// Wait for a write.
<-b.w
}
b.scrub()
b.mutex.RLock()
n, err = b.b.Read(p)
b.mutex.RUnlock()
return
}
// Write data to the buffer.
func (b *Buffer) Write(p []byte) (int, error) {
b.mutex.Lock()
n, err := b.b.Write(p)
b.mutex.Unlock()
if n > 0 {
// Signal write to channel.
select {
case b.w <- struct{}{}:
default:
}
// Scrub if enabled.
if b.scrubAlways || (b.scrubSize > 0 && b.b.Len() >= b.scrubSize) {
b.scrub()
}
}
return n, err
}
func (b *Buffer) WriteString(s string) (n int, err error) {
return b.Write([]byte(s))
}
func (b *Buffer) scrub() {
b.mutex.Lock()
n := new(bytes.Buffer)
if i := atomic.LoadInt64(&b.i); i > 0 {
// Copy unmodified contents of previous buffer to new buffer.
io.CopyN(n, b.b, i)
}
// Copy scrubbed contents of remaining contents of the buffer to new buffer.
n.WriteString(b.Scrubber.Scrub(b.b.String()))
// Swap buffers and advance the index to the new length.
b.b = n
atomic.StoreInt64(&b.i, int64(n.Len()))
b.mutex.Unlock()
}
// CancelFunc cancels a Buffer flusher.
type CancelFunc func()