gitech

1
// Copyright 2021 The Gitea Authors. All rights reserved.
2
// SPDX-License-Identifier: MIT
3

4
package filebuffer
5

6
import (
7
	"bytes"
8
	"errors"
9
	"io"
10
	"math"
11
	"os"
12
)
13

14
var (
15
	// ErrInvalidMemorySize occurs if the memory size is not in a valid range
16
	ErrInvalidMemorySize = errors.New("Memory size must be greater 0 and lower math.MaxInt32")
17
	// ErrWriteAfterRead occurs if Write is called after a read operation
18
	ErrWriteAfterRead = errors.New("Write is unsupported after a read operation")
19
)
20

21
type readAtSeeker interface {
22
	io.ReadSeeker
23
	io.ReaderAt
24
}
25

26
// FileBackedBuffer uses a memory buffer with a fixed size.
27
// If more data is written a temporary file is used instead.
28
// It implements io.ReadWriteCloser, io.ReadSeekCloser and io.ReaderAt
29
type FileBackedBuffer struct {
30
	maxMemorySize int64
31
	size          int64
32
	buffer        bytes.Buffer
33
	file          *os.File
34
	reader        readAtSeeker
35
}
36

37
// New creates a file backed buffer with a specific maximum memory size
38
func New(maxMemorySize int) (*FileBackedBuffer, error) {
39
	if maxMemorySize < 0 || maxMemorySize > math.MaxInt32 {
40
		return nil, ErrInvalidMemorySize
41
	}
42

43
	return &FileBackedBuffer{
44
		maxMemorySize: int64(maxMemorySize),
45
	}, nil
46
}
47

48
// CreateFromReader creates a file backed buffer and copies the provided reader data into it.
49
func CreateFromReader(r io.Reader, maxMemorySize int) (*FileBackedBuffer, error) {
50
	b, err := New(maxMemorySize)
51
	if err != nil {
52
		return nil, err
53
	}
54

55
	_, err = io.Copy(b, r)
56
	if err != nil {
57
		return nil, err
58
	}
59

60
	return b, nil
61
}
62

63
// Write implements io.Writer
64
func (b *FileBackedBuffer) Write(p []byte) (int, error) {
65
	if b.reader != nil {
66
		return 0, ErrWriteAfterRead
67
	}
68

69
	var n int
70
	var err error
71

72
	if b.file != nil {
73
		n, err = b.file.Write(p)
74
	} else {
75
		if b.size+int64(len(p)) > b.maxMemorySize {
76
			b.file, err = os.CreateTemp("", "gitea-buffer-")
77
			if err != nil {
78
				return 0, err
79
			}
80

81
			_, err = io.Copy(b.file, &b.buffer)
82
			if err != nil {
83
				return 0, err
84
			}
85

86
			return b.Write(p)
87
		}
88

89
		n, err = b.buffer.Write(p)
90
	}
91

92
	if err != nil {
93
		return n, err
94
	}
95
	b.size += int64(n)
96
	return n, nil
97
}
98

99
// Size returns the byte size of the buffered data
100
func (b *FileBackedBuffer) Size() int64 {
101
	return b.size
102
}
103

104
func (b *FileBackedBuffer) switchToReader() error {
105
	if b.reader != nil {
106
		return nil
107
	}
108

109
	if b.file != nil {
110
		if _, err := b.file.Seek(0, io.SeekStart); err != nil {
111
			return err
112
		}
113
		b.reader = b.file
114
	} else {
115
		b.reader = bytes.NewReader(b.buffer.Bytes())
116
	}
117
	return nil
118
}
119

120
// Read implements io.Reader
121
func (b *FileBackedBuffer) Read(p []byte) (int, error) {
122
	if err := b.switchToReader(); err != nil {
123
		return 0, err
124
	}
125

126
	return b.reader.Read(p)
127
}
128

129
// ReadAt implements io.ReaderAt
130
func (b *FileBackedBuffer) ReadAt(p []byte, off int64) (int, error) {
131
	if err := b.switchToReader(); err != nil {
132
		return 0, err
133
	}
134

135
	return b.reader.ReadAt(p, off)
136
}
137

138
// Seek implements io.Seeker
139
func (b *FileBackedBuffer) Seek(offset int64, whence int) (int64, error) {
140
	if err := b.switchToReader(); err != nil {
141
		return 0, err
142
	}
143

144
	return b.reader.Seek(offset, whence)
145
}
146

147
// Close implements io.Closer
148
func (b *FileBackedBuffer) Close() error {
149
	if b.file != nil {
150
		err := b.file.Close()
151
		os.Remove(b.file.Name())
152
		b.file = nil
153
		return err
154
	}
155
	return nil
156
}
157