Add stateful unpacker

This commit is contained in:
ChillerDragon 2024-06-17 13:11:59 +08:00
parent 309062c7d3
commit dcf26f818f
4 changed files with 247 additions and 59 deletions

View file

@ -1,9 +1,113 @@
package packer
import "errors"
import (
"slices"
)
func PackInt(num int) ([]byte, error) {
dstLen := 4
type Unpacker struct {
data []byte
idx int
}
func (u *Unpacker) Reset(data []byte) {
u.data = slices.Clone(data)
u.idx = 0
}
// first byte of the current buffer
func (u *Unpacker) byte() byte {
return u.data[u.idx]
}
// consume one byte
func (u *Unpacker) getByte() byte {
b := u.data[u.idx]
u.idx++
return b
}
const (
Sanitize = 1
SanitizeCC = 2
SanitizeSkipWhitespaces = 4
)
func (u *Unpacker) GetStringSanitized(sanitizeType int) string {
bytes := []byte{}
skipping := sanitizeType&SanitizeSkipWhitespaces != 0
for {
b := u.getByte()
if b == 0x00 {
break
}
if skipping {
if b == ' ' || b == '\t' || b == '\n' {
continue
}
skipping = false
}
if sanitizeType&SanitizeCC != 0 {
if b < 32 {
b = ' '
}
} else if sanitizeType&Sanitize != 0 {
if b < 32 && !(b == '\r') && !(b == '\n') && !(b == '\t') {
b = ' '
}
}
bytes = append(bytes, b)
}
return string(bytes)
}
func (u *Unpacker) GetString() string {
return u.GetStringSanitized(Sanitize)
}
func (u *Unpacker) GetInt() int {
sign := int(u.byte()>>6) & 1
res := int(u.byte() & 0x3F)
// fake loop should only loop once
// its the poor mans goto hack
for {
if (u.byte() & 0x80) == 0 {
break
}
u.idx += 1
res |= int(u.byte()&0x7F) << 6
if (u.byte() & 0x80) == 0 {
break
}
u.idx += 1
res |= int(u.byte()&0x7F) << (6 + 7)
if (u.byte() & 0x80) == 0 {
break
}
u.idx += 1
res |= int(u.byte()&0x7F) << (6 + 7 + 7)
if (u.byte() & 0x80) == 0 {
break
}
u.idx += 1
res |= int(u.byte()&0x7F) << (6 + 7 + 7 + 7)
break
}
u.idx += 1
res ^= -sign
return res
}
func PackInt(num int) []byte {
res := []byte{0x00}
idx := 0
if num < 0 {
@ -14,10 +118,6 @@ func PackInt(num int) ([]byte, error) {
res[0] |= byte(num & 0x3F) // pack 6 bit into dst
num >>= 6 // discard 6 bits
for num != 0 {
if idx > dstLen {
return nil, errors.New("Int too big")
}
res = append(res, 0x00)
res[idx] |= 0x80 // set extend bit
@ -26,7 +126,7 @@ func PackInt(num int) ([]byte, error) {
num >>= 7 // discard 7 bits
}
return res, nil
return res
}
func UnpackInt(data []byte) int {

121
packer/packer_state_test.go Normal file
View file

@ -0,0 +1,121 @@
package packer
import (
"reflect"
"testing"
)
// unpack with state
func TestUnpackSimpleInts(t *testing.T) {
u := Unpacker{}
u.Reset([]byte{0x01, 0x02, 0x03, 0x0f})
got := u.GetInt()
want := 1
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
got = u.GetInt()
want = 2
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
got = u.GetInt()
want = 3
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
got = u.GetInt()
want = 15
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
}
func TestUnpackString(t *testing.T) {
u := Unpacker{}
u.Reset([]byte{'f', 'o', 'o', 0x00})
got := u.GetString()
want := "foo"
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
}
func TestUnpackTwoStrings(t *testing.T) {
u := Unpacker{}
u.Reset([]byte{'f', 'o', 'o', 0x00, 'b', 'a', 'r', 0x00})
got := u.GetString()
want := "foo"
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
got = u.GetString()
want = "bar"
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
}
func TestUnpackMixed(t *testing.T) {
u := Unpacker{}
u.Reset([]byte{0x0F, 0x0F, 'f', 'o', 'o', 0x00, 'b', 'a', 'r', 0x00, 0x01})
// ints
{
got := u.GetInt()
want := 15
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
got = u.GetInt()
want = 15
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
}
// strings
{
got := u.GetString()
want := "foo"
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
got = u.GetString()
want = "bar"
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
}
// ints
{
got := u.GetInt()
want := 1
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
}
}

View file

@ -8,107 +8,71 @@ import (
// pack
func TestPackSmallPositiveInts(t *testing.T) {
got, err := PackInt(1)
got := PackInt(1)
want := []byte{0x01}
if err != nil {
t.Errorf("Unexpected error: %v", err)
}
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
}
func TestPackMultiBytePositiveInts(t *testing.T) {
got, err := PackInt(63)
got := PackInt(63)
want := []byte{0x3F}
if err != nil {
t.Errorf("Unexpected error: %v", err)
}
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
got, err = PackInt(64)
got = PackInt(64)
want = []byte{0x80, 0x01}
if err != nil {
t.Errorf("Unexpected error: %v", err)
}
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
got, err = PackInt(65)
got = PackInt(65)
want = []byte{0x81, 0x01}
if err != nil {
t.Errorf("Unexpected error: %v", err)
}
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
}
func TestPackSmallNegativeInts(t *testing.T) {
got, err := PackInt(-1)
got := PackInt(-1)
want := []byte{0x40}
if err != nil {
t.Errorf("Unexpected error: %v", err)
}
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
got, err = PackInt(-2)
got = PackInt(-2)
want = []byte{0x41}
if err != nil {
t.Errorf("Unexpected error: %v", err)
}
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
}
func TestPackMultiByteNegativeInts(t *testing.T) {
got, err := PackInt(-63)
got := PackInt(-63)
want := []byte{0x7E}
if err != nil {
t.Errorf("Unexpected error: %v", err)
}
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
got, err = PackInt(-64)
got = PackInt(-64)
want = []byte{0x7F}
if err != nil {
t.Errorf("Unexpected error: %v", err)
}
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}
got, err = PackInt(-65)
got = PackInt(-65)
want = []byte{0xC0, 0x01}
if err != nil {
t.Errorf("Unexpected error: %v", err)
}
if !reflect.DeepEqual(got, want) {
t.Errorf("got %v, wanted %v", got, want)
}

View file

@ -11,6 +11,7 @@ import (
"github.com/teeworlds-go/huffman"
"github.com/teeworlds-go/teeworlds/chunk"
"github.com/teeworlds-go/teeworlds/packer"
"github.com/teeworlds-go/teeworlds/packet"
)
@ -168,7 +169,7 @@ func byteSliceToString(s []byte) string {
return string(s)
}
func (client *TeeworldsClient) onSystemMsg(msg int, chunk chunk.Chunk) {
func (client *TeeworldsClient) onSystemMsg(msg int, chunk chunk.Chunk, u *packer.Unpacker) {
if msg == msgSysMapChange {
fmt.Println("got map change")
client.sendReady()
@ -180,7 +181,7 @@ func (client *TeeworldsClient) onSystemMsg(msg int, chunk chunk.Chunk) {
}
}
func (client *TeeworldsClient) onGameMsg(msg int, chunk chunk.Chunk) {
func (client *TeeworldsClient) onGameMsg(msg int, chunk chunk.Chunk, u *packer.Unpacker) {
if msg == msgGameReadyToEnter {
fmt.Println("got ready to enter")
client.sendEnterGame()
@ -192,16 +193,18 @@ func (client *TeeworldsClient) onGameMsg(msg int, chunk chunk.Chunk) {
func (client *TeeworldsClient) onMessage(chunk chunk.Chunk) {
fmt.Printf("got chunk size=%d data=%v\n", chunk.Header.Size, chunk.Data)
// TODO: we need to int unpack this because it can be 2 bytes long
msg := int(chunk.Data[0])
u := packer.Unpacker{}
u.Reset(chunk.Data)
msg := u.GetInt()
sys := msg&1 != 0
msg >>= 1
if sys {
client.onSystemMsg(msg, chunk)
client.onSystemMsg(msg, chunk, &u)
} else {
client.onGameMsg(msg, chunk)
client.onGameMsg(msg, chunk, &u)
}
}