From c5c28564e9ae97a41262c0f3a6df959691687ed0 Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Thu, 6 Aug 2020 12:07:07 +0200
Subject: [PATCH] lnrpc: add macaroon workaround for WebSockets in browsers

For security reasons, browsers are limited in the header fields they can
send when opening a WebSocket connection. Specifically, the macaroon
cannot be sent in the Grpc-Metadata-Macaroon header field as that would
be possible for normal REST requests. Instead we only have the special
field "Sec-Websocket-Protocol" that can be used to transport custom
data. We allow the macaroon to be sent there and transform it into a
proper header field for the target request.
---
 lnrpc/websocket_proxy.go | 63 +++++++++++++++++++++++++++++++++++-----
 1 file changed, 56 insertions(+), 7 deletions(-)

diff --git a/lnrpc/websocket_proxy.go b/lnrpc/websocket_proxy.go
index 921b2171..3cb701be 100644
--- a/lnrpc/websocket_proxy.go
+++ b/lnrpc/websocket_proxy.go
@@ -21,6 +21,16 @@ const (
 	// This is necessary because the WebSocket API specifies that a
 	// handshake request must always be done through a GET request.
 	MethodOverrideParam = "method"
+
+	// HeaderWebSocketProtocol is the name of the WebSocket protocol
+	// exchange header field that we use to transport additional header
+	// fields.
+	HeaderWebSocketProtocol = "Sec-Websocket-Protocol"
+
+	// WebSocketProtocolDelimiter is the delimiter we use between the
+	// additional header field and its value. We use the plus symbol because
+	// the default delimiters aren't allowed in the protocol names.
+	WebSocketProtocolDelimiter = "+"
 )
 
 var (
@@ -32,6 +42,13 @@ var (
 		"Referer":                true,
 		"Grpc-Metadata-Macaroon": true,
 	}
+
+	// defaultProtocolsToAllow are additional header fields that we allow
+	// to be transported inside of the Sec-Websocket-Protocol field to be
+	// forwarded to the backend.
+	defaultProtocolsToAllow = map[string]bool{
+		"Grpc-Metadata-Macaroon": true,
+	}
 )
 
 // NewWebSocketProxy attempts to expose the underlying handler as a response-
@@ -101,13 +118,13 @@ func (p *WebsocketProxy) upgradeToWebSocketProxy(w http.ResponseWriter,
 		p.logger.Errorf("WS: error preparing request:", err)
 		return
 	}
-	for header := range r.Header {
-		headerName := textproto.CanonicalMIMEHeaderKey(header)
-		forward, ok := defaultHeadersToForward[headerName]
-		if ok && forward {
-			request.Header.Set(headerName, r.Header.Get(header))
-		}
-	}
+
+	// Allow certain headers to be forwarded, either from source headers
+	// or the special Sec-Websocket-Protocol header field.
+	forwardHeaders(r.Header, request.Header)
+
+	// Also allow the target request method to be overwritten, as all
+	// WebSocket establishment calls MUST be GET requests.
 	if m := r.URL.Query().Get(MethodOverrideParam); m != "" {
 		request.Method = m
 	}
@@ -182,6 +199,38 @@ func (p *WebsocketProxy) upgradeToWebSocketProxy(w http.ResponseWriter,
 	}
 }
 
+// forwardHeaders forwards certain allowed header fields from the source request
+// to the target request. Because browsers are limited in what header fields
+// they can send on the WebSocket setup call, we also allow additional fields to
+// be transported in the special Sec-Websocket-Protocol field.
+func forwardHeaders(source, target http.Header) {
+	// Forward allowed header fields directly.
+	for header := range source {
+		headerName := textproto.CanonicalMIMEHeaderKey(header)
+		forward, ok := defaultHeadersToForward[headerName]
+		if ok && forward {
+			target.Set(headerName, source.Get(header))
+		}
+	}
+
+	// Browser aren't allowed to set custom header fields on WebSocket
+	// requests. We need to allow them to submit the macaroon as a WS
+	// protocol, which is the only allowed header. Set any "protocols" we
+	// declare valid as header fields on the forwarded request.
+	protocol := source.Get(HeaderWebSocketProtocol)
+	for key := range defaultProtocolsToAllow {
+		if strings.HasPrefix(protocol, key) {
+			// The format is "<protocol name>+<value>". We know the
+			// protocol string starts with the name so we only need
+			// to set the value.
+			values := strings.Split(
+				protocol, WebSocketProtocolDelimiter,
+			)
+			target.Set(key, values[1])
+		}
+	}
+}
+
 // newRequestForwardingReader creates a new request forwarding pipe.
 func newRequestForwardingReader() *requestForwardingReader {
 	r, w := io.Pipe()