tools/router/buffer_with_file.go - master - githubmirror/pocketbase

tools/router/buffer_with_file.go
125 строк3 KB
Gani Georgiev
[#7575] use memory+file buffer when rereading the request body (fix #7572)
09 мар 2026, 18:19
09 мар 2026, 18:19ba8b51a
package router

import (
	"bytes"
	"errors"
	"io"
	"os"
)

var _ io.ReadWriteCloser = (*bufferWithFile)(nil)

// newBufferWithFile initializes and returns a new bufferWithFile with the specified memoryLimit.
//
// If memoryLimit is negative or zero, defaults to [DefaultMaxMemory].
func newBufferWithFile(memoryLimit int64) *bufferWithFile {
	if memoryLimit <= 0 {
		memoryLimit = DefaultMaxMemory
	}

	return &bufferWithFile{
		buf:         new(bytes.Buffer),
		memoryLimit: memoryLimit,
	}
}

// bufferWithFile is similar to [bytes.Buffer] but after the limit it
// fallbacks to a temporary file to minimize excessive memory usage.
type bufferWithFile struct {
	buf            *bytes.Buffer
	file           *os.File
	memoryLimit    int64
	fileReadOffset int64
}

// Read implements the standard [io.Reader] interface by reading
// up to len(p) bytes into p.
func (b *bufferWithFile) Read(p []byte) (n int, err error) {
	if b.buf == nil {
		return 0, errors.New("[bufferWithFile.Read] not initialized or already closed")
	}

	// eagerly get length because bytes.Buffer may resize and change it
	maxToRead := len(p)

	// read first from the memory buffer
	if b.buf.Len() > 0 {
		n, err = b.buf.Read(p)
		if err != nil && err != io.EOF {
			return n, err
		}
	}

	// continue reading from the file to fill the remaining bytes
	if n < maxToRead && b.file != nil {
		fileN, fileErr := b.file.ReadAt(p[n:maxToRead], b.fileReadOffset)
		b.fileReadOffset += int64(fileN)
		n += fileN
		err = fileErr
	}

	// return EOF if the buffers are empty and nothing has been read
	// (to minimize potential breaking changes and for consistency with the bytes.Buffer rules)
	if n == 0 && maxToRead > 0 && err == nil {
		return 0, io.EOF
	}

	return n, err
}

// Write implements the standard [io.Writer] interface by writing the
// content of p into the buffer.
//
// If the current memory buffer doesn't have enough space to hold len(p),
// it write p into a temp disk file.
func (b *bufferWithFile) Write(p []byte) (int, error) {
	if b.buf == nil {
		return 0, errors.New("[bufferWithFile.Write] not initialized or already closed")
	}

	// already above the limit -> continue with the file
	if b.file != nil {
		return b.file.Write(p)
	}

	// above limit -> create and write to file
	if int64(b.buf.Len()+len(p)) > b.memoryLimit {
		if b.file == nil {
			var err error
			b.file, err = os.CreateTemp("", "pb_buffer_file_*")
			if err != nil {
				return 0, err
			}
		}

		return b.file.Write(p)
	}

	// write in memory
	return b.buf.Write(p)
}

// Close implements the standard [io.Closer] interface.
//
// It unsets the memory buffer and will cleanup after the fallback
// temporary file (if exists).
//
// It is safe to call Close multiple times.
// Once Close is invoked the buffer no longer can be used and should be discarded.
func (b *bufferWithFile) Close() error {
	if b.file != nil {
		err := errors.Join(
			b.file.Close(),
			os.Remove(b.file.Name()),
		)
		if err != nil {
			return err
		}

		b.file = nil
	}

	b.buf = nil

	return nil
}