1// Copyright 2018 The Prometheus Authors 2// Licensed under the Apache License, Version 2.0 (the "License"); 3// you may not use this file except in compliance with the License. 4// You may obtain a copy of the License at 5// 6// http://www.apache.org/licenses/LICENSE-2.0 7// 8// Unless required by applicable law or agreed to in writing, software 9// distributed under the License is distributed on an "AS IS" BASIS, 10// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 11// See the License for the specific language governing permissions and 12// limitations under the License. 13 14package encoding 15 16import ( 17 "encoding/binary" 18 "hash" 19 "hash/crc32" 20 "math" 21 "unsafe" 22 23 "github.com/dennwc/varint" 24 "github.com/pkg/errors" 25) 26 27var ( 28 ErrInvalidSize = errors.New("invalid size") 29 ErrInvalidChecksum = errors.New("invalid checksum") 30) 31 32// Encbuf is a helper type to populate a byte slice with various types. 33type Encbuf struct { 34 B []byte 35 C [binary.MaxVarintLen64]byte 36} 37 38func (e *Encbuf) Reset() { e.B = e.B[:0] } 39func (e *Encbuf) Get() []byte { return e.B } 40func (e *Encbuf) Len() int { return len(e.B) } 41 42func (e *Encbuf) PutString(s string) { e.B = append(e.B, s...) } 43func (e *Encbuf) PutByte(c byte) { e.B = append(e.B, c) } 44func (e *Encbuf) PutBytes(b []byte) { e.B = append(e.B, b...) } 45 46func (e *Encbuf) PutBE32int(x int) { e.PutBE32(uint32(x)) } 47func (e *Encbuf) PutUvarint32(x uint32) { e.PutUvarint64(uint64(x)) } 48func (e *Encbuf) PutBE64int64(x int64) { e.PutBE64(uint64(x)) } 49func (e *Encbuf) PutUvarint(x int) { e.PutUvarint64(uint64(x)) } 50 51func (e *Encbuf) PutBE32(x uint32) { 52 binary.BigEndian.PutUint32(e.C[:], x) 53 e.B = append(e.B, e.C[:4]...) 54} 55 56func (e *Encbuf) PutBE64(x uint64) { 57 binary.BigEndian.PutUint64(e.C[:], x) 58 e.B = append(e.B, e.C[:8]...) 59} 60 61func (e *Encbuf) PutBEFloat64(x float64) { 62 e.PutBE64(math.Float64bits(x)) 63} 64 65func (e *Encbuf) PutUvarint64(x uint64) { 66 n := binary.PutUvarint(e.C[:], x) 67 e.B = append(e.B, e.C[:n]...) 68} 69 70func (e *Encbuf) PutVarint64(x int64) { 71 n := binary.PutVarint(e.C[:], x) 72 e.B = append(e.B, e.C[:n]...) 73} 74 75// PutUvarintStr writes a string to the buffer prefixed by its varint length (in bytes!). 76func (e *Encbuf) PutUvarintStr(s string) { 77 b := *(*[]byte)(unsafe.Pointer(&s)) 78 e.PutUvarint(len(b)) 79 e.PutString(s) 80} 81 82// PutUvarintBytes writes a a variable length byte buffer. 83func (e *Encbuf) PutUvarintBytes(b []byte) { 84 e.PutUvarint(len(b)) 85 e.PutBytes(b) 86} 87 88// PutHash appends a hash over the buffers current contents to the buffer. 89func (e *Encbuf) PutHash(h hash.Hash) { 90 h.Reset() 91 e.WriteToHash(h) 92 e.PutHashSum(h) 93} 94 95// WriteToHash writes the current buffer contents to the given hash. 96func (e *Encbuf) WriteToHash(h hash.Hash) { 97 _, err := h.Write(e.B) 98 if err != nil { 99 panic(err) // The CRC32 implementation does not error 100 } 101} 102 103// PutHashSum writes the Sum of the given hash to the buffer. 104func (e *Encbuf) PutHashSum(h hash.Hash) { 105 e.B = h.Sum(e.B) 106} 107 108// Decbuf provides safe methods to extract data from a byte slice. It does all 109// necessary bounds checking and advancing of the byte slice. 110// Several datums can be extracted without checking for errors. However, before using 111// any datum, the err() method must be checked. 112type Decbuf struct { 113 B []byte 114 E error 115} 116 117// NewDecbufAt returns a new decoding buffer. It expects the first 4 bytes 118// after offset to hold the big endian encoded content length, followed by the contents and the expected 119// checksum. 120func NewDecbufAt(bs ByteSlice, off int, castagnoliTable *crc32.Table) Decbuf { 121 if bs.Len() < off+4 { 122 return Decbuf{E: ErrInvalidSize} 123 } 124 b := bs.Range(off, off+4) 125 l := int(binary.BigEndian.Uint32(b)) 126 127 if bs.Len() < off+4+l+4 { 128 return Decbuf{E: ErrInvalidSize} 129 } 130 131 // Load bytes holding the contents plus a CRC32 checksum. 132 b = bs.Range(off+4, off+4+l+4) 133 dec := Decbuf{B: b[:len(b)-4]} 134 135 if castagnoliTable != nil { 136 137 if exp := binary.BigEndian.Uint32(b[len(b)-4:]); dec.Crc32(castagnoliTable) != exp { 138 return Decbuf{E: ErrInvalidChecksum} 139 } 140 } 141 return dec 142} 143 144// NewDecbufUvarintAt returns a new decoding buffer. It expects the first bytes 145// after offset to hold the uvarint-encoded buffers length, followed by the contents and the expected 146// checksum. 147func NewDecbufUvarintAt(bs ByteSlice, off int, castagnoliTable *crc32.Table) Decbuf { 148 // We never have to access this method at the far end of the byte slice. Thus just checking 149 // against the MaxVarintLen32 is sufficient. 150 if bs.Len() < off+binary.MaxVarintLen32 { 151 return Decbuf{E: ErrInvalidSize} 152 } 153 b := bs.Range(off, off+binary.MaxVarintLen32) 154 155 l, n := varint.Uvarint(b) 156 if n <= 0 || n > binary.MaxVarintLen32 { 157 return Decbuf{E: errors.Errorf("invalid uvarint %d", n)} 158 } 159 160 if bs.Len() < off+n+int(l)+4 { 161 return Decbuf{E: ErrInvalidSize} 162 } 163 164 // Load bytes holding the contents plus a CRC32 checksum. 165 b = bs.Range(off+n, off+n+int(l)+4) 166 dec := Decbuf{B: b[:len(b)-4]} 167 168 if dec.Crc32(castagnoliTable) != binary.BigEndian.Uint32(b[len(b)-4:]) { 169 return Decbuf{E: ErrInvalidChecksum} 170 } 171 return dec 172} 173 174// NewDecbufRaw returns a new decoding buffer of the given length. 175func NewDecbufRaw(bs ByteSlice, length int) Decbuf { 176 if bs.Len() < length { 177 return Decbuf{E: ErrInvalidSize} 178 } 179 return Decbuf{B: bs.Range(0, length)} 180} 181 182func (d *Decbuf) Uvarint() int { return int(d.Uvarint64()) } 183func (d *Decbuf) Be32int() int { return int(d.Be32()) } 184func (d *Decbuf) Be64int64() int64 { return int64(d.Be64()) } 185 186// Crc32 returns a CRC32 checksum over the remaining bytes. 187func (d *Decbuf) Crc32(castagnoliTable *crc32.Table) uint32 { 188 return crc32.Checksum(d.B, castagnoliTable) 189} 190 191func (d *Decbuf) Skip(l int) { 192 if len(d.B) < l { 193 d.E = ErrInvalidSize 194 return 195 } 196 d.B = d.B[l:] 197} 198 199func (d *Decbuf) UvarintStr() string { 200 return string(d.UvarintBytes()) 201} 202 203// The return value becomes invalid if the byte slice goes away. 204// Compared to UvarintStr, this avoid allocations. 205func (d *Decbuf) UvarintBytes() []byte { 206 l := d.Uvarint64() 207 if d.E != nil { 208 return []byte{} 209 } 210 if len(d.B) < int(l) { 211 d.E = ErrInvalidSize 212 return []byte{} 213 } 214 s := d.B[:l] 215 d.B = d.B[l:] 216 return s 217} 218 219func (d *Decbuf) Varint64() int64 { 220 if d.E != nil { 221 return 0 222 } 223 // Decode as unsigned first, since that's what the varint library implements. 224 ux, n := varint.Uvarint(d.B) 225 if n < 1 { 226 d.E = ErrInvalidSize 227 return 0 228 } 229 // Now decode "ZigZag encoding" https://developers.google.com/protocol-buffers/docs/encoding#signed_integers. 230 x := int64(ux >> 1) 231 if ux&1 != 0 { 232 x = ^x 233 } 234 d.B = d.B[n:] 235 return x 236} 237 238func (d *Decbuf) Uvarint64() uint64 { 239 if d.E != nil { 240 return 0 241 } 242 x, n := varint.Uvarint(d.B) 243 if n < 1 { 244 d.E = ErrInvalidSize 245 return 0 246 } 247 d.B = d.B[n:] 248 return x 249} 250 251func (d *Decbuf) Be64() uint64 { 252 if d.E != nil { 253 return 0 254 } 255 if len(d.B) < 8 { 256 d.E = ErrInvalidSize 257 return 0 258 } 259 x := binary.BigEndian.Uint64(d.B) 260 d.B = d.B[8:] 261 return x 262} 263 264func (d *Decbuf) Be64Float64() float64 { 265 return math.Float64frombits(d.Be64()) 266} 267 268func (d *Decbuf) Be32() uint32 { 269 if d.E != nil { 270 return 0 271 } 272 if len(d.B) < 4 { 273 d.E = ErrInvalidSize 274 return 0 275 } 276 x := binary.BigEndian.Uint32(d.B) 277 d.B = d.B[4:] 278 return x 279} 280 281func (d *Decbuf) Byte() byte { 282 if d.E != nil { 283 return 0 284 } 285 if len(d.B) < 1 { 286 d.E = ErrInvalidSize 287 return 0 288 } 289 x := d.B[0] 290 d.B = d.B[1:] 291 return x 292} 293 294func (d *Decbuf) ConsumePadding() { 295 if d.E != nil { 296 return 297 } 298 for len(d.B) > 1 && d.B[0] == '\x00' { 299 d.B = d.B[1:] 300 } 301 if len(d.B) < 1 { 302 d.E = ErrInvalidSize 303 } 304} 305 306func (d *Decbuf) Err() error { return d.E } 307func (d *Decbuf) Len() int { return len(d.B) } 308func (d *Decbuf) Get() []byte { return d.B } 309 310// ByteSlice abstracts a byte slice. 311type ByteSlice interface { 312 Len() int 313 Range(start, end int) []byte 314} 315