From ddaee90378be50f15eb82bd44ebfdd5acb23b3d6 Mon Sep 17 00:00:00 2001
From: meows <b5c6@protonmail.com>
Date: Wed, 18 Dec 2019 08:59:07 -0600
Subject: [PATCH] rpc: refactor to install module validity checker for WS
 server too

Signed-off-by: meows <b5c6@protonmail.com>
---
 rpc/endpoints.go | 55 ++++++++++++++++++++++++++++--------------------
 1 file changed, 32 insertions(+), 23 deletions(-)

diff --git a/rpc/endpoints.go b/rpc/endpoints.go
index 845acf849b..dae5e0099b 100644
--- a/rpc/endpoints.go
+++ b/rpc/endpoints.go
@@ -23,34 +23,42 @@ import (
 	"github.com/ethereum/go-ethereum/log"
 )
 
+// mustAvailableModule enforces that requested api modules (eg. via --rpcapi) are actually
+// available API services. If an invalid module is given (ie API "foo" wanted which does not exist),
+// then log.Crit is used to cause program to exit, logging the invalid module and a list of available
+// API service names.
+func mustAvailableModule(module string, apis []API) {
+	apiExists := false
+	for _, api := range apis {
+		if module == api.Namespace {
+			apiExists = true
+			break
+		}
+	}
+	if !apiExists {
+		log.Crit("invalid api module", "module", module, "available", func() string {
+			available := []string{}
+		outer:
+			for _, api := range apis {
+				// Only include unique api names
+				for _, av := range available {
+					if av == api.Namespace {
+						continue outer
+					}
+				}
+				available = append(available, api.Namespace)
+			}
+			return strings.Join(available, ",")
+		}())
+	}
+}
+
 // StartHTTPEndpoint starts the HTTP RPC endpoint, configured with cors/vhosts/modules
 func StartHTTPEndpoint(endpoint string, apis []API, modules []string, cors []string, vhosts []string, timeouts HTTPTimeouts) (net.Listener, *Server, error) {
 	// Generate the whitelist based on the allowed modules
 	whitelist := make(map[string]bool)
 	for _, module := range modules {
-		apiExists := false
-		for _, api := range apis {
-			if module == api.Namespace {
-				apiExists = true
-				break
-			}
-		}
-		if !apiExists {
-			log.Crit("invalid api module", "module", module, "available", func() string {
-				available := []string{}
-				outer:
-				for _, api := range apis {
-					// Only include unique api names
-					for _, av := range available {
-						if av == api.Namespace {
-							continue outer
-						}
-					}
-					available = append(available, api.Namespace)
-				}
-				return strings.Join(available, ",")
-			}())
-		}
+		mustAvailableModule(module, apis)
 		whitelist[module] = true
 	}
 	// Register all the APIs exposed by the services
@@ -81,6 +89,7 @@ func StartWSEndpoint(endpoint string, apis []API, modules []string, wsOrigins []
 	// Generate the whitelist based on the allowed modules
 	whitelist := make(map[string]bool)
 	for _, module := range modules {
+		mustAvailableModule(module, apis)
 		whitelist[module] = true
 	}
 	// Register all the APIs exposed by the services