From 1debc193a42542347f42db0f5bbb7016f50768c4 Mon Sep 17 00:00:00 2001
From: Jens Neuse <jens.neuse@gmx.de>
Date: Mon, 17 Apr 2023 09:05:44 +0200
Subject: [PATCH] fix: avoid panic if client immediately disconnects on request

---
 .../graphql_datasource/graphql_sse_handler.go | 28 ++++++++++++++++++-
 .../datasource/httpclient/httpclient.go       |  7 +++--
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/pkg/engine/datasource/graphql_datasource/graphql_sse_handler.go b/pkg/engine/datasource/graphql_datasource/graphql_sse_handler.go
index 9ddb26587f..9b810df52d 100644
--- a/pkg/engine/datasource/graphql_datasource/graphql_sse_handler.go
+++ b/pkg/engine/datasource/graphql_datasource/graphql_sse_handler.go
@@ -65,7 +65,28 @@ func (h *gqlSSEConnectionHandler) StartBlocking(sub Subscription) {
 }
 
 func (h *gqlSSEConnectionHandler) subscribe(ctx context.Context, sub Subscription, dataCh, errCh chan []byte) {
-	resp, err := h.performSubscriptionRequest(ctx)
+	// if we used the downstream context, we got a panic if the downstream client disconnects immediately after the request was sent
+	// this happens, e.g. with React strict mode which renders the component twice
+	// to solve the issue, we use a separate context for the origin request
+	// with a goroutine that cancels the origin request if the downstream client disconnects
+	// in order to free resources after the initial handshake, we cancel the goroutine after we've received a response
+	originCtx, cancelOriginRequest := context.WithCancel(context.Background())
+	defer cancelOriginRequest()
+	waitForResponse, cancelWaitForResponse := context.WithCancel(context.Background())
+	go func() {
+		select {
+		case <-ctx.Done():
+			// cancel the origin request if the downstream client disconnected
+			cancelOriginRequest()
+		case <-waitForResponse.Done():
+			// end the goroutine to free resources
+		}
+	}()
+	resp, err := h.performSubscriptionRequest(originCtx)
+	// cancel the goroutine to free resources
+	// the originRequest will be canceled through defer cancelOriginRequest()
+	// as we check on every iteration (below) if the downstream ctx is done
+	cancelWaitForResponse()
 	if err != nil {
 		h.log.Error("failed to perform subscription request", log.Error(err))
 
@@ -117,6 +138,11 @@ func (h *gqlSSEConnectionHandler) subscribe(ctx context.Context, sub Subscriptio
 					continue
 				}
 
+				if ctx.Err() != nil {
+					// request context was canceled do not send an error as channel will be closed
+					return
+				}
+
 				dataCh <- data
 			case bytes.HasPrefix(line, headerEvent):
 				event := trim(line[len(headerEvent):])
diff --git a/pkg/engine/datasource/httpclient/httpclient.go b/pkg/engine/datasource/httpclient/httpclient.go
index bc865682a4..b1b12b252f 100644
--- a/pkg/engine/datasource/httpclient/httpclient.go
+++ b/pkg/engine/datasource/httpclient/httpclient.go
@@ -55,8 +55,11 @@ func CtxSetUndefinedVariables(ctx context.Context, undefinedVariables []string)
 }
 
 func CtxGetUndefinedVariables(ctx context.Context) []string {
-	undefinedVariables, _ := ctx.Value(removeUndefinedVariables).([]string)
-	return undefinedVariables
+	undefinedVariables := ctx.Value(removeUndefinedVariables)
+	if undefinedVariables, ok := undefinedVariables.([]string); ok {
+		return undefinedVariables
+	}
+	return nil
 }
 
 func wrapQuotesIfString(b []byte) []byte {