113 lines
3.4 KiB
Go
113 lines
3.4 KiB
Go
// Package bufferwritter provides a contentWriter implementation using a shared buffer pool for memory management.
|
|
package bufferwriter
|
|
|
|
import (
|
|
"fmt"
|
|
"io"
|
|
"time"
|
|
|
|
"github.com/trufflesecurity/trufflehog/v3/pkg/buffers/buffer"
|
|
"github.com/trufflesecurity/trufflehog/v3/pkg/buffers/pool"
|
|
"github.com/trufflesecurity/trufflehog/v3/pkg/context"
|
|
)
|
|
|
|
type metrics struct{}
|
|
|
|
func (metrics) recordDataProcessed(size int64, dur time.Duration) {
|
|
writeSize.Observe(float64(size))
|
|
totalWriteDuration.Add(float64(dur.Microseconds()))
|
|
}
|
|
|
|
func init() { bufferPool = pool.NewBufferPool() }
|
|
|
|
// bufferPool is the shared Buffer pool used by all BufferedFileWriters.
|
|
// This allows for efficient reuse of buffers across multiple writers.
|
|
var bufferPool *pool.Pool
|
|
|
|
// state represents the current mode of buffer.
|
|
type state uint8
|
|
|
|
const (
|
|
// writeOnly indicates the buffer is in write-only mode.
|
|
writeOnly state = iota
|
|
// readOnly indicates the buffer has been closed and is in read-only mode.
|
|
readOnly
|
|
)
|
|
|
|
// BufferWriter implements contentWriter, using a shared buffer pool for memory management.
|
|
type BufferWriter struct {
|
|
buf *buffer.Buffer // The current buffer in use.
|
|
bufPool *pool.Pool // The buffer pool used to manage the buffer.
|
|
size int // The total size of the content written to the buffer.
|
|
state state // The current state of the buffer.
|
|
|
|
metrics metrics
|
|
}
|
|
|
|
// New creates a new instance of BufferWriter.
|
|
func New(_ context.Context) *BufferWriter {
|
|
return &BufferWriter{state: writeOnly, bufPool: bufferPool}
|
|
}
|
|
|
|
// NewFromReader creates a new instance of BufferWriter and writes the content from the provided reader to the buffer.
|
|
func NewFromReader(ctx context.Context, r io.Reader) (*BufferWriter, error) {
|
|
buf := New(ctx)
|
|
n, err := io.Copy(buf, r)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("error writing to buffer writer: %w", err)
|
|
}
|
|
|
|
ctx.Logger().V(3).Info("file written to buffer writer", "bytes", n)
|
|
|
|
return buf, nil
|
|
}
|
|
|
|
// Write delegates the writing operation to the underlying bytes.Buffer.
|
|
func (b *BufferWriter) Write(data []byte) (int, error) {
|
|
if b.state != writeOnly {
|
|
return 0, fmt.Errorf("buffer must be in write-only mode to write data; current state: %d", b.state)
|
|
}
|
|
if b.buf == nil {
|
|
b.buf = b.bufPool.Get()
|
|
if b.buf == nil {
|
|
b.buf = buffer.NewBuffer()
|
|
}
|
|
}
|
|
|
|
size := len(data)
|
|
b.size += size
|
|
start := time.Now()
|
|
defer func(start time.Time) {
|
|
b.metrics.recordDataProcessed(int64(size), time.Since(start))
|
|
}(start)
|
|
return b.buf.Write(data)
|
|
}
|
|
|
|
// ReadCloser provides a read-closer for the buffer's content.
|
|
// It wraps the buffer's content in a NopCloser to provide a ReadCloser without additional closing behavior,
|
|
// as closing a bytes.Buffer is a no-op.
|
|
func (b *BufferWriter) ReadCloser() (io.ReadCloser, error) {
|
|
if b.state != readOnly {
|
|
return nil, fmt.Errorf("buffer is in read-only mode")
|
|
}
|
|
|
|
return buffer.ReadCloser(b.buf.Bytes(), func() { b.bufPool.Put(b.buf) }), nil
|
|
}
|
|
|
|
// CloseForWriting is a no-op for buffer, as there is no resource cleanup needed for bytes.Buffer.
|
|
func (b *BufferWriter) CloseForWriting() error {
|
|
b.state = readOnly
|
|
return nil
|
|
}
|
|
|
|
// String returns the buffer's content as a string.
|
|
func (b *BufferWriter) String() (string, error) {
|
|
if b.buf == nil {
|
|
return "", fmt.Errorf("buffer is nil")
|
|
}
|
|
return b.buf.String(), nil
|
|
}
|
|
|
|
// Len returns the length of the buffer's content.
|
|
func (b *BufferWriter) Len() int { return b.size }
|