github.com/hashicorp/vault/sdk@v0.13.0/helper/compressutil/compress.go

// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0

package compressutil

import (
	"bytes"
	"compress/gzip"
	"compress/lzw"
	"fmt"
	"io"

	"github.com/golang/snappy"
	"github.com/pierrec/lz4"
)

const (
	// A byte value used as a canary prefix for the compressed information
	// which is used to distinguish if a JSON input is compressed or not.
	// The value of this constant should not be a first character of any
	// valid JSON string.

	CompressionTypeGzip        = "gzip"
	CompressionCanaryGzip byte = 'G'

	CompressionTypeLZW        = "lzw"
	CompressionCanaryLZW byte = 'L'

	CompressionTypeSnappy        = "snappy"
	CompressionCanarySnappy byte = 'S'

	CompressionTypeLZ4        = "lz4"
	CompressionCanaryLZ4 byte = '4'
)

// CompressUtilReadCloser embeds the snappy reader which implements the io.Reader
// interface. The decompress procedure in this utility expects an
// io.ReadCloser. This type implements the io.Closer interface to retain the
// generic way of decompression.
type CompressUtilReadCloser struct {
	io.Reader
}

// Close is a noop method implemented only to satisfy the io.Closer interface
func (c *CompressUtilReadCloser) Close() error {
	return nil
}

// CompressionConfig is used to select a compression type to be performed by
// Compress and Decompress utilities.
// Supported types are:
// * CompressionTypeLZW
// * CompressionTypeGzip
// * CompressionTypeSnappy
// * CompressionTypeLZ4
//
// When using CompressionTypeGzip, the compression levels can also be chosen:
// * gzip.DefaultCompression
// * gzip.BestSpeed
// * gzip.BestCompression
type CompressionConfig struct {
	// Type of the compression algorithm to be used
	Type string

	// When using Gzip format, the compression level to employ
	GzipCompressionLevel int
}

// Compress places the canary byte in a buffer and uses the same buffer to fill
// in the compressed information of the given input. The configuration supports
// two type of compression: LZW and Gzip. When using Gzip compression format,
// if GzipCompressionLevel is not specified, the 'gzip.DefaultCompression' will
// be assumed.
func Compress(data []byte, config *CompressionConfig) ([]byte, error) {
	var buf bytes.Buffer
	var writer io.WriteCloser
	var err error

	if config == nil {
		return nil, fmt.Errorf("config is nil")
	}

	// Write the canary into the buffer and create writer to compress the
	// input data based on the configured type
	switch config.Type {
	case CompressionTypeLZW:
		buf.Write([]byte{CompressionCanaryLZW})
		writer = lzw.NewWriter(&buf, lzw.LSB, 8)

	case CompressionTypeGzip:
		buf.Write([]byte{CompressionCanaryGzip})

		switch {
		case config.GzipCompressionLevel == gzip.BestCompression,
			config.GzipCompressionLevel == gzip.BestSpeed,
			config.GzipCompressionLevel == gzip.DefaultCompression:
			// These are valid compression levels
		default:
			// If compression level is set to NoCompression or to
			// any invalid value, fallback to DefaultCompression
			config.GzipCompressionLevel = gzip.DefaultCompression
		}
		writer, err = gzip.NewWriterLevel(&buf, config.GzipCompressionLevel)

	case CompressionTypeSnappy:
		buf.Write([]byte{CompressionCanarySnappy})
		writer = snappy.NewBufferedWriter(&buf)

	case CompressionTypeLZ4:
		buf.Write([]byte{CompressionCanaryLZ4})
		writer = lz4.NewWriter(&buf)

	default:
		return nil, fmt.Errorf("unsupported compression type")
	}

	if err != nil {
		return nil, fmt.Errorf("failed to create a compression writer: %w", err)
	}

	if writer == nil {
		return nil, fmt.Errorf("failed to create a compression writer")
	}

	// Compress the input and place it in the same buffer containing the
	// canary byte.
	if _, err = writer.Write(data); err != nil {
		return nil, fmt.Errorf("failed to compress input data: err: %w", err)
	}

	// Close the io.WriteCloser
	if err = writer.Close(); err != nil {
		return nil, err
	}

	// Return the compressed bytes with canary byte at the start
	return buf.Bytes(), nil
}

// Decompress checks if the first byte in the input matches the canary byte.
// If the first byte is a canary byte, then the input past the canary byte
// will be decompressed using the method specified in the given configuration.
// If the first byte isn't a canary byte, then the utility returns a boolean
// value indicating that the input was not compressed.
func Decompress(data []byte) ([]byte, bool, error) {
	bytes, _, notCompressed, err := DecompressWithCanary(data)
	return bytes, notCompressed, err
}

// DecompressWithCanary checks if the first byte in the input matches the canary byte.
// If the first byte is a canary byte, then the input past the canary byte
// will be decompressed using the method specified in the given configuration. The type of compression used is also
// returned. If the first byte isn't a canary byte, then the utility returns a boolean
// value indicating that the input was not compressed.
func DecompressWithCanary(data []byte) ([]byte, string, bool, error) {
	var err error
	var reader io.ReadCloser
	var compressionType string
	if data == nil || len(data) == 0 {
		return nil, "", false, fmt.Errorf("'data' being decompressed is empty")
	}

	canary := data[0]
	cData := data[1:]

	switch canary {
	// If the first byte matches the canary byte, remove the canary
	// byte and try to decompress the data that is after the canary.
	case CompressionCanaryGzip:
		if len(data) < 2 {
			return nil, "", false, fmt.Errorf("invalid 'data' after the canary")
		}
		reader, err = gzip.NewReader(bytes.NewReader(cData))
		compressionType = CompressionTypeGzip

	case CompressionCanaryLZW:
		if len(data) < 2 {
			return nil, "", false, fmt.Errorf("invalid 'data' after the canary")
		}
		reader = lzw.NewReader(bytes.NewReader(cData), lzw.LSB, 8)
		compressionType = CompressionTypeLZW

	case CompressionCanarySnappy:
		if len(data) < 2 {
			return nil, "", false, fmt.Errorf("invalid 'data' after the canary")
		}
		reader = &CompressUtilReadCloser{
			Reader: snappy.NewReader(bytes.NewReader(cData)),
		}
		compressionType = CompressionTypeSnappy

	case CompressionCanaryLZ4:
		if len(data) < 2 {
			return nil, "", false, fmt.Errorf("invalid 'data' after the canary")
		}
		reader = &CompressUtilReadCloser{
			Reader: lz4.NewReader(bytes.NewReader(cData)),
		}
		compressionType = CompressionTypeLZ4

	default:
		// If the first byte doesn't match the canary byte, it means
		// that the content was not compressed at all. Indicate the
		// caller that the input was not compressed.
		return nil, "", true, nil
	}
	if err != nil {
		return nil, "", false, fmt.Errorf("failed to create a compression reader: %w", err)
	}
	if reader == nil {
		return nil, "", false, fmt.Errorf("failed to create a compression reader")
	}

	// Close the io.ReadCloser
	defer reader.Close()

	// Read all the compressed data into a buffer
	var buf bytes.Buffer

	// Read the compressed data into a buffer, but do so
	// slowly to prevent reading all the data into memory
	// at once (protecting against e.g. zip bombs).
	for {
		_, err := io.CopyN(&buf, reader, 1024)
		if err != nil {
			if err == io.EOF {
				break
			}
			return nil, "", false, err
		}
	}

	return buf.Bytes(), compressionType, false, nil
}