From 2dd887de7d8ceeea6e892c1c4b778488beb8d3b4 Mon Sep 17 00:00:00 2001
From: ChillerDragon <ChillerDragon@gmail.com>
Date: Fri, 7 Jun 2024 15:38:35 +0800
Subject: [PATCH] Add chunk header parsing

---
 chunk/chunk.go         | 36 ++++++++++++++++
 chunk/chunk_test.go    | 25 +++++++++++
 chunk/splitter.go      | 28 +++++++++++++
 chunk/splitter_test.go | 95 ++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 184 insertions(+)
 create mode 100644 chunk/chunk.go
 create mode 100644 chunk/chunk_test.go
 create mode 100644 chunk/splitter.go
 create mode 100644 chunk/splitter_test.go

diff --git a/chunk/chunk.go b/chunk/chunk.go
new file mode 100644
index 0000000..0fa32ee
--- /dev/null
+++ b/chunk/chunk.go
@@ -0,0 +1,36 @@
+package chunk
+
+const (
+	chunkFlagVital = 1
+	chunkFlagResend = 2
+)
+
+type ChunkFlags struct {
+	Vital bool
+	Resend bool
+}
+
+type ChunkHeader struct {
+	Flags ChunkFlags
+	Size  int
+	// sequence number
+	// will be acknowledged in the packet header ack
+	Seq   int
+}
+
+type Chunk struct {
+	Header ChunkHeader
+	Data []byte
+}
+
+func (header *ChunkHeader) Unpack(data []byte) {
+	flagBits := (data[0] >> 6) & 0x03
+	header.Flags.Vital = (flagBits & chunkFlagVital) != 0
+	header.Flags.Resend = (flagBits & chunkFlagResend) != 0
+	header.Size = (int(data[0] & 0x3F) << 6) | (int(data[1]) & 0x3F)
+
+	if header.Flags.Vital {
+		header.Seq = int((data[1] & 0xC0) << 2) | int(data[2])
+	}
+}
+
diff --git a/chunk/chunk_test.go b/chunk/chunk_test.go
new file mode 100644
index 0000000..11fc060
--- /dev/null
+++ b/chunk/chunk_test.go
@@ -0,0 +1,25 @@
+package chunk
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestVitalHeader(t *testing.T) {
+	header := ChunkHeader{}
+	header.Unpack([]byte{0x40, 0x10, 0x0a})
+
+	want := ChunkHeader {
+		Flags: ChunkFlags {
+			Vital: true,
+			Resend: false,
+		},
+		Size: 16,
+		Seq: 10,
+	}
+
+	if !reflect.DeepEqual(header, want) {
+		t.Errorf("got %v, wanted %v", header, want)
+	}
+}
+
diff --git a/chunk/splitter.go b/chunk/splitter.go
new file mode 100644
index 0000000..0373f7f
--- /dev/null
+++ b/chunk/splitter.go
@@ -0,0 +1,28 @@
+package chunk
+
+// data has to be the uncompressed payload of a teeworlds packet
+// without the packet header
+//
+// It will return all the chunks (messages) in that packet
+func UnpackChunks(data []byte) []Chunk {
+	chunks := []Chunk{}
+	payloadSize := len(data)
+	i := 0
+
+	for i < payloadSize {
+		chunk := Chunk{}
+		chunk.Header.Unpack(data[i:])
+		i += 2 // header
+		if chunk.Header.Flags.Vital {
+			i++
+		}
+		end := i + chunk.Header.Size
+		chunk.Data = make([]byte, end - i)
+		copy(chunk.Data[:], data[i:end])
+		i += chunk.Header.Size
+		chunks = append(chunks, chunk)
+	}
+
+	return chunks
+}
+
diff --git a/chunk/splitter_test.go b/chunk/splitter_test.go
new file mode 100644
index 0000000..92ff608
--- /dev/null
+++ b/chunk/splitter_test.go
@@ -0,0 +1,95 @@
+package chunk
+
+import (
+	"fmt"
+	"reflect"
+	"testing"
+)
+
+func TestSplitter(t *testing.T) {
+	// packet payload of real traffic
+	//
+	// Teeworlds 0.7 Protocol packet
+	//     Flags: none (..00 00..)
+	//     Acknowledged sequence number: 3 (.... ..00 0000 0011)
+	//     Number of chunks: 3
+	//     Token: 560baebb
+	//     Payload (80 bytes)
+	// Teeworlds 0.7 Protocol chunk: game.sv_vote_clear_options
+	//	Header (vital: 5)
+	// 	Flags: vital (01.. ....)
+	// 	Size: 1 byte (..00 0000 ..00 0001)
+	// 	Sequence number: 5 (00.. .... 0000 0101)
+	// Teeworlds 0.7 Protocol chunk: game.sv_tune_params
+	//	Header (vital: 6)
+	//	Flags: vital (01.. ....)
+	//	Size: 69 bytes (..00 0001 ..00 0101)
+	//	Sequence number: 6 (00.. .... 0000 0110)
+	// Teeworlds 0.7 Protocol chunk: game.sv_ready_to_enter
+	//	Header (vital: 7)
+	//	Flags: vital (01.. ....)
+	//	Size: 1 byte (..00 0000 ..00 0001)
+	//	Sequence number: 7 (00.. .... 0000 0111)
+	data := []byte{
+		0x40, 0x01, 0x05, 0x16, 0x41, 0x05, 0x06, 0x0c, 0xa8, 0x0f, 0x88, 0x03, 0x32, 0xa8, 0x14, 0xb0,
+		0x12, 0xb4, 0x07, 0x96, 0x02, 0x9f, 0x01, 0xb0, 0xd1, 0x04, 0x80, 0x7d, 0xac, 0x04, 0x9c, 0x17,
+		0x32, 0x98, 0xdb, 0x06, 0x80, 0xb5, 0x18, 0x8c, 0x02, 0xbd, 0x01, 0xa0, 0xed, 0x1a, 0x88, 0x03,
+		0xbd, 0x01, 0xb8, 0xc8, 0x21, 0x90, 0x01, 0x14, 0xbc, 0x0a, 0xa0, 0x9a, 0x0c, 0x88, 0x03, 0x80,
+		0xe2, 0x09, 0x98, 0xea, 0x01, 0xa4, 0x01, 0x00, 0xa4, 0x01, 0xa4, 0x01, 0x40, 0x01, 0x07, 0x10,
+	}
+
+	chunks := UnpackChunks(data)
+
+	{
+		got := len(chunks)
+		want := 3
+
+		if want != got {
+			t.Errorf("got %v, wanted %v", got, want)
+		}
+	}
+
+	{
+		fmt.Printf("%v\n", chunks)
+
+		got := len(chunks[0].Data)
+		want := 1
+
+		if want != got {
+			t.Errorf("got %v, wanted %v", got, want)
+		}
+
+		want = 0x16
+		got = int(chunks[0].Data[0])
+
+		if want != got {
+			t.Errorf("got %v, wanted %v", got, want)
+		}
+	}
+
+	{
+		got := len(chunks[1].Data)
+		want := 69
+
+		if want != got {
+			t.Errorf("got %v, wanted %v", got, want)
+		}
+	}
+
+
+	{
+		want := ChunkHeader {
+			Flags: ChunkFlags {
+				Vital: true,
+				Resend: false,
+			},
+			Size: 1,
+			Seq: 5,
+		}
+
+		if !reflect.DeepEqual(chunks[0].Header, want) {
+			t.Errorf("got %v, wanted %v", chunks[0].Header, want)
+		}
+	}
+}
+