github.com/hamba/avro@v1.8.0/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/ioutil"
    11  
    12  	"github.com/golang/snappy"
    13  )
    14  
    15  // CodecName represents a compression codec name.
    16  type CodecName string
    17  
    18  // Supported compression codecs.
    19  const (
    20  	Null    CodecName = "null"
    21  	Deflate CodecName = "deflate"
    22  	Snappy  CodecName = "snappy"
    23  )
    24  
    25  func resolveCodec(name CodecName, lvl int) (Codec, error) {
    26  	switch name {
    27  	case Null, "":
    28  		return &NullCodec{}, nil
    29  
    30  	case Deflate:
    31  		return &DeflateCodec{compLvl: lvl}, nil
    32  
    33  	case Snappy:
    34  		return &SnappyCodec{}, nil
    35  
    36  	default:
    37  		return nil, fmt.Errorf("unknown codec %s", name)
    38  	}
    39  }
    40  
    41  // Codec represents a compression codec.
    42  type Codec interface {
    43  	// Decode decodes the given bytes.
    44  	Decode([]byte) ([]byte, error)
    45  	// Encode encodes the given bytes.
    46  	Encode([]byte) []byte
    47  }
    48  
    49  // NullCodec is a no op codec.
    50  type NullCodec struct{}
    51  
    52  // Decode decodes the given bytes.
    53  func (*NullCodec) Decode(b []byte) ([]byte, error) {
    54  	return b, nil
    55  }
    56  
    57  // Encode encodes the given bytes.
    58  func (*NullCodec) Encode(b []byte) []byte {
    59  	return b
    60  }
    61  
    62  // DeflateCodec is a flate compression codec.
    63  type DeflateCodec struct {
    64  	compLvl int
    65  }
    66  
    67  // Decode decodes the given bytes.
    68  func (c *DeflateCodec) Decode(b []byte) ([]byte, error) {
    69  	r := flate.NewReader(bytes.NewBuffer(b))
    70  	data, err := ioutil.ReadAll(r)
    71  	if err != nil {
    72  		_ = r.Close()
    73  		return nil, err
    74  	}
    75  	_ = r.Close()
    76  
    77  	return data, nil
    78  }
    79  
    80  // Encode encodes the given bytes.
    81  func (c *DeflateCodec) Encode(b []byte) []byte {
    82  	data := bytes.NewBuffer(make([]byte, 0, len(b)))
    83  
    84  	w, _ := flate.NewWriter(data, c.compLvl)
    85  	_, _ = w.Write(b)
    86  	_ = w.Close()
    87  
    88  	return data.Bytes()
    89  }
    90  
    91  // SnappyCodec is a snappy compression codec.
    92  type SnappyCodec struct{}
    93  
    94  // Decode decodes the given bytes.
    95  func (*SnappyCodec) Decode(b []byte) ([]byte, error) {
    96  	l := len(b)
    97  	if l < 5 {
    98  		return nil, errors.New("block does not contain snappy checksum")
    99  	}
   100  
   101  	dst, err := snappy.Decode(nil, b[:l-4])
   102  	if err != nil {
   103  		return nil, err
   104  	}
   105  
   106  	crc := binary.BigEndian.Uint32(b[l-4:])
   107  	if crc32.ChecksumIEEE(dst) != crc {
   108  		return nil, errors.New("snappy checksum mismatch")
   109  	}
   110  
   111  	return dst, nil
   112  }
   113  
   114  // Encode encodes the given bytes.
   115  func (*SnappyCodec) Encode(b []byte) []byte {
   116  	dst := snappy.Encode(nil, b)
   117  
   118  	dst = append(dst, 0, 0, 0, 0)
   119  	binary.BigEndian.PutUint32(dst[len(dst)-4:], crc32.ChecksumIEEE(b))
   120  
   121  	return dst
   122  }