From 8d63f2afb1f2eb360dbe5d72ff30712ef779fc50 Mon Sep 17 00:00:00 2001
From: Huiping Guo <huiping192@gmail.com>
Date: Sun, 5 May 2024 01:08:57 +0900
Subject: [PATCH] fix bug

---
 Sources/HPRTMP/Network/NetworkClient.swift | 24 +++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/Sources/HPRTMP/Network/NetworkClient.swift b/Sources/HPRTMP/Network/NetworkClient.swift
index 153b0a9..c6272f2 100644
--- a/Sources/HPRTMP/Network/NetworkClient.swift
+++ b/Sources/HPRTMP/Network/NetworkClient.swift
@@ -18,6 +18,7 @@ final class RTMPClientHandler: ChannelInboundHandler {
   
   func channelRead(context: ChannelHandlerContext, data: NIOAny) {
     let data = self.unwrapInboundIn(data)
+    guard !data.isEmpty else { return }
     responseCallback(data)
   }
 }
@@ -52,6 +53,7 @@ class NetworkClient: NetworkConnectable {
   private var host: String?
   private var port: Int?
   
+  private var cachedReceivedData: Data = .init()
   private var dataPromise: EventLoopPromise<Data>?
   
   init() {
@@ -89,15 +91,27 @@ class NetworkClient: NetworkConnectable {
   }
   
   func receiveData() async throws -> Data {
-    guard let promise = self.dataPromise else {
-      throw NSError(domain: "RTMPClientError", code: 2, userInfo: [NSLocalizedDescriptionKey: "Data promise is not initialized."])
+    guard let channel = self.channel else {
+      print("Connection not established")
+      throw NSError(domain: "RTMPClientError", code: 1, userInfo: [NSLocalizedDescriptionKey: "Connection not established"])
     }
-    self.dataPromise = channel!.eventLoop.makePromise(of: Data.self)
-    return try await promise.futureResult.get()
+    if !cachedReceivedData.isEmpty {
+      let data = cachedReceivedData
+      cachedReceivedData = Data()
+      return data
+    }
+    
+    self.dataPromise = channel.eventLoop.makePromise(of: Data.self)
+    return try await dataPromise!.futureResult.get()
   }
   
   private func responseReceived(data: Data) {
-    self.dataPromise?.succeed(data)
+    cachedReceivedData.append(data)
+    if let dataPromise {
+      dataPromise.succeed(cachedReceivedData)
+      cachedReceivedData = Data()
+      self.dataPromise = nil
+    }
   }
   
   func shutdown() {