compress/zstd_compressor.go

// Copyright (C) 2019-2025, Lux Industries, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package compress

import (
	"fmt"
	"math"

	"github.com/klauspost/compress/zstd"
)

var _ Compressor = (*ZstdCompressor)(nil)

func NewZstdCompressor(maxSize int64) (*ZstdCompressor, error) {
	return NewZstdCompressorWithLevel(maxSize, zstd.SpeedDefault)
}

func NewZstdCompressorWithLevel(maxSize int64, level zstd.EncoderLevel) (*ZstdCompressor, error) {
	if maxSize == math.MaxInt64 {
		// "Decompress" creates "io.LimitReader" with max size + 1:
		// if the max size + 1 overflows, "io.LimitReader" reads nothing
		// returning 0 byte for the decompress call
		// require max size < math.MaxInt64 to prevent int64 overflows
		return nil, ErrInvalidMaxSizeCompressor
	}
	// Configure decoder with maximum decompressed size limit
	decoder, err := zstd.NewReader(nil, zstd.WithDecoderMaxMemory(uint64(maxSize)))
	if err != nil {
		return nil, err
	}
	return &ZstdCompressor{
		maxSize: maxSize,
		level:   level,
		decoder: decoder,
	}, nil
}

// ZstdCompressor implements Compressor using zstd compression
type ZstdCompressor struct {
	maxSize int64
	level   zstd.EncoderLevel
	decoder *zstd.Decoder
}

func (z *ZstdCompressor) Compress(msg []byte) ([]byte, error) {
	if int64(len(msg)) > z.maxSize {
		return nil, fmt.Errorf("%w: (%d) > (%d)", ErrMsgTooLarge, len(msg), z.maxSize)
	}
	encoder, err := zstd.NewWriter(nil, zstd.WithEncoderLevel(z.level))
	if err != nil {
		return nil, err
	}
	return encoder.EncodeAll(msg, nil), nil
}

func (z *ZstdCompressor) Decompress(msg []byte) ([]byte, error) {
	decompressed, err := z.decoder.DecodeAll(msg, nil)
	if err != nil {
		// If the decoder returns an error about size limit, wrap it with our error
		if err.Error() == "decompressed size exceeds configured limit" {
			return nil, fmt.Errorf("%w: decompression stopped due to size limit", ErrDecompressedMsgTooLarge)
		}
		return nil, err
	}
	if int64(len(decompressed)) > z.maxSize {
		return nil, fmt.Errorf("%w: (%d) > (%d)", ErrDecompressedMsgTooLarge, len(decompressed), z.maxSize)
	}
	return decompressed, nil
}