From 215ab162abae74b9dc6de2134964cad2a4cb12ac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Reis?= <joao.r.reis@outlook.com>
Date: Wed, 19 Oct 2022 19:17:12 +0100
Subject: [PATCH] fix #69

---
 proxy/pkg/zdmproxy/clientconn.go    | 18 ++++++++++--------
 proxy/pkg/zdmproxy/clienthandler.go |  1 -
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/proxy/pkg/zdmproxy/clientconn.go b/proxy/pkg/zdmproxy/clientconn.go
index b5039643..236597f0 100644
--- a/proxy/pkg/zdmproxy/clientconn.go
+++ b/proxy/pkg/zdmproxy/clientconn.go
@@ -165,6 +165,7 @@ func (cc *ClientConnector) listenForRequests() {
 		bufferedReader := bufio.NewReaderSize(cc.connection, cc.conf.RequestWriteBufferSizeBytes)
 		connectionAddr := cc.connection.RemoteAddr().String()
 		protocolErrOccurred := false
+		var alreadySentProtocolErr *frame.RawFrame
 		for cc.clientHandlerContext.Err() == nil {
 			f, err := readRawFrame(bufferedReader, connectionAddr, cc.clientHandlerContext)
 
@@ -174,14 +175,15 @@ func (cc *ClientConnector) listenForRequests() {
 					err, cc.clientHandlerContext, cc.clientHandlerCancelFunc, ClientConnectorLogPrefix, "reading", connectionAddr)
 				break
 			} else if protocolErrResponseFrame != nil {
-				f = protocolErrResponseFrame
-				if !protocolErrOccurred {
-					protocolErrOccurred = true
-					cc.sendResponseToClient(protocolErrResponseFrame)
-					cc.clientHandlerShutdownRequestCancelFn()
-					setDrainModeNowFunc()
-					continue
-				}
+				alreadySentProtocolErr = protocolErrResponseFrame
+				protocolErrOccurred = true
+				cc.sendResponseToClient(protocolErrResponseFrame)
+				continue
+			} else if alreadySentProtocolErr != nil {
+				clonedProtocolErr := alreadySentProtocolErr.Clone()
+				clonedProtocolErr.Header.StreamId = f.Header.StreamId
+				cc.sendResponseToClient(clonedProtocolErr)
+				continue
 			}
 
 			wg.Add(1)
diff --git a/proxy/pkg/zdmproxy/clienthandler.go b/proxy/pkg/zdmproxy/clienthandler.go
index 85e9a17f..d24f9f0c 100644
--- a/proxy/pkg/zdmproxy/clienthandler.go
+++ b/proxy/pkg/zdmproxy/clienthandler.go
@@ -643,7 +643,6 @@ func (ch *ClientHandler) tryProcessProtocolError(response *Response, protocolErr
 					errMsg, response.connectorType)
 			}
 			ch.clientConnector.sendResponseToClient(response.responseFrame)
-			ch.clientHandlerShutdownRequestCancelFn()
 		}
 		return true
 	}