Merge pull request #10 from aaronvb/bind-webui-address

Add a CLI flag to allow binding of web UI address

Author: aaronvb

README.md

@@ -43,14 +43,15 @@ Available Commands:
   version     Print version number of Request Hole
 
 Flags:
-  -a, --address string      sets the address for the endpoint (default "localhost")
-      --details             shows header details in the request
-  -h, --help                help for rh
-      --log string          writes incoming requests to the specified log file (example: --log rh.log)
-  -p, --port int            sets the port for the endpoint (default 8080)
-  -r, --response_code int   sets the response code (default 200)
-      --web                 runs a web server to show incoming requests
-      --web_port int        sets the port for the web server (default 8081)
+  -a, --address string       sets the address for the endpoint (default "localhost")
+      --details              shows header details in the request
+  -h, --help                 help for rh
+      --log string           writes incoming requests to the specified log file (example: --log rh.log)
+  -p, --port int             sets the port for the endpoint (default 8080)
+  -r, --response_code int    sets the response code (default 200)
+      --web                  runs the web UI to show incoming requests
+      --web_address string   sets the address for the web UI (default "localhost")
+      --web_port int         sets the port for the web UI (default 8081)
 
 Use "rh [command] --help" for more information about a command.
 ```

cmd/protocol.go

@@ -32,11 +32,13 @@ func httpCommand(cmd *cobra.Command, args []string) {
 		Port:         Port,
 		ResponseCode: ResponseCode,
 		Web:          Web,
+		WebAddress:   WebAddress,
 		WebPort:      WebPort,
 	}
 
 	if Web {
 		web := &renderer.Web{
+			Address:      WebAddress,
 			Port:         WebPort,
 			StaticFiles:  StaticFS,
 			RequestAddr:  Address,

cmd/root.go

@@ -15,6 +15,7 @@ var (
 	Port         int
 	ResponseCode int
 	Web          bool
+	WebAddress   string
 	WebPort      int
 	StaticFS     http.FileSystem
 )
@@ -40,6 +41,7 @@ func init() {
 	rootCmd.PersistentFlags().StringVar(&LogFile, "log", "", "writes incoming requests to the specified log file (example: --log rh.log)")
 
 	// Web server renderer
-	rootCmd.PersistentFlags().BoolVar(&Web, "web", false, "runs a web server to show incoming requests")
-	rootCmd.PersistentFlags().IntVar(&WebPort, "web_port", 8081, "sets the port for the web server")
+	rootCmd.PersistentFlags().BoolVar(&Web, "web", false, "runs the web UI to show incoming requests")
+	rootCmd.PersistentFlags().StringVar(&WebAddress, "web_address", "localhost", "sets the address for the web UI")
+	rootCmd.PersistentFlags().IntVar(&WebPort, "web_port", 8081, "sets the port for the web UI")
 }

pkg/renderer/web.go

@@ -24,15 +24,31 @@ import (
 	"github.com/rs/cors"
 )
 
+// Web is the renderer for the web UI.
 type Web struct {
-	BuildInfo     map[string]string
-	Port          int
-	ResponseCode  int
-	RequestAddr   string
-	RequestPort   int
-	StaticFiles   http.FileSystem
-	mu            sync.Mutex
-	requests      []*protocol.RequestPayload
+	// Address is the address the web UI server will bind to.
+	Address string
+
+	// BuildInfo contains build information.
+	BuildInfo map[string]string
+
+	// Port is the port the web UI server will run on.
+	Port int
+
+	// These are used for showing information about the request endpoint in the
+	// web UI.
+	ResponseCode int
+	RequestAddr  string
+	RequestPort  int
+
+	StaticFiles http.FileSystem
+
+	mu sync.Mutex
+
+	// requests contain the incoming requests.
+	requests []*protocol.RequestPayload
+
+	// subscriptions contain the websocket connections to our graphql subscribers.
 	subscriptions map[string]chan *protocol.RequestPayload
 }
 
@@ -42,7 +58,7 @@ func (web *Web) Start(wg *sync.WaitGroup, rp chan protocol.RequestPayload, q cha
 	web.requests = make([]*protocol.RequestPayload, 0)
 	web.subscriptions = make(map[string]chan *protocol.RequestPayload)
 
-	addr := fmt.Sprintf("localhost:%d", web.Port)
+	addr := fmt.Sprintf("%s:%d", web.Address, web.Port)
 	errorLog := log.New(&httpErrorLog{}, "", 0)
 
 	srv := &http.Server{

pkg/server/server.go

@@ -48,6 +48,9 @@ type FlagData struct {
 	// Web determines if we use the web renderer, otherwise defaults to the printer renderer.
 	Web bool
 
+	// WebAddress is the address the web UI will bind to.
+	WebAddress string
+
 	// WebPort defines which port we host the web renderer at, defaults to 8081.
 	WebPort int
 }
@@ -119,7 +122,7 @@ func (s *Server) startText() string {
 	text := fmt.Sprintf("%s %s\nListening on http://%s:%d", primary, version, s.FlagData.Addr, s.FlagData.Port)
 
 	if s.FlagData.Web {
-		text = fmt.Sprintf("%s\nWeb running on: http://localhost:%d", text, s.FlagData.WebPort)
+		text = fmt.Sprintf("%s\nWeb running on: http://%s:%d", text, s.FlagData.WebAddress, s.FlagData.WebPort)
 	}
 
 	if s.FlagData.Details {

pkg/server/server_test.go

@@ -55,3 +55,43 @@ func TestStartTextWithLogFile(t *testing.T) {
 		t.Errorf("Expected %s, got %s", expected, result)
 	}
 }
+
+func TestStartTextWithWebUIDefault(t *testing.T) {
+	pterm.DisableColor()
+	flags := FlagData{
+		Addr:      "localhost",
+		Port:      8080,
+		BuildInfo: map[string]string{"version": "dev"},
+		Web:       true,
+	}
+	server := Server{FlagData: flags}
+	result := server.startText()
+	expected := fmt.Sprintf(
+		"Request Hole %s\nListening on http://%s:%d\nWeb running on: http://%s:%d", "dev",
+		server.FlagData.Addr, server.FlagData.Port, server.FlagData.WebAddress, server.FlagData.WebPort)
+
+	if result != expected {
+		t.Errorf("Expected %s, got %s", expected, result)
+	}
+}
+
+func TestStartTextWithWebUICustomFlags(t *testing.T) {
+	pterm.DisableColor()
+	flags := FlagData{
+		Addr:       "localhost",
+		Port:       8080,
+		BuildInfo:  map[string]string{"version": "dev"},
+		Web:        true,
+		WebAddress: "0.0.0.0",
+		WebPort:    8082,
+	}
+	server := Server{FlagData: flags}
+	result := server.startText()
+	expected := fmt.Sprintf(
+		"Request Hole %s\nListening on http://%s:%d\nWeb running on: http://%s:%d", "dev",
+		server.FlagData.Addr, server.FlagData.Port, server.FlagData.WebAddress, server.FlagData.WebPort)
+
+	if result != expected {
+		t.Errorf("Expected %s, got %s", expected, result)
+	}
+}