github.com/hamba/avro/v2@v2.22.1-0.20240518180522-aff3955acf7d/ocf/codec.go (about)

     1  package ocf
     2  
     3  import (
     4  	"bytes"
     5  	"compress/flate"
     6  	"encoding/binary"
     7  	"errors"
     8  	"fmt"
     9  	"hash/crc32"
    10  	"io"
    11  
    12  	"github.com/golang/snappy"
    13  	"github.com/klauspost/compress/zstd"
    14  )
    15  
    16  // CodecName represents a compression codec name.
    17  type CodecName string
    18  
    19  // Supported compression codecs.
    20  const (
    21  	Null      CodecName = "null"
    22  	Deflate   CodecName = "deflate"
    23  	Snappy    CodecName = "snappy"
    24  	ZStandard CodecName = "zstandard"
    25  )
    26  
    27  func resolveCodec(name CodecName, lvl int) (Codec, error) {
    28  	switch name {
    29  	case Null, "":
    30  		return &NullCodec{}, nil
    31  
    32  	case Deflate:
    33  		return &DeflateCodec{compLvl: lvl}, nil
    34  
    35  	case Snappy:
    36  		return &SnappyCodec{}, nil
    37  
    38  	case ZStandard:
    39  		return &ZStandardCodec{}, nil
    40  
    41  	default:
    42  		return nil, fmt.Errorf("unknown codec %s", name)
    43  	}
    44  }
    45  
    46  // Codec represents a compression codec.
    47  type Codec interface {
    48  	// Decode decodes the given bytes.
    49  	Decode([]byte) ([]byte, error)
    50  	// Encode encodes the given bytes.
    51  	Encode([]byte) []byte
    52  }
    53  
    54  // NullCodec is a no op codec.
    55  type NullCodec struct{}
    56  
    57  // Decode decodes the given bytes.
    58  func (*NullCodec) Decode(b []byte) ([]byte, error) {
    59  	return b, nil
    60  }
    61  
    62  // Encode encodes the given bytes.
    63  func (*NullCodec) Encode(b []byte) []byte {
    64  	return b
    65  }
    66  
    67  // DeflateCodec is a flate compression codec.
    68  type DeflateCodec struct {
    69  	compLvl int
    70  }
    71  
    72  // Decode decodes the given bytes.
    73  func (c *DeflateCodec) Decode(b []byte) ([]byte, error) {
    74  	r := flate.NewReader(bytes.NewBuffer(b))
    75  	data, err := io.ReadAll(r)
    76  	if err != nil {
    77  		_ = r.Close()
    78  		return nil, err
    79  	}
    80  	_ = r.Close()
    81  
    82  	return data, nil
    83  }
    84  
    85  // Encode encodes the given bytes.
    86  func (c *DeflateCodec) Encode(b []byte) []byte {
    87  	data := bytes.NewBuffer(make([]byte, 0, len(b)))
    88  
    89  	w, _ := flate.NewWriter(data, c.compLvl)
    90  	_, _ = w.Write(b)
    91  	_ = w.Close()
    92  
    93  	return data.Bytes()
    94  }
    95  
    96  // SnappyCodec is a snappy compression codec.
    97  type SnappyCodec struct{}
    98  
    99  // Decode decodes the given bytes.
   100  func (*SnappyCodec) Decode(b []byte) ([]byte, error) {
   101  	l := len(b)
   102  	if l < 5 {
   103  		return nil, errors.New("block does not contain snappy checksum")
   104  	}
   105  
   106  	dst, err := snappy.Decode(nil, b[:l-4])
   107  	if err != nil {
   108  		return nil, err
   109  	}
   110  
   111  	crc := binary.BigEndian.Uint32(b[l-4:])
   112  	if crc32.ChecksumIEEE(dst) != crc {
   113  		return nil, errors.New("snappy checksum mismatch")
   114  	}
   115  
   116  	return dst, nil
   117  }
   118  
   119  // Encode encodes the given bytes.
   120  func (*SnappyCodec) Encode(b []byte) []byte {
   121  	dst := snappy.Encode(nil, b)
   122  
   123  	dst = append(dst, 0, 0, 0, 0)
   124  	binary.BigEndian.PutUint32(dst[len(dst)-4:], crc32.ChecksumIEEE(b))
   125  
   126  	return dst
   127  }
   128  
   129  // ZStandardCodec is a zstandard compression codec.
   130  type ZStandardCodec struct{}
   131  
   132  // Decode decodes the given bytes.
   133  func (*ZStandardCodec) Decode(b []byte) ([]byte, error) {
   134  	dec, _ := zstd.NewReader(nil)
   135  	defer dec.Close()
   136  
   137  	return dec.DecodeAll(b, nil)
   138  }
   139  
   140  // Encode encodes the given bytes.
   141  func (*ZStandardCodec) Encode(b []byte) []byte {
   142  	enc, _ := zstd.NewWriter(nil)
   143  	defer func() { _ = enc.Close() }()
   144  
   145  	return enc.EncodeAll(b, nil)
   146  }