/
arg2_kv.go
51 lines (46 loc) · 1.27 KB
/
arg2_kv.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
package thriftarg2test
import (
"fmt"
"testing"
"github.com/stretchr/testify/require"
"github.com/uber/tchannel-go/typed"
)
// BuildKVBuffer builds an thrift Arg2 KV buffer.
func BuildKVBuffer(kv map[string]string) []byte {
// Scan once to know size of buffer
var bufSize int
for k, v := range kv {
// k~2 v~2
bufSize += 2 + len(k) + 2 + len(v)
}
bufSize += 2 // nh:2
buf := make([]byte, bufSize)
wb := typed.NewWriteBuffer(buf)
wb.WriteUint16(uint16(len(kv)))
for k, v := range kv {
wb.WriteLen16String(k)
wb.WriteLen16String(v)
}
return buf[:wb.BytesWritten()]
}
// ReadKVBuffer converts an arg2 buffer to a string map
func ReadKVBuffer(b []byte) (map[string]string, error) {
rbuf := typed.NewReadBuffer(b)
nh := rbuf.ReadUint16()
retMap := make(map[string]string, nh)
for i := uint16(0); i < nh; i++ {
key := rbuf.ReadLen16String()
val := rbuf.ReadLen16String()
retMap[key] = val
}
if rbuf.BytesRemaining() > 0 {
return nil, fmt.Errorf("kv buffer wasn't fully consumed (%d bytes remaining)", rbuf.BytesRemaining())
}
return retMap, nil
}
// MustReadKVBuffer calls require.NoError on the error returned by ReadKVBuffer
func MustReadKVBuffer(tb testing.TB, b []byte) map[string]string {
m, err := ReadKVBuffer(b)
require.NoError(tb, err)
return m
}