1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
|
package compress
import (
"encoding/binary"
"fmt"
"io"
"github.com/go-faster/city"
"github.com/go-faster/errors"
"github.com/klauspost/compress/zstd"
"github.com/pierrec/lz4/v4"
)
// Reader decodes compressed blocks.
type Reader struct {
reader io.Reader
data []byte
pos int64
raw []byte
header []byte
zstd *zstd.Decoder
}
// FormatU128 formats city.U128 as hex.
func FormatU128(v city.U128) string {
var buf [16]byte
binary.LittleEndian.PutUint64(buf[:8], v.Low)
binary.LittleEndian.PutUint64(buf[8:], v.High)
return fmt.Sprintf("%x", buf)
}
// readBlock reads next compressed data into raw and decompresses into data.
func (r *Reader) readBlock() error {
r.pos = 0
_ = r.header[headerSize-1]
if _, err := io.ReadFull(r.reader, r.header); err != nil {
return errors.Wrap(err, "header")
}
var (
rawSize = int(binary.LittleEndian.Uint32(r.header[hRawSize:])) - compressHeaderSize
dataSize = int(binary.LittleEndian.Uint32(r.header[hDataSize:]))
)
if dataSize < 0 || dataSize > maxDataSize {
return errors.Errorf("data size should be %d < %d < %d", 0, dataSize, maxDataSize)
}
if rawSize < 0 || rawSize > maxBlockSize {
return errors.Errorf("raw size should be %d < %d < %d", 0, rawSize, maxBlockSize)
}
r.data = append(r.data[:0], make([]byte, dataSize)...)
r.raw = append(r.raw[:0], r.header...)
r.raw = append(r.raw, make([]byte, rawSize)...)
_ = r.raw[:rawSize+headerSize-1]
if _, err := io.ReadFull(r.reader, r.raw[headerSize:]); err != nil {
return errors.Wrap(err, "read raw")
}
hGot := city.U128{
Low: binary.LittleEndian.Uint64(r.raw[0:8]),
High: binary.LittleEndian.Uint64(r.raw[8:16]),
}
h := city.CH128(r.raw[hMethod:])
if hGot != h {
return errors.Wrap(&CorruptedDataErr{
Actual: h,
Reference: hGot,
RawSize: rawSize,
DataSize: dataSize,
}, "mismatch")
}
switch m := Method(r.header[hMethod]); m {
case LZ4:
n, err := lz4.UncompressBlock(r.raw[headerSize:], r.data)
if err != nil {
return errors.Wrap(err, "uncompress")
}
if n != dataSize {
return errors.Errorf("unexpected uncompressed data size: %d (actual) != %d (got in header)",
n, dataSize,
)
}
case ZSTD:
if r.zstd == nil {
// Lazily initializing to prevent spawning goroutines in NewReader.
// See https://github.com/golang/go/issues/47056#issuecomment-997436820
zstdReader, err := zstd.NewReader(nil,
zstd.WithDecoderConcurrency(1),
zstd.WithDecoderLowmem(true),
)
if err != nil {
return errors.Wrap(err, "zstd")
}
r.zstd = zstdReader
}
data, err := r.zstd.DecodeAll(r.raw[headerSize:], r.data[:0])
if err != nil {
return errors.Wrap(err, "uncompress")
}
if len(data) != dataSize {
return errors.Errorf("unexpected uncompressed data size: %d (actual) != %d (got in header)",
len(data), dataSize,
)
}
r.data = data
case None:
copy(r.data, r.raw[headerSize:])
default:
return errors.Errorf("compression 0x%02x not implemented", m)
}
return nil
}
// Read implements io.Reader.
func (r *Reader) Read(p []byte) (n int, err error) {
if r.pos >= int64(len(r.data)) {
if err := r.readBlock(); err != nil {
return 0, errors.Wrap(err, "read next block")
}
}
n = copy(p, r.data[r.pos:])
r.pos += int64(n)
return n, nil
}
// NewReader returns new *Reader from r.
func NewReader(r io.Reader) *Reader {
return &Reader{
zstd: nil, // lazily initialized
reader: r,
header: make([]byte, headerSize),
}
}
|
