1 files changed, 52 insertions, 0 deletions
diff --git a/workhorse/cmd/gitlab-zip-metadata/limit/reader.go b/workhorse/cmd/gitlab-zip-metadata/limit/reader.go
new file mode 100644
index 00000000000..c1e9bd20993
--- /dev/null
+++ b/workhorse/cmd/gitlab-zip-metadata/limit/reader.go
@@ -0,0 +1,52 @@
+package limit
+
+import (
+	"errors"
+	"io"
+	"sync/atomic"
+)
+
+var ErrLimitExceeded = errors.New("reader limit exceeded")
+
+const megabyte = 1 << 20
+
+// LimitedReaderAt supports running a callback in case of reaching a read limit
+// (bytes), and allows using a smaller limit than a defined offset for a read.
+type LimitedReaderAt struct {
+	read      int64
+	limit     int64
+	parent    io.ReaderAt
+	limitFunc func(int64)
+}
+
+func (r *LimitedReaderAt) ReadAt(p []byte, off int64) (int, error) {
+	if max := r.limit - r.read; int64(len(p)) > max {
+		p = p[0:max]
+	}
+
+	n, err := r.parent.ReadAt(p, off)
+	atomic.AddInt64(&r.read, int64(n))
+
+	if r.read >= r.limit {
+		r.limitFunc(r.read)
+
+		return n, ErrLimitExceeded
+	}
+
+	return n, err
+}
+
+func NewLimitedReaderAt(reader io.ReaderAt, limit int64, limitFunc func(int64)) io.ReaderAt {
+	return &LimitedReaderAt{parent: reader, limit: limit, limitFunc: limitFunc}
+}
+
+// SizeToLimit tries to dermine an appropriate limit in bytes for an archive of
+// a given size. If the size is less than 1 gigabyte we always limit a reader
+// to 100 megabytes, otherwise the limit is 10% of a given size.
+func SizeToLimit(size int64) int64 {
+	if size <= 1024*megabyte {
+		return 100 * megabyte
+	}
+
+	return size / 10
+}