From 5f94ebbd7df1d658847af1435165422124223826 Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Thu, 22 Jul 2021 15:43:12 +0200
Subject: [PATCH] lnrpc: use request context in WebSocket proxy

The request context was not properly used to pass it along to the gRPC
endpoint which caused streaming calls to still be active on the gRPC
side even if the WS side already hung up.
We also issue an explicit close on the forwarding writer to signal when
the WS side was closed.
---
 lnrpc/websocket_proxy.go | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/lnrpc/websocket_proxy.go b/lnrpc/websocket_proxy.go
index 8c5001daa..9d4d2e9a7 100644
--- a/lnrpc/websocket_proxy.go
+++ b/lnrpc/websocket_proxy.go
@@ -147,12 +147,12 @@ func (p *WebsocketProxy) upgradeToWebSocketProxy(w http.ResponseWriter,
 		}
 	}()
 
-	ctx, cancelFn := context.WithCancel(context.Background())
+	ctx, cancelFn := context.WithCancel(r.Context())
 	defer cancelFn()
 
 	requestForwarder := newRequestForwardingReader()
 	request, err := http.NewRequestWithContext(
-		r.Context(), r.Method, r.URL.String(), requestForwarder,
+		ctx, r.Method, r.URL.String(), requestForwarder,
 	)
 	if err != nil {
 		p.logger.Errorf("WS: error preparing request:", err)
@@ -181,6 +181,7 @@ func (p *WebsocketProxy) upgradeToWebSocketProxy(w http.ResponseWriter,
 	go func() {
 		<-ctx.Done()
 		responseForwarder.Close()
+		requestForwarder.CloseWriter()
 	}()
 
 	go func() {