reader.go

package decodedsource

import (
	"bufio"
	"bytes"
	"compress/bzip2"
	"compress/gzip"
	"errors"
	"fmt"
	"io"

	"github.com/klauspost/compress/zstd"
	"github.com/pierrec/lz4/v4"
	"github.com/ulikunitz/xz"
)

// Compression identifies a whole-stream compression envelope.
type Compression uint8

const (
	// CompressionAuto detects a supported compression format from its magic.
	// Inputs without a recognized magic are treated as plain streams.
	CompressionAuto Compression = iota
	// CompressionPlain passes the input through unchanged.
	CompressionPlain
	// CompressionGzip decodes a gzip stream, including concatenated members.
	CompressionGzip
	// CompressionZstd decodes a zstd stream.
	CompressionZstd
	// CompressionBzip2 decodes a bzip2 stream.
	CompressionBzip2
	// CompressionXZ decodes an xz stream.
	CompressionXZ
	// CompressionLZ4 decodes an LZ4 frame stream.
	CompressionLZ4
)

// ErrUnsupportedCompression reports an unknown Compression value.
var ErrUnsupportedCompression = errors.New("unsupported compression")

func (c Compression) String() string {
	switch c {
	case CompressionAuto:
		return "auto"
	case CompressionPlain:
		return "plain"
	case CompressionGzip:
		return "gzip"
	case CompressionZstd:
		return "zstd"
	case CompressionBzip2:
		return "bzip2"
	case CompressionXZ:
		return "xz"
	case CompressionLZ4:
		return "lz4"
	default:
		return fmt.Sprintf("compression-%d", uint8(c))
	}
}

// DetectCompression detects a supported compression envelope from prefix.
// Unrecognized input is classified as CompressionPlain.
func DetectCompression(prefix []byte) Compression {
	switch {
	case len(prefix) >= 3 && bytes.Equal(prefix[:3], []byte{0x1f, 0x8b, 0x08}):
		return CompressionGzip
	case len(prefix) >= 4 && bytes.Equal(prefix[:4], []byte{0x28, 0xb5, 0x2f, 0xfd}):
		return CompressionZstd
	case len(prefix) >= 4 && bytes.Equal(prefix[:3], []byte("BZh")) && prefix[3] >= '1' && prefix[3] <= '9':
		return CompressionBzip2
	case len(prefix) >= 6 && bytes.Equal(prefix[:6], []byte{0xfd, '7', 'z', 'X', 'Z', 0x00}):
		return CompressionXZ
	case len(prefix) >= 4 && bytes.Equal(prefix[:4], []byte{0x04, 0x22, 0x4d, 0x18}):
		return CompressionLZ4
	default:
		return CompressionPlain
	}
}

// NewReader returns a reader for the decoded stream. CompressionAuto detects
// the envelope from its magic. Closing the returned reader releases decoder
// resources but does not close r; ownership of r remains with the caller.
func NewReader(r io.Reader, compression Compression) (io.ReadCloser, error) {
	if r == nil {
		return nil, errors.New("nil input reader")
	}

	if compression == CompressionAuto {
		buffered := bufio.NewReader(r)
		prefix, err := buffered.Peek(6)
		if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, bufio.ErrBufferFull) {
			return nil, err
		}
		compression = DetectCompression(prefix)
		r = buffered
	}

	switch compression {
	case CompressionPlain:
		return io.NopCloser(r), nil
	case CompressionGzip:
		return gzip.NewReader(r)
	case CompressionZstd:
		decoder, err := zstd.NewReader(r)
		if err != nil {
			return nil, err
		}
		return &readerCloser{
			Reader: decoder,
			close: func() error {
				decoder.Close()
				return nil
			},
		}, nil
	case CompressionBzip2:
		return io.NopCloser(bzip2.NewReader(r)), nil
	case CompressionXZ:
		reader, err := xz.NewReader(r)
		if err != nil {
			return nil, err
		}
		return io.NopCloser(reader), nil
	case CompressionLZ4:
		return io.NopCloser(lz4.NewReader(r)), nil
	default:
		return nil, fmt.Errorf("%w: %s", ErrUnsupportedCompression, compression)
	}
}

type readerCloser struct {
	io.Reader
	close func() error
}

func (r *readerCloser) Close() error {
	if r.close == nil {
		return nil
	}
	return r.close()
}