[#7575] use memory+file buffer when rereading the request body (fix #7572)

tools/router/buffer_with_file.go

@@ -0,0 +1,125 @@
+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
+}