From 06b751940fc28951fa3f042e7f4d35bef6db9aeb Mon Sep 17 00:00:00 2001
From: ChillerDragon <ChillerDragon@gmail.com>
Date: Sun, 23 Apr 2023 11:46:25 +0200
Subject: [PATCH] Fix chunk header packing

---
 tests/chunk_header7_test.py  | 154 ++++++++++++++++++++++++++++++++---
 twnet_parser/chunk_header.py |   4 +-
 2 files changed, 143 insertions(+), 15 deletions(-)

diff --git a/tests/chunk_header7_test.py b/tests/chunk_header7_test.py
index 3b8efac..5217127 100644
--- a/tests/chunk_header7_test.py
+++ b/tests/chunk_header7_test.py
@@ -18,19 +18,147 @@ def cmp_bytes_as_bin(got: bytes, expected: bytes) -> None:
     e_bin = [format(int(x), '#010b') for x in expected]
     assert g_bin == e_bin
 
-# TODO: uncomment when working
-# def test_chunk_header_pack_vital1() -> None:
-#     header: ChunkHeader = ChunkHeader()
-# 
-#     header.flags.resend = False
-#     header.flags.vital = True
-#     header.size = 16
-#     header.seq = 10
-# 
-#     assert bytes([0b01000000, 0b00010000, 0b00001010]) == b'\x40\x10\x0a'
-# 
-#     cmp_bytes_as_bin(header.pack(), b'\x40\x10\x0a')
-#     assert header.pack() == b'\x40\x10\x0a'
+def test_chunk_header_pack_vital1() -> None:
+    header: ChunkHeader = ChunkHeader()
+
+    header.flags.resend = False
+    header.flags.vital = True
+    header.size = 16
+    header.seq = 10
+
+    assert bytes([0b01000000, 0b00010000, 0b00001010]) == b'\x40\x10\x0a'
+
+    cmp_bytes_as_bin(header.pack(), b'\x40\x10\x0a')
+    assert header.pack() == b'\x40\x10\x0a'
+
+def test_chunk_header_pack_resend() -> None:
+    header: ChunkHeader = ChunkHeader()
+
+    header.flags.resend = True
+    header.flags.vital = False
+    header.size = 16
+    header.seq = 0
+
+    assert bytes([0b10000000, 0b00010000]) == b'\x80\x10'
+
+    cmp_bytes_as_bin(header.pack(), b'\x80\x10')
+    assert header.pack() == b'\x80\x10'
+
+def test_chunk_header_pack_no_flags1() -> None:
+    header: ChunkHeader = ChunkHeader()
+
+    header.flags.resend = False
+    header.flags.vital = False
+    header.size = 16
+    header.seq = 0
+
+    assert bytes([0b00000000, 0b00010000]) == b'\x00\x10'
+
+    cmp_bytes_as_bin(header.pack(), b'\x00\x10')
+    assert header.pack() == b'\x00\x10'
+
+def test_chunk_header_pack_no_flags2() -> None:
+    header: ChunkHeader = ChunkHeader()
+
+    header.flags.resend = False
+    header.flags.vital = False
+    header.seq = 0
+
+    header.size = 17
+    assert bytes([0b00000000, 0b00010001]) == b'\x00\x11'
+    cmp_bytes_as_bin(header.pack(), b'\x00\x11')
+    assert header.pack() == b'\x00\x11'
+
+    header.size = 18
+    assert bytes([0b00000000, 0b00010010]) == b'\x00\x12'
+    cmp_bytes_as_bin(header.pack(), b'\x00\x12')
+    assert header.pack() == b'\x00\x12'
+
+    header.size = 19
+    assert bytes([0b00000000, 0b00010011]) == b'\x00\x13'
+    cmp_bytes_as_bin(header.pack(), b'\x00\x13')
+    assert header.pack() == b'\x00\x13'
+
+    header.size = 20
+    assert bytes([0b00000000, 0b00010100]) == b'\x00\x14'
+    cmp_bytes_as_bin(header.pack(), b'\x00\x14')
+    assert header.pack() == b'\x00\x14'
+
+    header.size = 21
+    assert bytes([0b00000000, 0b00010101]) == b'\x00\x15'
+    cmp_bytes_as_bin(header.pack(), b'\x00\x15')
+    assert header.pack() == b'\x00\x15'
+
+def test_chunk_header_pack_no_flags_max_size() -> None:
+    header: ChunkHeader = ChunkHeader()
+
+    header.flags.resend = False
+    header.flags.vital = False
+    header.seq = 0
+
+    header.size = 4094
+    #
+    # from libtw2 docs
+    # https://github.com/heinrich5991/libtw2/blob/7885c99974ee445ce13297b72ae3e7c6ea3b969d/doc/packet7.md
+    #
+    # chunk7_header_nonvital:
+    #     [ 1] flag_resend
+    #     [ 1] flag_vital
+    #     [ 6] <----------
+    #     [ 2] padding   |-- size
+    #     [ 6] <----------
+    #
+    #               FFss ssss    PPss ssss
+    assert bytes([0b0011_1111, 0b0011_1110]) == b'\x3f\x3e'
+    cmp_bytes_as_bin(header.pack(), b'\x3f\x3e')
+    assert header.pack() == b'\x3f\x3e'
+
+    header.size = 4095
+    #
+    # from libtw2 docs
+    # https://github.com/heinrich5991/libtw2/blob/7885c99974ee445ce13297b72ae3e7c6ea3b969d/doc/packet7.md
+    #
+    # chunk7_header_nonvital:
+    #     [ 1] flag_resend
+    #     [ 1] flag_vital
+    #     [ 6] <----------
+    #     [ 2] padding   |-- size
+    #     [ 6] <----------
+    #
+    #               FFss ssss    PPss ssss
+    assert bytes([0b0011_1111, 0b0011_1111]) == b'\x3f\x3f'
+    cmp_bytes_as_bin(header.pack(), b'\x3f\x3f')
+    assert header.pack() == b'\x3f\x3f'
+
+
+def test_chunk_header_pack_no_flags_out_of_bounds_size() -> None:
+    header: ChunkHeader = ChunkHeader()
+
+    header.flags.resend = False
+    header.flags.vital = False
+    header.seq = 0
+
+    header.size = 4096
+    #
+    # from libtw2 docs
+    # https://github.com/heinrich5991/libtw2/blob/7885c99974ee445ce13297b72ae3e7c6ea3b969d/doc/packet7.md
+    #
+    # chunk7_header_nonvital:
+    #     [ 1] flag_resend
+    #     [ 1] flag_vital
+    #     [ 6] <----------
+    #     [ 2] padding   |-- size
+    #     [ 6] <----------
+    #
+    #               FFss ssss    PPss ssss
+    assert bytes([0b0000_0000, 0b0000_0000]) == b'\x00\x00'
+    #                 ^
+    #                 does overflow to zero
+    #                 should we error here instead?
+    #                 TODO: https://gitlab.com/teeworlds-network/twnet_parser/-/issues/2
+    cmp_bytes_as_bin(header.pack(), b'\x00\x00')
+    assert header.pack() == b'\x00\x00'
+
 
 def test_chunk_header_unpack_vital2() -> None:
     parser = ChunkHeaderParser()
diff --git a/twnet_parser/chunk_header.py b/twnet_parser/chunk_header.py
index 144db89..8a7afb2 100644
--- a/twnet_parser/chunk_header.py
+++ b/twnet_parser/chunk_header.py
@@ -28,9 +28,9 @@ class ChunkHeader(PrettyPrint):
     def pack(self) -> bytes:
         flags: int = 0
         if self.flags.resend:
-            flags |= 1
-        if self.flags.vital:
             flags |= 2
+        if self.flags.vital:
+            flags |= 1
         data: bytearray = bytearray([ \
                 ((flags & 0x03) << 6) | \
                 ((self.size >> 6) & 0x3f),