feat(thrift): skipdecoder

protocol/thrift/skipdecoder.go

@@ -0,0 +1,162 @@
+package thrift
+
+import (
+	"encoding/binary"
+	"fmt"
+	"io"
+	"sync"
+)
+
+// SkipDecoder scans the underlying io.Reader and returns the bytes of a type
+type SkipDecoder struct {
+	r io.Reader
+
+	p int
+	b []byte
+}
+
+var poolSkipDecoder = sync.Pool{
+	New: func() interface{} {
+		return &SkipDecoder{b: make([]byte, 1024)}
+	},
+}
+
+// NewSkipDecoder ... call Release if no longer use
+func NewSkipDecoder(r io.Reader) *SkipDecoder {
+	p := poolSkipDecoder.Get().(*SkipDecoder)
+	p.Reset(r)
+	return p
+}
+
+// Reset ...
+func (p *SkipDecoder) Reset(r io.Reader) {
+	p.r = r
+	p.p = 0
+}
+
+// Release ...
+func (p *SkipDecoder) Release() {
+	poolSkipDecoder.Put(p)
+}
+
+func (p *SkipDecoder) grow(n int) {
+	if len(p.b)-p.p >= n {
+		return
+	}
+	sz := 2 * cap(p.b)
+	if sz < p.p+n {
+		sz = p.p + n
+	}
+	b := make([]byte, sz)
+	copy(b, p.b[:p.p])
+	p.b = b
+}
+
+// Bytes ...
+func (p *SkipDecoder) Bytes() []byte {
+	return p.b[:p.p]
+}
+
+// Next skips a specific type and returns its bytes
+func (p *SkipDecoder) Next(t TType) (buf []byte, err error) {
+	p.p = 0
+	if err := p.skip(t, defaultRecursionDepth); err != nil {
+		return nil, err
+	}
+	return p.Bytes(), nil
+}
+
+func (p *SkipDecoder) next(n int) ([]byte, error) {
+	p.grow(n)
+	n, err := io.ReadFull(p.r, p.b[p.p:p.p+n])
+	if err != nil {
+		return nil, NewProtocolExceptionWithErr(err)
+	}
+	ret := p.b[p.p : p.p+n]
+	p.p += n
+	return ret, nil
+}
+
+func (p *SkipDecoder) skip(t TType, maxdepth int) error {
+	if maxdepth == 0 {
+		return errDepthLimitExceeded
+	}
+	if sz := typeToSize[t]; sz > 0 {
+		_, err := p.next(int(sz))
+		return err
+	}
+	switch t {
+	case STRING:
+		b, err := p.next(4)
+		if err != nil {
+			return err
+		}
+		sz := int(binary.BigEndian.Uint32(b))
+		if sz < 0 {
+			return errNegativeSize
+		}
+		if _, err := p.next(sz); err != nil {
+			return err
+		}
+	case STRUCT:
+		for {
+			b, err := p.next(1) // TType
+			if err != nil {
+				return err
+			}
+			tp := TType(b[0])
+			if tp == STOP {
+				break
+			}
+			if _, err := p.next(2); err != nil { // Field ID
+				return err
+			}
+			if err := p.skip(tp, maxdepth-1); err != nil {
+				return err
+			}
+		}
+	case MAP:
+		b, err := p.next(6) // 1 byte key TType, 1 byte value TType, 4 bytes Len
+		if err != nil {
+			return err
+		}
+		kt, vt, sz := TType(b[0]), TType(b[1]), int32(binary.BigEndian.Uint32(b[2:]))
+		if sz < 0 {
+			return errNegativeSize
+		}
+		ksz, vsz := int(typeToSize[kt]), int(typeToSize[vt])
+		if ksz > 0 && vsz > 0 {
+			_, err := p.next(int(sz) * (ksz + vsz))
+			return err
+		}
+		for i := int32(0); i < sz; i++ {
+			if err := p.skip(kt, maxdepth-1); err != nil {
+				return err
+			}
+			if err := p.skip(vt, maxdepth-1); err != nil {
+				return err
+			}
+		}
+	case SET, LIST:
+		b, err := p.next(5) // 1 byte value type, 4 bytes Len
+		if err != nil {
+			return err
+		}
+		vt, sz := TType(b[0]), int32(binary.BigEndian.Uint32(b[1:]))
+		if sz < 0 {
+			return errNegativeSize
+		}
+		if vsz := typeToSize[vt]; vsz > 0 {
+			_, err := p.next(int(sz) * int(vsz))
+			return err
+		}
+		for i := int32(0); i < sz; i++ {
+			if err := p.skip(vt, maxdepth-1); err != nil {
+				return err
+			}
+		}
+	default:
+		return NewProtocolException(INVALID_DATA, fmt.Sprintf("unknown data type %d", t))
+	}
+	return nil
+}

protocol/thrift/skipdecoder_test.go

