From 70752bfb4f4f08cef08b52bfc8c20573438a029f Mon Sep 17 00:00:00 2001
From: Alex <alex@roese.dev>
Date: Mon, 3 Jan 2022 22:56:12 +0100
Subject: [PATCH] consumer message handling

---
 .../listener/SocketClientMessageListener.java | 25 +++++++++++--------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/src/main/java/net/krakatoa/proxy/listener/SocketClientMessageListener.java b/src/main/java/net/krakatoa/proxy/listener/SocketClientMessageListener.java
index d57eb46..1897178 100644
--- a/src/main/java/net/krakatoa/proxy/listener/SocketClientMessageListener.java
+++ b/src/main/java/net/krakatoa/proxy/listener/SocketClientMessageListener.java
@@ -14,6 +14,7 @@ import earth.krakatao.protocol.KraSocketClientProtocolStatus;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 import net.krakatoa.proxy.ProxySystem;
 import net.md_5.bungee.api.ProxyServer;
 import net.md_5.bungee.api.chat.TextComponent;
@@ -41,18 +42,20 @@ public class SocketClientMessageListener implements KraSocketClientEventInterfac
     System.out.println("consumer ids " + ProxySystem.getInstance().getKraSocketClientProtocol()
         .getConsumerHashMap().keySet());
 
-    if (kraProtocolMessage.getStatus() == KraSocketClientProtocolStatus.REPLY.getStatus()) {
-      ProxySystem.getInstance().getKraSocketClientProtocol().getConsumerHashMap()
-          .forEach((cmdID, consumer) -> {
-            if (cmdID == kraProtocolMessage.getCmdID()) {
-              ProxySystem.getInstance().getLogger()
-                  .info("accept " + cmdID + " " + kraProtocolMessage.getCmdID());
-              consumer.accept(kraProtocolMessage);
+    if (kraProtocolMessage.getStatus() == KraSocketClientProtocolStatus.REPLY.getStatus()
+        && ProxySystem.getInstance().getKraSocketClientProtocol().getConsumerHashMap()
+        .containsKey(kraProtocolMessage.getCmdID())) {
+      Consumer<KraSocketClientProtocolMessage> consumer = ProxySystem.getInstance()
+          .getKraSocketClientProtocol().getConsumerHashMap().get(kraProtocolMessage.getCmdID());
 
-              ProxySystem.getInstance().getKraSocketClientProtocol().getConsumerHashMap()
-                  .remove(cmdID);
-            }
-          });
+      if (consumer != null) {
+        ProxySystem.getInstance().getLogger()
+            .info("accept " + kraProtocolMessage.getCmdID());
+        consumer.accept(kraProtocolMessage);
+
+        ProxySystem.getInstance().getKraSocketClientProtocol()
+            .removeData(kraProtocolMessage.getCmdID());
+      }
     }
 
     switch (kraProtocolMessage.getCmdNumber()) {