cubefs

1
// Copyright 2020+ Klaus Post. All rights reserved.
2
// License information can be found in the LICENSE file.
3

4
package zstd
5

6
import (
7
	"bytes"
8
	"encoding/binary"
9
	"errors"
10
	"io"
11
)
12

13
// HeaderMaxSize is the maximum size of a Frame and Block Header.
14
// If less is sent to Header.Decode it *may* still contain enough information.
15
const HeaderMaxSize = 14 + 3
16

17
// Header contains information about the first frame and block within that.
18
type Header struct {
19
	// SingleSegment specifies whether the data is to be decompressed into a
20
	// single contiguous memory segment.
21
	// It implies that WindowSize is invalid and that FrameContentSize is valid.
22
	SingleSegment bool
23

24
	// WindowSize is the window of data to keep while decoding.
25
	// Will only be set if SingleSegment is false.
26
	WindowSize uint64
27

28
	// Dictionary ID.
29
	// If 0, no dictionary.
30
	DictionaryID uint32
31

32
	// HasFCS specifies whether FrameContentSize has a valid value.
33
	HasFCS bool
34

35
	// FrameContentSize is the expected uncompressed size of the entire frame.
36
	FrameContentSize uint64
37

38
	// Skippable will be true if the frame is meant to be skipped.
39
	// This implies that FirstBlock.OK is false.
40
	Skippable bool
41

42
	// SkippableID is the user-specific ID for the skippable frame.
43
	// Valid values are between 0 to 15, inclusive.
44
	SkippableID int
45

46
	// SkippableSize is the length of the user data to skip following
47
	// the header.
48
	SkippableSize uint32
49

50
	// HeaderSize is the raw size of the frame header.
51
	//
52
	// For normal frames, it includes the size of the magic number and
53
	// the size of the header (per section 3.1.1.1).
54
	// It does not include the size for any data blocks (section 3.1.1.2) nor
55
	// the size for the trailing content checksum.
56
	//
57
	// For skippable frames, this counts the size of the magic number
58
	// along with the size of the size field of the payload.
59
	// It does not include the size of the skippable payload itself.
60
	// The total frame size is the HeaderSize plus the SkippableSize.
61
	HeaderSize int
62

63
	// First block information.
64
	FirstBlock struct {
65
		// OK will be set if first block could be decoded.
66
		OK bool
67

68
		// Is this the last block of a frame?
69
		Last bool
70

71
		// Is the data compressed?
72
		// If true CompressedSize will be populated.
73
		// Unfortunately DecompressedSize cannot be determined
74
		// without decoding the blocks.
75
		Compressed bool
76

77
		// DecompressedSize is the expected decompressed size of the block.
78
		// Will be 0 if it cannot be determined.
79
		DecompressedSize int
80

81
		// CompressedSize of the data in the block.
82
		// Does not include the block header.
83
		// Will be equal to DecompressedSize if not Compressed.
84
		CompressedSize int
85
	}
86

87
	// If set there is a checksum present for the block content.
88
	// The checksum field at the end is always 4 bytes long.
89
	HasCheckSum bool
90
}
91

