first-commit

modules/util/filebuffer/file_backed_buffer.go

@@ -0,0 +1,133 @@
+// Copyright 2021 The Gitea Authors. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+package filebuffer
+
+import (
+	"bytes"
+	"errors"
+	"io"
+	"os"
+)
+
+var ErrWriteAfterRead = errors.New("write is unsupported after a read operation") // occurs if Write is called after a read operation
+
+type readAtSeeker interface {
+	io.ReadSeeker
+	io.ReaderAt
+}
+
+// FileBackedBuffer uses a memory buffer with a fixed size.
+// If more data is written a temporary file is used instead.
+// It implements io.ReadWriteCloser, io.ReadSeekCloser and io.ReaderAt
+type FileBackedBuffer struct {
+	maxMemorySize int64
+	size          int64
+	buffer        bytes.Buffer
+	tempDir       string
+	file          *os.File
+	reader        readAtSeeker
+}
+
+// New creates a file backed buffer with a specific maximum memory size
+func New(maxMemorySize int, tempDir string) *FileBackedBuffer {
+	return &FileBackedBuffer{
+		maxMemorySize: int64(maxMemorySize),
+		tempDir:       tempDir,
+	}
+}
+
+// Write implements io.Writer
+func (b *FileBackedBuffer) Write(p []byte) (int, error) {
+	if b.reader != nil {
+		return 0, ErrWriteAfterRead
+	}
+
+	var n int
+	var err error
+
+	if b.file != nil {
+		n, err = b.file.Write(p)
+	} else {
+		if b.size+int64(len(p)) > b.maxMemorySize {
+			b.file, err = os.CreateTemp(b.tempDir, "gitea-buffer-")
+			if err != nil {
+				return 0, err
+			}
+
+			_, err = io.Copy(b.file, &b.buffer)
+			if err != nil {
+				return 0, err
+			}
+
+			return b.Write(p)
+		}
+
+		n, err = b.buffer.Write(p)
+	}
+
+	if err != nil {
+		return n, err
+	}
+	b.size += int64(n)
+	return n, nil
+}
+
+// Size returns the byte size of the buffered data
+func (b *FileBackedBuffer) Size() int64 {
+	return b.size
+}
+
+func (b *FileBackedBuffer) switchToReader() error {
+	if b.reader != nil {
+		return nil
+	}
+
+	if b.file != nil {
+		if _, err := b.file.Seek(0, io.SeekStart); err != nil {
+			return err
+		}
+		b.reader = b.file
+	} else {
+		b.reader = bytes.NewReader(b.buffer.Bytes())
+	}
+	return nil
+}
+
+// Read implements io.Reader
+func (b *FileBackedBuffer) Read(p []byte) (int, error) {
+	if err := b.switchToReader(); err != nil {
+		return 0, err
+	}
+
+	return b.reader.Read(p)
+}
+
+// ReadAt implements io.ReaderAt
+func (b *FileBackedBuffer) ReadAt(p []byte, off int64) (int, error) {
+	if err := b.switchToReader(); err != nil {
+		return 0, err
+	}
+
+	return b.reader.ReadAt(p, off)
+}
+
+// Seek implements io.Seeker
+func (b *FileBackedBuffer) Seek(offset int64, whence int) (int64, error) {
+	if err := b.switchToReader(); err != nil {
+		return 0, err
+	}
+
+	return b.reader.Seek(offset, whence)
+}
+
+// Close implements io.Closer
+func (b *FileBackedBuffer) Close() error {
+	if b.file != nil {
+		err := b.file.Close()
+		_ = os.Remove(b.file.Name())
+		b.file = nil
+		return err
+	}
+	return nil
+}

modules/util/filebuffer/file_backed_buffer_test.go

@@ -0,0 +1,36 @@
+// Copyright 2022 The Gitea Authors. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+package filebuffer
+
+import (
+	"io"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestFileBackedBuffer(t *testing.T) {
+	cases := []struct {
+		MaxMemorySize int
+		Data          string
+	}{
+		{5, "test"},
+		{5, "testtest"},
+	}
+
+	for _, c := range cases {
+		buf := New(c.MaxMemorySize, t.TempDir())
+		_, err := io.Copy(buf, strings.NewReader(c.Data))
+		assert.NoError(t, err)
+
+		assert.EqualValues(t, len(c.Data), buf.Size())
+
+		data, err := io.ReadAll(buf)
+		assert.NoError(t, err)
+		assert.Equal(t, c.Data, string(data))
+
+		assert.NoError(t, buf.Close())
+	}
+}