From fb3f67a59757f66fe161fd30b2288a389f2de471 Mon Sep 17 00:00:00 2001
From: Klaus Post <klauspost@gmail.com>
Date: Fri, 18 Apr 2025 11:10:55 +0200
Subject: [PATCH] Fix shared error buffer (#21203)

v.cancelFn(RemoteErr(m.Payload)) would use an already returned buffer.

Simplify code a bit as well by returning on errors.
---
 internal/grid/connection.go | 25 +++++++++++++++----------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/internal/grid/connection.go b/internal/grid/connection.go
index 981573365..a193b0e79 100644
--- a/internal/grid/connection.go
+++ b/internal/grid/connection.go
@@ -1625,23 +1625,28 @@ func (c *Connection) handleMuxServerMsg(ctx context.Context, m message) {
 			Msg: nil,
 			Err: RemoteErr(m.Payload),
 		})
+		if v.cancelFn != nil {
+			v.cancelFn(RemoteErr(m.Payload))
+		}
 		PutByteBuffer(m.Payload)
-	} else if m.Payload != nil {
+		v.close()
+		c.outgoing.Delete(m.MuxID)
+		return
+	}
+	// Return payload.
+	if m.Payload != nil {
 		v.response(m.Seq, Response{
 			Msg: m.Payload,
 			Err: nil,
 		})
 	}
+	// Close when EOF.
 	if m.Flags&FlagEOF != 0 {
-		if v.cancelFn != nil && m.Flags&FlagPayloadIsErr == 0 {
-			// We must obtain the lock before closing
-			// Otherwise others may pick up the error before close is called.
-			v.respMu.Lock()
-			v.closeLocked()
-			v.respMu.Unlock()
-		} else {
-			v.close()
-		}
+		// We must obtain the lock before closing
+		// Otherwise others may pick up the error before close is called.
+		v.respMu.Lock()
+		v.closeLocked()
+		v.respMu.Unlock()
 		if debugReqs {
 			fmt.Println(m.MuxID, c.String(), "handleMuxServerMsg: DELETING MUX")
 		}