From ef32942dad8a9171729c9149980d9be44ec4066c Mon Sep 17 00:00:00 2001
From: Justin McCarthy <justin@strongdm.com>
Date: Sat, 16 Mar 2019 00:02:50 -0700
Subject: [PATCH] Serve now provides an HTTP server, listening on port 4000 by
 default. Listen port can be set with the --port flag to the serve command.
 Websocket refresh remains intact.

---
 internal/cli/serve.go    | 11 +++++++++--
 internal/render/html.go  |  9 ++++++---
 internal/render/site.go  | 14 ++++++++++++++
 internal/render/watch.go |  1 -
 4 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/internal/cli/serve.go b/internal/cli/serve.go
index 19184c6..82109a8 100644
--- a/internal/cli/serve.go
+++ b/internal/cli/serve.go
@@ -7,8 +7,15 @@ import (
 )
 
 var serveCommand = cli.Command{
-	Name:   "serve",
-	Usage:  "live updating version of the build command",
+	Name:  "serve",
+	Usage: "live updating version of the build command",
+	Flags: []cli.Flag{
+		cli.IntFlag{
+			Name:        "port",
+			Value:       4000,
+			Destination: &render.ServePort,
+		},
+	},
 	Action: serveAction,
 	Before: beforeAll(pandocMustExist, cleanContainers),
 }
diff --git a/internal/render/html.go b/internal/render/html.go
index 170ec0e..a51cb8c 100644
--- a/internal/render/html.go
+++ b/internal/render/html.go
@@ -15,7 +15,10 @@ import (
 
 const websocketReloader = `<script>
 (function(){
-	var ws = new WebSocket("ws://localhost:5122/ws")
+	var ws = new WebSocket("ws://localhost:%d/ws")
+	if (location.host != "") {
+		ws = new WebSocket("ws://"+location.host+"/ws")
+	}
 	var connected = false
 	ws.onopen = function(e) {
 		connected = true
@@ -74,7 +77,7 @@ func html(output string, live bool, errCh chan error, wg *sync.WaitGroup) {
 			}
 
 			if live {
-				w.Write([]byte(websocketReloader))
+				w.Write([]byte(fmt.Sprintf(websocketReloader, ServePort)))
 			}
 			w.Close()
 		}
@@ -82,7 +85,7 @@ func html(output string, live bool, errCh chan error, wg *sync.WaitGroup) {
 		if live {
 			if !opened {
 				opened = true
-				open.Run(filepath.Join(".", "output", "index.html"))
+				open.Run(fmt.Sprintf("http://127.0.0.1:%d/", ServePort))
 			}
 		} else {
 			wg.Done()
diff --git a/internal/render/site.go b/internal/render/site.go
index b967666..ca354ac 100644
--- a/internal/render/site.go
+++ b/internal/render/site.go
@@ -1,8 +1,10 @@
 package render
 
 import (
+	"fmt"
 	"net/http"
 	"os"
+	"path/filepath"
 	"sync"
 	"time"
 
@@ -11,6 +13,8 @@ import (
 	"github.com/yosssi/ace"
 )
 
+var ServePort int
+
 var upgrader = websocket.Upgrader{
 	ReadBufferSize:  1024,
 	WriteBufferSize: 1024,
@@ -87,6 +91,16 @@ func Build(output string, live bool) error {
 
 	if live {
 		watch(errCh)
+
+		go func() {
+			http.Handle("/", http.FileServer(http.Dir(filepath.Join(".", "output"))))
+			err := http.ListenAndServe(fmt.Sprintf("127.0.0.1:%d", ServePort), nil)
+			if err != nil {
+				panic(err)
+			}
+		}()
+
+		fmt.Printf("Serving content of output/ at http://127.0.0.1:%d (ctrl-c to quit)\n", ServePort)
 	}
 	// PDF
 	wg.Add(1)
diff --git a/internal/render/watch.go b/internal/render/watch.go
index a69f35e..07a6c30 100644
--- a/internal/render/watch.go
+++ b/internal/render/watch.go
@@ -45,7 +45,6 @@ func watch(errCh chan error) {
 	}
 
 	http.HandleFunc("/ws", serveWs)
-	go http.ListenAndServe("127.0.0.1:5122", nil)
 
 	return
 }