kesavand | c71914f | 2022-03-25 11:19:03 +0530 | [diff] [blame^] | 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 | } |