92
// Decode the header from the beginning of the stream.
93
// This will decode the frame header and the first block header if enough bytes are provided.
94
// It is recommended to provide at least HeaderMaxSize bytes.
95
// If the frame header cannot be read an error will be returned.
96
// If there isn't enough input, io.ErrUnexpectedEOF is returned.
97
// The FirstBlock.OK will indicate if enough information was available to decode the first block header.
98
func (h *Header) Decode(in []byte) error {
99
	*h = Header{}
100
	if len(in) < 4 {
101
		return io.ErrUnexpectedEOF
102
	}
103
	h.HeaderSize += 4
104
	b, in := in[:4], in[4:]
105
	if !bytes.Equal(b, frameMagic) {
106
		if !bytes.Equal(b[1:4], skippableFrameMagic) || b[0]&0xf0 != 0x50 {
107
			return ErrMagicMismatch
108
		}
109
		if len(in) < 4 {
110
			return io.ErrUnexpectedEOF
111
		}
112
		h.HeaderSize += 4
113
		h.Skippable = true
114
		h.SkippableID = int(b[0] & 0xf)
115
		h.SkippableSize = binary.LittleEndian.Uint32(in)
116
		return nil
117
	}
118

119
	// Read Window_Descriptor
120
	// https://github.com/facebook/zstd/blob/dev/doc/zstd_compression_format.md#window_descriptor
121
	if len(in) < 1 {
122
		return io.ErrUnexpectedEOF
123
	}
124
	fhd, in := in[0], in[1:]
125
	h.HeaderSize++
126
	h.SingleSegment = fhd&(1<<5) != 0
127
	h.HasCheckSum = fhd&(1<<2) != 0
128
	if fhd&(1<<3) != 0 {
129
		return errors.New("reserved bit set on frame header")
130
	}
131

132
	if !h.SingleSegment {
133
		if len(in) < 1 {
134
			return io.ErrUnexpectedEOF
135
		}
136
		var wd byte
137
		wd, in = in[0], in[1:]
138
		h.HeaderSize++
139
		windowLog := 10 + (wd >> 3)
140
		windowBase := uint64(1) << windowLog
141
		windowAdd := (windowBase / 8) * uint64(wd&0x7)
142
		h.WindowSize = windowBase + windowAdd
143
	}
144

145
	// Read Dictionary_ID
146
	// https://github.com/facebook/zstd/blob/dev/doc/zstd_compression_format.md#dictionary_id
147
	if size := fhd & 3; size != 0 {
148
		if size == 3 {
149
			size = 4
150
		}
151
		if len(in) < int(size) {
152
			return io.ErrUnexpectedEOF
153
		}
154
		b, in = in[:size], in[size:]
155
		h.HeaderSize += int(size)
156
		switch size {
157
		case 1:
158
			h.DictionaryID = uint32(b[0])
159
		case 2:
160
			h.DictionaryID = uint32(b[0]) | (uint32(b[1]) << 8)
161
		case 4:
162
			h.DictionaryID = uint32(b[0]) | (uint32(b[1]) << 8) | (uint32(b[2]) << 16) | (uint32(b[3]) << 24)
163
		}
164
	}
165

166
	// Read Frame_Content_Size
167
	// https://github.com/facebook/zstd/blob/dev/doc/zstd_compression_format.md#frame_content_size
168
	var fcsSize int
169
	v := fhd >> 6
170
	switch v {
171
	case 0:
172
		if h.SingleSegment {
173
			fcsSize = 1
174
		}
175
	default:
176
		fcsSize = 1 << v
177
	}
178

179
	if fcsSize > 0 {
180
		h.HasFCS = true
181
		if len(in) < fcsSize {
182
			return io.ErrUnexpectedEOF
183
		}
184
		b, in = in[:fcsSize], in[fcsSize:]
185
		h.HeaderSize += int(fcsSize)
186
		switch fcsSize {
187
		case 1:
188
			h.FrameContentSize = uint64(b[0])
189
		case 2:
190
			// When FCS_Field_Size is 2, the offset of 256 is added.
191
			h.FrameContentSize = uint64(b[0]) | (uint64(b[1]) << 8) + 256
192
		case 4:
193
			h.FrameContentSize = uint64(b[0]) | (uint64(b[1]) << 8) | (uint64(b[2]) << 16) | (uint64(b[3]) << 24)
194
		case 8:
195
			d1 := uint32(b[0]) | (uint32(b[1]) << 8) | (uint32(b[2]) << 16) | (uint32(b[3]) << 24)
196
			d2 := uint32(b[4]) | (uint32(b[5]) << 8) | (uint32(b[6]) << 16) | (uint32(b[7]) << 24)
197
			h.FrameContentSize = uint64(d1) | (uint64(d2) << 32)
198
		}
199
	}
200

201
	// Frame Header done, we will not fail from now on.
202
	if len(in) < 3 {
203
		return nil
204
	}
205
	tmp := in[:3]
206
	bh := uint32(tmp[0]) | (uint32(tmp[1]) << 8) | (uint32(tmp[2]) << 16)
207
	h.FirstBlock.Last = bh&1 != 0
208
	blockType := blockType((bh >> 1) & 3)
209
	// find size.
210
	cSize := int(bh >> 3)
211
	switch blockType {
212
	case blockTypeReserved:
213
		return nil
214
	case blockTypeRLE:
215
		h.FirstBlock.Compressed = true
216
		h.FirstBlock.DecompressedSize = cSize
217
		h.FirstBlock.CompressedSize = 1
218
	case blockTypeCompressed:
219
		h.FirstBlock.Compressed = true
220
		h.FirstBlock.CompressedSize = cSize
221
	case blockTypeRaw:
222
		h.FirstBlock.DecompressedSize = cSize
223
		h.FirstBlock.CompressedSize = cSize
224
	default:
225
		panic("Invalid block type")
226
	}
227

228
	h.FirstBlock.OK = true
229
	return nil
230
}
231