From b533b3a03a6d630a013a7124cf6c52535685a0b0 Mon Sep 17 00:00:00 2001
From: MacOMNI <414294494@qq.com>
Date: Thu, 24 Oct 2024 16:48:21 +0800
Subject: [PATCH] update peer test

---
 Networking/Sources/Networking/Connection.swift   | 14 +++++++++++++-
 Networking/Tests/NetworkingTests/PeerTests.swift |  4 ++--
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/Networking/Sources/Networking/Connection.swift b/Networking/Sources/Networking/Connection.swift
index 169990a4..1d1473ef 100644
--- a/Networking/Sources/Networking/Connection.swift
+++ b/Networking/Sources/Networking/Connection.swift
@@ -12,6 +12,10 @@ public protocol ConnectionInfoProtocol {
     var remoteAddress: NetAddr { get }
 }
 
+enum ConnectionError: Error {
+    case receiveFailed
+}
+
 public final class Connection<Handler: StreamHandler>: Sendable, ConnectionInfoProtocol {
     let connection: QuicConnection
     let impl: PeerImpl<Handler>
@@ -49,7 +53,15 @@ public final class Connection<Handler: StreamHandler>: Sendable, ConnectionInfoP
             response.append(nextData)
             break
         }
-        return response
+        guard response.count >= 4 else {
+            stream.close(abort: true)
+            throw ConnectionError.receiveFailed
+        }
+        let lengthData = response.prefix(4)
+        let length = UInt32(
+            littleEndian: lengthData.withUnsafeBytes { $0.loadUnaligned(as: UInt32.self) }
+        )
+        return response.dropFirst(4).prefix(Int(length))
     }
 
     @discardableResult
diff --git a/Networking/Tests/NetworkingTests/PeerTests.swift b/Networking/Tests/NetworkingTests/PeerTests.swift
index 94e95026..16bb8066 100644
--- a/Networking/Tests/NetworkingTests/PeerTests.swift
+++ b/Networking/Tests/NetworkingTests/PeerTests.swift
@@ -211,7 +211,7 @@ struct PeerTests {
         let data1 = try await connection1.request(
             MockRequest(kind: .typeA, data: Data("hello world".utf8))
         )
-        #expect(data1.suffix(from: 4) == Data("hello world".utf8))
+        #expect(data1 == Data("hello world".utf8))
 
         let connection2 = try peer2.connect(
             to: NetAddr(ipAddress: "127.0.0.1", port: 8083)!, role: .validator
@@ -221,6 +221,6 @@ struct PeerTests {
         let data2 = try await connection2.request(
             MockRequest(kind: .typeB, data: Data("I am jam".utf8))
         )
-        #expect(data2.suffix(from: 4) == Data("I am jam".utf8))
+        #expect(data2 == Data("I am jam".utf8))
     }
 }