1 files changed, 163 insertions, 0 deletions
diff --git a/modules/zstd/zstd.go b/modules/zstd/zstd.go
new file mode 100644
index 0000000..d224944
--- /dev/null
+++ b/modules/zstd/zstd.go
@@ -0,0 +1,163 @@
+// Copyright 2024 The Gitea Authors. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+// Package zstd provides a high-level API for reading and writing zstd-compressed data.
+// It supports both regular and seekable zstd streams.
+// It's not a new wheel, but a wrapper around the zstd and zstd-seekable-format-go packages.
+package zstd
+
+import (
+	"errors"
+	"io"
+
+	seekable "github.com/SaveTheRbtz/zstd-seekable-format-go/pkg"
+	"github.com/klauspost/compress/zstd"
+)
+
+type Writer zstd.Encoder
+
+var _ io.WriteCloser = (*Writer)(nil)
+
+// NewWriter returns a new zstd writer.
+func NewWriter(w io.Writer, opts ...WriterOption) (*Writer, error) {
+	zstdW, err := zstd.NewWriter(w, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return (*Writer)(zstdW), nil
+}
+
+func (w *Writer) Write(p []byte) (int, error) {
+	return (*zstd.Encoder)(w).Write(p)
+}
+
+func (w *Writer) Close() error {
+	return (*zstd.Encoder)(w).Close()
+}
+
+type Reader zstd.Decoder
+
+var _ io.ReadCloser = (*Reader)(nil)
+
+// NewReader returns a new zstd reader.
+func NewReader(r io.Reader, opts ...ReaderOption) (*Reader, error) {
+	zstdR, err := zstd.NewReader(r, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return (*Reader)(zstdR), nil
+}
+
+func (r *Reader) Read(p []byte) (int, error) {
+	return (*zstd.Decoder)(r).Read(p)
+}
+
+func (r *Reader) Close() error {
+	(*zstd.Decoder)(r).Close() // no error returned
+	return nil
+}
+
+type SeekableWriter struct {
+	buf []byte
+	n   int
+	w   seekable.Writer
+}
+
+var _ io.WriteCloser = (*SeekableWriter)(nil)
+
+// NewSeekableWriter returns a zstd writer to compress data to seekable format.
+// blockSize is an important parameter, it should be decided according to the actual business requirements.
+// If it's too small, the compression ratio could be very bad, even no compression at all.
+// If it's too large, it could cost more traffic when reading the data partially from underlying storage.
+func NewSeekableWriter(w io.Writer, blockSize int, opts ...WriterOption) (*SeekableWriter, error) {
+	zstdW, err := zstd.NewWriter(nil, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	seekableW, err := seekable.NewWriter(w, zstdW)
+	if err != nil {
+		return nil, err
+	}
+
+	return &SeekableWriter{
+		buf: make([]byte, blockSize),
+		w:   seekableW,
+	}, nil
+}
+
+func (w *SeekableWriter) Write(p []byte) (int, error) {
+	written := 0
+	for len(p) > 0 {
+		n := copy(w.buf[w.n:], p)
+		w.n += n
+		written += n
+		p = p[n:]
+
+		if w.n == len(w.buf) {
+			if _, err := w.w.Write(w.buf); err != nil {
+				return written, err
+			}
+			w.n = 0
+		}
+	}
+	return written, nil
+}
+
+func (w *SeekableWriter) Close() error {
+	if w.n > 0 {
+		if _, err := w.w.Write(w.buf[:w.n]); err != nil {
+			return err
+		}
+	}
+	return w.w.Close()
+}
+
+type SeekableReader struct {
+	r seekable.Reader
+	c func() error
+}
+
+var _ io.ReadSeekCloser = (*SeekableReader)(nil)
+
+// NewSeekableReader returns a zstd reader to decompress data from seekable format.
+func NewSeekableReader(r io.ReadSeeker, opts ...ReaderOption) (*SeekableReader, error) {
+	zstdR, err := zstd.NewReader(nil, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	seekableR, err := seekable.NewReader(r, zstdR)
+	if err != nil {
+		return nil, err
+	}
+
+	ret := &SeekableReader{
+		r: seekableR,
+	}
+	if closer, ok := r.(io.Closer); ok {
+		ret.c = closer.Close
+	}
+
+	return ret, nil
+}
+
+func (r *SeekableReader) Read(p []byte) (int, error) {
+	return r.r.Read(p)
+}
+
+func (r *SeekableReader) Seek(offset int64, whence int) (int64, error) {
+	return r.r.Seek(offset, whence)
+}
+
+func (r *SeekableReader) Close() error {
+	return errors.Join(
+		func() error {
+			if r.c != nil {
+				return r.c()
+			}
+			return nil
+		}(),
+		r.r.Close(),
+	)
+}