Просмотр исходного кода

fix: use verified cq2 header semantics

wuzilong 1 месяц назад
Родитель
Сommit
5bb8a3415f

+ 7 - 7
cq2_sdk/src/cq2_sdk/transport/framing.py

@@ -19,12 +19,12 @@ HEADER_SIZE = WIRE_HEADER_SIZE
 
 
 class FrameHeader(BaseModel):
-    pre: int
+    param3: int
     msg_id: int
     recog: int
     param1: int
     param2: int
-    param3: int
+    sdata: int
     codec: int
     body_len: int
     sequence: int
@@ -38,7 +38,7 @@ def encode_frame(
     param3: int,
     body: bytes,
     *,
-    pre: int = 0,
+    sdata: int = 0,
     codec: int = 0,
     sequence: int = 0,
 ) -> bytes:
@@ -47,12 +47,12 @@ def encode_frame(
         raise ValueError(f"body too large: {len(body)} bytes")
     plain_header = struct.pack(
         PLAIN_HEADER_FORMAT,
-        pre,
+        param3,
         msg_id,
         recog,
         param1,
         param2,
-        param3,
+        sdata,
         codec,
         len(body),
         sequence,
@@ -69,12 +69,12 @@ def decode_header(raw: bytes) -> FrameHeader:
         decode6(raw[:WIRE_HEADER_SIZE], expected_size=PLAIN_HEADER_SIZE),
     )
     return FrameHeader(
-        pre=fields[0],
+        param3=fields[0],
         msg_id=fields[1],
         recog=fields[2],
         param1=fields[3],
         param2=fields[4],
-        param3=fields[5],
+        sdata=fields[5],
         codec=fields[6],
         body_len=fields[7],
         sequence=fields[8],

+ 31 - 1
cq2_sdk/src/cq2_sdk/transport/tcp.py

@@ -11,6 +11,7 @@ via `await recv(timeout=...)`.
 """
 from __future__ import annotations
 import asyncio
+import time
 from typing import Optional
 from .framing import WIRE_HEADER_SIZE, FrameHeader, decode_frame, decode_header, encode_frame
 
@@ -26,6 +27,7 @@ class TcpClient:
         self.host = host
         self.port = port
         self._next_sequence = initial_sequence
+        self._sdata = 0
         self._reader: Optional[asyncio.StreamReader] = None
         self._writer: Optional[asyncio.StreamWriter] = None
         self._buffer = bytearray()
@@ -36,6 +38,10 @@ class TcpClient:
         self._closed = False
 
     async def connect(self) -> None:
+        if self._writer is not None and not self._writer.is_closing():
+            raise RuntimeError("already connected")
+        self._closed = False
+        self._recv_error = None
         self._reader, self._writer = await asyncio.open_connection(self.host, self.port)
         self._recv_task = asyncio.create_task(self._recv_loop())
 
@@ -44,18 +50,41 @@ class TcpClient:
             return
         self._closed = True
         await self.stop_heartbeat()
+        await self._disconnect_transport(reset_sequence=True)
+
+    async def _disconnect_transport(self, *, reset_sequence: bool) -> None:
         if self._recv_task and not self._recv_task.done():
             self._recv_task.cancel()
             try:
                 await self._recv_task
             except (asyncio.CancelledError, Exception):
                 pass
+        self._recv_task = None
         if self._writer:
             try:
                 self._writer.close()
                 await self._writer.wait_closed()
             except Exception:
                 pass
+        self._reader = None
+        self._writer = None
+        self._buffer.clear()
+        self._msg_queue = asyncio.Queue()
+        self._recv_error = None
+        if reset_sequence:
+            self._next_sequence = 0
+
+    async def reconnect(self, host: str, port: int) -> None:
+        """Move to another CQ2 endpoint while retaining session data."""
+        await self.stop_heartbeat()
+        await self._disconnect_transport(reset_sequence=True)
+        self.host = host
+        self.port = port
+        await self.connect()
+
+    def set_sdata(self, value: int) -> None:
+        """Set the session value emitted in decoded header word five."""
+        self._sdata = value
 
     async def _recv_loop(self) -> None:
         assert self._reader is not None
@@ -100,6 +129,7 @@ class TcpClient:
             param2,
             param3,
             body,
+            sdata=self._sdata,
             sequence=self._next_sequence,
         )
         self._next_sequence += 1
@@ -124,7 +154,7 @@ class TcpClient:
                     if self._closed:
                         return
                     try:
-                        await self.send(60000, 0, 0, 0, 0, b"")
+                        await self.send(60000, int(time.time()), 0, 0, 0, b"")
                     except Exception:
                         return
             except asyncio.CancelledError:

+ 18 - 4
cq2_sdk/tests/unit/test_framing.py

@@ -29,6 +29,20 @@ def test_decode6_restores_captured_heartbeat_header():
     assert decode6(HEARTBEAT_WIRE_HEADER, expected_size=36) == HEARTBEAT_HEADER
 
 
+def test_decode_captured_heartbeat_uses_verified_header_field_order():
+    header = decode_header(HEARTBEAT_WIRE_HEADER)
+
+    assert header.param3 == 0
+    assert header.msg_id == C.CS_60000
+    assert header.recog == 1789709744
+    assert header.param1 == 0
+    assert header.param2 == 0
+    assert header.sdata == 149
+    assert header.codec == 0
+    assert header.body_len == 0
+    assert header.sequence == 112
+
+
 def test_decode6_rejects_non_wire_character():
     with pytest.raises(ValueError, match="outside 6-bit wire range"):
         decode6(b"<~")
@@ -49,19 +63,19 @@ def test_encode_decode_roundtrip():
         recog=900000,
         param1=12345,
         param2=0,
-        param3=0,
+        param3=7,
         body=body,
-        pre=7,
+        sdata=11,
         sequence=42,
     )
     assert len(raw) == WIRE_HEADER_SIZE + len(body)
     hdr, decoded_body = decode_frame(raw)
-    assert hdr.pre == 7
+    assert hdr.param3 == 7
     assert hdr.msg_id == C.CS_2001
     assert hdr.recog == 900000
     assert hdr.param1 == 12345
     assert hdr.param2 == 0
-    assert hdr.param3 == 0
+    assert hdr.sdata == 11
     assert hdr.body_len == 5
     assert hdr.codec == 0
     assert hdr.sequence == 42

+ 3 - 3
cq2_sdk/tests/unit/test_message_registry.py

@@ -109,9 +109,9 @@ async def test_wait_for_sc_returns_matching():
         def __init__(self):
             self.queue = asyncio.Queue()
             # Queue some frames: 999 (drop), 55 (match), 100 (drop)
-            self.queue.put_nowait((FrameHeader(pre=0, msg_id=999, recog=0, param1=0, param2=0, param3=0, codec=0, body_len=0, sequence=0), b""))
-            self.queue.put_nowait((FrameHeader(pre=0, msg_id=55, recog=0, param1=0, param2=0, param3=0, codec=0, body_len=0, sequence=1), b"matched"))
-            self.queue.put_nowait((FrameHeader(pre=0, msg_id=100, recog=0, param1=0, param2=0, param3=0, codec=0, body_len=0, sequence=2), b""))
+            self.queue.put_nowait((FrameHeader(param3=0, msg_id=999, recog=0, param1=0, param2=0, sdata=0, codec=0, body_len=0, sequence=0), b""))
+            self.queue.put_nowait((FrameHeader(param3=0, msg_id=55, recog=0, param1=0, param2=0, sdata=0, codec=0, body_len=0, sequence=1), b"matched"))
+            self.queue.put_nowait((FrameHeader(param3=0, msg_id=100, recog=0, param1=0, param2=0, sdata=0, codec=0, body_len=0, sequence=2), b""))
 
         async def recv(self, timeout=30.0):
             return await self.queue.get()

+ 92 - 0
cq2_sdk/tests/unit/test_tcp_client.py

@@ -1,5 +1,6 @@
 import asyncio
 import pytest
+import cq2_sdk.transport.tcp as tcp_module
 from cq2_sdk.transport.tcp import TcpClient
 from cq2_sdk.transport.framing import WIRE_HEADER_SIZE, decode_header, encode_frame
 
@@ -131,6 +132,97 @@ async def test_send_uses_and_increments_initial_sequence():
         await server.wait_closed()
 
 
+@pytest.mark.asyncio
+async def test_send_uses_current_sdata():
+    received = asyncio.Future()
+
+    async def handler(reader, writer):
+        header = decode_header(await reader.readexactly(WIRE_HEADER_SIZE))
+        if header.body_len:
+            await reader.readexactly(header.body_len)
+        received.set_result(header)
+        writer.close()
+
+    server = await asyncio.start_server(handler, "127.0.0.1", 0)
+    port = server.sockets[0].getsockname()[1]
+    client = TcpClient("127.0.0.1", port)
+    client.set_sdata(149)
+    await client.connect()
+    try:
+        await client.send(2001, 900000, 1, 2, 3, b"login")
+        header = await asyncio.wait_for(received, timeout=1.0)
+        assert header.param3 == 3
+        assert header.sdata == 149
+    finally:
+        await client.close()
+        server.close()
+        await server.wait_closed()
+
+
+@pytest.mark.asyncio
+async def test_reconnect_resets_sequence_and_preserves_sdata():
+    received: list[tuple[int, int]] = []
+    complete = asyncio.Event()
+
+    async def handler(reader, writer):
+        header = decode_header(await reader.readexactly(WIRE_HEADER_SIZE))
+        if header.body_len:
+            await reader.readexactly(header.body_len)
+        received.append((header.sequence, header.sdata))
+        if len(received) == 2:
+            complete.set()
+        writer.close()
+        await writer.wait_closed()
+
+    server = await asyncio.start_server(handler, "127.0.0.1", 0)
+    port = server.sockets[0].getsockname()[1]
+    client = TcpClient("127.0.0.1", port)
+    client.set_sdata(149)
+    await client.connect()
+    try:
+        await client.send(100, 0, 0, 0, 0, b"role")
+        await client.reconnect("127.0.0.1", port)
+        await client.send(20, 0, 0, 0, 0, b"world")
+        await asyncio.wait_for(complete.wait(), timeout=1.0)
+        assert received == [(0, 149), (0, 149)]
+    finally:
+        await client.close()
+        server.close()
+        await server.wait_closed()
+
+
+@pytest.mark.asyncio
+async def test_heartbeat_uses_current_unix_time_in_recog(monkeypatch):
+    received = asyncio.Future()
+
+    async def handler(reader, writer):
+        header = decode_header(await reader.readexactly(WIRE_HEADER_SIZE))
+        received.set_result(header)
+        writer.close()
+
+    class FixedTime:
+        @staticmethod
+        def time():
+            return 1_789_709_744
+
+    monkeypatch.setattr(tcp_module, "time", FixedTime, raising=False)
+    server = await asyncio.start_server(handler, "127.0.0.1", 0)
+    port = server.sockets[0].getsockname()[1]
+    client = TcpClient("127.0.0.1", port)
+    client.set_sdata(149)
+    await client.connect()
+    try:
+        await client.start_heartbeat(interval=0.01)
+        header = await asyncio.wait_for(received, timeout=1.0)
+        assert header.msg_id == 60000
+        assert header.recog == 1_789_709_744
+        assert header.sdata == 149
+    finally:
+        await client.close()
+        server.close()
+        await server.wait_closed()
+
+
 @pytest.mark.asyncio
 async def test_recv_surfaces_unsupported_body_codec():
     async def handler(reader, writer):