@@ -0,0 +1,141 @@
+package thrift
+
+import (
+	"bytes"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestSkipDecoder(t *testing.T) {
+	x := BinaryProtocol{}
+	// byte
+	b := x.AppendByte([]byte(nil), 1)
+	sz0 := len(b)
+
+	// string
+	b = x.AppendString(b, strings.Repeat("hello", 500)) // larger than buffer
+	sz1 := len(b)
+
+	// list<i32>
+	b = x.AppendListBegin(b, I32, 1)
+	b = x.AppendI32(b, 1)
+	sz2 := len(b)
+
+	// list<string>
+	b = x.AppendListBegin(b, STRING, 1)
+	b = x.AppendString(b, "hello")
+	sz3 := len(b)
+
+	// list<list<i32>>
+	b = x.AppendListBegin(b, LIST, 1)
+	b = x.AppendListBegin(b, I32, 1)
+	b = x.AppendI32(b, 1)
+	sz4 := len(b)
+
+	// map<i32, i64>
+	b = x.AppendMapBegin(b, I32, I64, 1)
+	b = x.AppendI32(b, 1)
+	b = x.AppendI64(b, 2)
+	sz5 := len(b)
+
+	// map<i32, string>
+	b = x.AppendMapBegin(b, I32, STRING, 1)
+	b = x.AppendI32(b, 1)
+	b = x.AppendString(b, "hello")
+	sz6 := len(b)
+
+	// map<string, i64>
+	b = x.AppendMapBegin(b, STRING, I64, 1)
+	b = x.AppendString(b, "hello")
+	b = x.AppendI64(b, 2)
+	sz7 := len(b)
+
+	// map<i32, list<i32>>
+	b = x.AppendMapBegin(b, I32, LIST, 1)
+	b = x.AppendI32(b, 1)
+	b = x.AppendListBegin(b, I32, 1)
+	b = x.AppendI32(b, 1)
+	sz8 := len(b)
+
+	// map<list<i32>, i32>
+	b = x.AppendMapBegin(b, LIST, I32, 1)
+	b = x.AppendListBegin(b, I32, 1)
+	b = x.AppendI32(b, 1)
+	b = x.AppendI32(b, 1)
+	sz9 := len(b)
+
+	// struct i32, list<i32>
+	b = x.AppendFieldBegin(b, I32, 1)
+	b = x.AppendI32(b, 1)
+	b = x.AppendFieldBegin(b, LIST, 1)
+	b = x.AppendListBegin(b, I32, 1)
+	b = x.AppendI32(b, 1)
+	b = x.AppendFieldStop(b)
+	sz10 := len(b)
+
+	r := NewSkipDecoder(bytes.NewReader(b))
+	defer r.Release()
+
+	readn := 0
+	b, err := r.Next(BYTE) // byte
+	require.NoError(t, err)
+	readn += len(b)
+	require.Equal(t, sz0, readn)
+	b, err = r.Next(STRING) // string
+	require.NoError(t, err)
+	readn += len(b)
+	require.Equal(t, sz1, readn)
+	b, err = r.Next(LIST) // list<i32>
+	require.NoError(t, err)
+	readn += len(b)
+	require.Equal(t, sz2, readn)
+	b, err = r.Next(LIST) // list<string>
+	require.NoError(t, err)
+	readn += len(b)
+	require.Equal(t, sz3, readn)
+	b, err = r.Next(LIST) // list<list<i32>>
+	require.NoError(t, err)
+	readn += len(b)
+	require.Equal(t, sz4, readn)
+	b, err = r.Next(MAP) // map<i32, i64>
+	require.NoError(t, err)
+	readn += len(b)
+	require.Equal(t, sz5, readn)
+	b, err = r.Next(MAP) // map<i32, string>
+	require.NoError(t, err)
+	readn += len(b)
+	require.Equal(t, sz6, readn)
+	b, err = r.Next(MAP) // map<string, i64>
+	require.NoError(t, err)
+	readn += len(b)
+	require.Equal(t, sz7, readn)
+	b, err = r.Next(MAP) // map<i32, list<i32>>
+	require.NoError(t, err)
+	readn += len(b)
+	require.Equal(t, sz8, readn)
+	b, err = r.Next(MAP) // map<list<i32>, i32>
+	require.NoError(t, err)
+	readn += len(b)
+	require.Equal(t, sz9, readn)
+	b, err = r.Next(STRUCT) // struct i32, list<i32>
+	require.NoError(t, err)
+	readn += len(b)
+	require.Equal(t, sz10, readn)
+
+	{ // other cases
+		// errDepthLimitExceeded
+		b = b[:0]
+		for i := 0; i < defaultRecursionDepth+1; i++ {
+			b = x.AppendFieldBegin(b, STRUCT, 1)
+		}
+		r := NewSkipDecoder(bytes.NewReader(b))
+		_, err := r.Next(STRUCT)
+		require.Same(t, errDepthLimitExceeded, err)
+
+		// unknown type
+		b, err = r.Next(TType(122))
+		require.Error(t, err)
+	}
+}