1
0
mirror of https://github.com/astaxie/beego.git synced 2024-11-14 22:50:55 +00:00
Beego/vendor/github.com/siddontang/rdb/reader.go
2018-11-09 12:37:28 +08:00

333 lines
6.4 KiB
Go

// Copyright 2014 Wandoujia Inc. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.
package rdb
import (
"bytes"
"encoding/binary"
"fmt"
"io"
"math"
"strconv"
)
const (
Version = 6
)
const (
rdbTypeString = 0
rdbTypeList = 1
rdbTypeSet = 2
rdbTypeZSet = 3
rdbTypeHash = 4
rdbTypeHashZipmap = 9
rdbTypeListZiplist = 10
rdbTypeSetIntset = 11
rdbTypeZSetZiplist = 12
rdbTypeHashZiplist = 13
rdbFlagExpiryMS = 0xfc
rdbFlagExpiry = 0xfd
rdbFlagSelectDB = 0xfe
rdbFlagEOF = 0xff
)
const (
rdb6bitLen = 0
rdb14bitLen = 1
rdb32bitLen = 2
rdbEncVal = 3
rdbEncInt8 = 0
rdbEncInt16 = 1
rdbEncInt32 = 2
rdbEncLZF = 3
rdbZiplist6bitlenString = 0
rdbZiplist14bitlenString = 1
rdbZiplist32bitlenString = 2
rdbZiplistInt16 = 0xc0
rdbZiplistInt32 = 0xd0
rdbZiplistInt64 = 0xe0
rdbZiplistInt24 = 0xf0
rdbZiplistInt8 = 0xfe
rdbZiplistInt4 = 15
)
type rdbReader struct {
raw io.Reader
buf [8]byte
nread int64
}
func newRdbReader(r io.Reader) *rdbReader {
return &rdbReader{raw: r}
}
func (r *rdbReader) Read(p []byte) (int, error) {
n, err := r.raw.Read(p)
r.nread += int64(n)
return n, err
}
func (r *rdbReader) offset() int64 {
return r.nread
}
func (r *rdbReader) readObject(otype byte) ([]byte, error) {
var b bytes.Buffer
r = newRdbReader(io.TeeReader(r, &b))
switch otype {
default:
return nil, fmt.Errorf("unknown object-type %02x", otype)
case rdbTypeHashZipmap:
fallthrough
case rdbTypeListZiplist:
fallthrough
case rdbTypeSetIntset:
fallthrough
case rdbTypeZSetZiplist:
fallthrough
case rdbTypeHashZiplist:
fallthrough
case rdbTypeString:
if _, err := r.readString(); err != nil {
return nil, err
}
case rdbTypeList, rdbTypeSet:
if n, err := r.readLength(); err != nil {
return nil, err
} else {
for i := 0; i < int(n); i++ {
if _, err := r.readString(); err != nil {
return nil, err
}
}
}
case rdbTypeZSet:
if n, err := r.readLength(); err != nil {
return nil, err
} else {
for i := 0; i < int(n); i++ {
if _, err := r.readString(); err != nil {
return nil, err
}
if _, err := r.readFloat(); err != nil {
return nil, err
}
}
}
case rdbTypeHash:
if n, err := r.readLength(); err != nil {
return nil, err
} else {
for i := 0; i < int(n); i++ {
if _, err := r.readString(); err != nil {
return nil, err
}
if _, err := r.readString(); err != nil {
return nil, err
}
}
}
}
return b.Bytes(), nil
}
func (r *rdbReader) readString() ([]byte, error) {
length, encoded, err := r.readEncodedLength()
if err != nil {
return nil, err
}
if !encoded {
return r.readBytes(int(length))
}
switch t := uint8(length); t {
default:
return nil, fmt.Errorf("invalid encoded-string %02x", t)
case rdbEncInt8:
i, err := r.readInt8()
return []byte(strconv.FormatInt(int64(i), 10)), err
case rdbEncInt16:
i, err := r.readInt16()
return []byte(strconv.FormatInt(int64(i), 10)), err
case rdbEncInt32:
i, err := r.readInt32()
return []byte(strconv.FormatInt(int64(i), 10)), err
case rdbEncLZF:
var inlen, outlen uint32
if inlen, err = r.readLength(); err != nil {
return nil, err
}
if outlen, err = r.readLength(); err != nil {
return nil, err
}
if in, err := r.readBytes(int(inlen)); err != nil {
return nil, err
} else {
return lzfDecompress(in, int(outlen))
}
}
}
func (r *rdbReader) readEncodedLength() (length uint32, encoded bool, err error) {
var u uint8
if u, err = r.readUint8(); err != nil {
return
}
length = uint32(u & 0x3f)
switch u >> 6 {
case rdb6bitLen:
case rdb14bitLen:
u, err = r.readUint8()
length = (length << 8) + uint32(u)
case rdbEncVal:
encoded = true
default:
length, err = r.readUint32BigEndian()
}
return
}
func (r *rdbReader) readLength() (uint32, error) {
length, encoded, err := r.readEncodedLength()
if err == nil && encoded {
err = fmt.Errorf("encoded-length")
}
return length, err
}
func (r *rdbReader) readFloat() (float64, error) {
u, err := r.readUint8()
if err != nil {
return 0, err
}
switch u {
case 253:
return math.NaN(), nil
case 254:
return math.Inf(0), nil
case 255:
return math.Inf(-1), nil
default:
if b, err := r.readBytes(int(u)); err != nil {
return 0, err
} else {
v, err := strconv.ParseFloat(string(b), 64)
return v, err
}
}
}
func (r *rdbReader) readByte() (byte, error) {
b := r.buf[:1]
_, err := r.Read(b)
return b[0], err
}
func (r *rdbReader) readFull(p []byte) error {
_, err := io.ReadFull(r, p)
return err
}
func (r *rdbReader) readBytes(n int) ([]byte, error) {
p := make([]byte, n)
return p, r.readFull(p)
}
func (r *rdbReader) readUint8() (uint8, error) {
b, err := r.readByte()
return uint8(b), err
}
func (r *rdbReader) readUint16() (uint16, error) {
b := r.buf[:2]
err := r.readFull(b)
return binary.LittleEndian.Uint16(b), err
}
func (r *rdbReader) readUint32() (uint32, error) {
b := r.buf[:4]
err := r.readFull(b)
return binary.LittleEndian.Uint32(b), err
}
func (r *rdbReader) readUint64() (uint64, error) {
b := r.buf[:8]
err := r.readFull(b)
return binary.LittleEndian.Uint64(b), err
}
func (r *rdbReader) readUint32BigEndian() (uint32, error) {
b := r.buf[:4]
err := r.readFull(b)
return binary.BigEndian.Uint32(b), err
}
func (r *rdbReader) readInt8() (int8, error) {
u, err := r.readUint8()
return int8(u), err
}
func (r *rdbReader) readInt16() (int16, error) {
u, err := r.readUint16()
return int16(u), err
}
func (r *rdbReader) readInt32() (int32, error) {
u, err := r.readUint32()
return int32(u), err
}
func (r *rdbReader) readInt64() (int64, error) {
u, err := r.readUint64()
return int64(u), err
}
func (r *rdbReader) readInt32BigEndian() (int32, error) {
u, err := r.readUint32BigEndian()
return int32(u), err
}
func lzfDecompress(in []byte, outlen int) (out []byte, err error) {
defer func() {
if x := recover(); x != nil {
err = fmt.Errorf("decompress exception: %v", x)
}
}()
out = make([]byte, outlen)
i, o := 0, 0
for i < len(in) {
ctrl := int(in[i])
i++
if ctrl < 32 {
for x := 0; x <= ctrl; x++ {
out[o] = in[i]
i++
o++
}
} else {
length := ctrl >> 5
if length == 7 {
length = length + int(in[i])
i++
}
ref := o - ((ctrl & 0x1f) << 8) - int(in[i]) - 1
i++
for x := 0; x <= length+1; x++ {
out[o] = out[ref]
ref++
o++
}
}
}
if o != outlen {
return nil, fmt.Errorf("decompress length is %d != expected %d", o, outlen)
}
return out, nil
}