From a76e752f3b33f774ce516f0e0f767449f4e4daf4 Mon Sep 17 00:00:00 2001
From: Anthony Ronning <anthonyronning@gmail.com>
Date: Mon, 6 Apr 2020 15:16:02 +0200
Subject: [PATCH] config+lnd+rpcserver: add CORS origin config for REST

---
 config.go    |  1 +
 rpcserver.go | 55 +++++++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/config.go b/config.go
index 11916a1a..b23bb5e4 100644
--- a/config.go
+++ b/config.go
@@ -159,6 +159,7 @@ type Config struct {
 	RawExternalIPs   []string `long:"externalip" description:"Add an ip:port to the list of local addresses we claim to listen on to peers. If a port is not specified, the default (9735) will be used regardless of other parameters"`
 	RPCListeners     []net.Addr
 	RESTListeners    []net.Addr
+	RestCORS         []string `long:"restcors" description:"Add an ip:port/hostname to allow cross origin access from. To allow all origins, set as \"*\"."`
 	Listeners        []net.Addr
 	ExternalIPs      []net.Addr
 	DisableListen    bool          `long:"nolisten" description:"Disable listening for incoming peer connections"`
diff --git a/rpcserver.go b/rpcserver.go
index 1a1d4d6d..abe3dd22 100644
--- a/rpcserver.go
+++ b/rpcserver.go
@@ -802,6 +802,12 @@ func (r *rpcServer) Start() error {
 	// Wrap the default grpc-gateway handler with the WebSocket handler.
 	restHandler := lnrpc.NewWebSocketProxy(restMux, rpcsLog)
 
+	// Set the CORS headers if configured. This wraps the HTTP handler with
+	// another handler.
+	if len(r.cfg.RestCORS) > 0 {
+		restHandler = allowCORS(restHandler, r.cfg.RestCORS)
+	}
+
 	// With our custom REST proxy mux created, register our main RPC and
 	// give all subservers a chance to register as well.
 	err := lnrpc.RegisterLightningHandlerFromEndpoint(
@@ -855,7 +861,8 @@ func (r *rpcServer) Start() error {
 
 			// Create our proxy chain now. A request will pass
 			// through the following chain:
-			// req ---> WS proxy ---> REST proxy --> gRPC endpoint
+			// req ---> CORS handler --> WS proxy --->
+			//   REST proxy --> gRPC endpoint
 			err := http.Serve(lis, restHandler)
 			if err != nil && !lnrpc.IsClosedConnError(err) {
 				rpcsLog.Error(err)
@@ -922,6 +929,52 @@ func addrPairsToOutputs(addrPairs map[string]int64) ([]*wire.TxOut, error) {
 	return outputs, nil
 }
 
+// allowCORS wraps the given http.Handler with a function that adds the
+// Access-Control-Allow-Origin header to the response.
+func allowCORS(handler http.Handler, origins []string) http.Handler {
+	allowHeaders := "Access-Control-Allow-Headers"
+	allowMethods := "Access-Control-Allow-Methods"
+	allowOrigin := "Access-Control-Allow-Origin"
+
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		origin := r.Header.Get("Origin")
+
+		// Skip everything if the browser doesn't send the Origin field.
+		if origin == "" {
+			handler.ServeHTTP(w, r)
+			return
+		}
+
+		// Set the static header fields first.
+		w.Header().Set(
+			allowHeaders,
+			"Content-Type, Accept, Grpc-Metadata-Macaroon",
+		)
+		w.Header().Set(allowMethods, "GET, POST, DELETE")
+
+		// Either we allow all origins or the incoming request matches
+		// a specific origin in our list of allowed origins.
+		for _, allowedOrigin := range origins {
+			if allowedOrigin == "*" || origin == allowedOrigin {
+				// Only set allowed origin to requested origin.
+				w.Header().Set(allowOrigin, origin)
+
+				break
+			}
+		}
+
+		// For a pre-flight request we only need to send the headers
+		// back. No need to call the rest of the chain.
+		if r.Method == "OPTIONS" {
+			return
+		}
+
+		// Everything's prepared now, we can pass the request along the
+		// chain of handlers.
+		handler.ServeHTTP(w, r)
+	})
+}
+
 // sendCoinsOnChain makes an on-chain transaction in or to send coins to one or
 // more addresses specified in the passed payment map. The payment map maps an
 // address to a specified output value to be sent to that address.