From a79f4779487ead30125f311ef35596b88469fd8d Mon Sep 17 00:00:00 2001
From: Antoni Sawicki <tenox@google.com>
Date: Thu, 29 Oct 2020 07:16:14 -0700
Subject: [PATCH] add support for built-in wrp.html

---
 .gitignore |  1 +
 Makefile   |  9 ++++++---
 go.mod     |  1 +
 go.sum     |  2 ++
 wrp.go     | 47 ++++++++++++++++++++++++++++++++++++++++++-----
 5 files changed, 52 insertions(+), 8 deletions(-)

diff --git a/.gitignore b/.gitignore
index 7a6b458..b176066 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
 wrp-*
 wrp
 wrp.exe
+statik
diff --git a/Makefile b/Makefile
index d408cd0..eac4b98 100644
--- a/Makefile
+++ b/Makefile
@@ -1,9 +1,9 @@
 all: wrp
 
-wrp: wrp.go
+wrp: wrp.go statik
 	go build wrp.go
 
-cross: 
+cross: statik
 	GOOS=linux GOARCH=amd64 go build -a -o wrp-amd64-linux wrp.go
 	GOOS=freebsd GOARCH=amd64 go build -a -o wrp-amd64-freebsd wrp.go
 	GOOS=openbsd GOARCH=amd64 go build -a -o wrp-amd64-openbsd wrp.go
@@ -12,6 +12,9 @@ cross:
 	GOOS=linux GOARCH=arm go build -a -o wrp-arm-linux wrp.go
 	GOOS=linux GOARCH=arm64 go build -a -o wrp-arm64-linux wrp.go
 
+statik: wrp.html
+	go generate
+
 docker: wrp
 	docker build -t tenox7/wrp:latest .
 
@@ -23,4 +26,4 @@ gcrio:
 	docker push gcr.io/tenox7/wrp
 
 clean:
-	rm -rf wrp-* wrp
+	rm -rf wrp-* wrp statik
diff --git a/go.mod b/go.mod
index 2830299..71a265d 100644
--- a/go.mod
+++ b/go.mod
@@ -9,5 +9,6 @@ require (
 	github.com/ericpauley/go-quantize v0.0.0-20200331213906-ae555eb2afa4
 	github.com/gobwas/ws v1.0.3 // indirect
 	github.com/mailru/easyjson v0.7.1 // indirect
+	github.com/rakyll/statik v0.1.7
 	golang.org/x/sys v0.0.0-20200420163511-1957bb5e6d1f // indirect
 )
diff --git a/go.sum b/go.sum
index 96beb6c..9f437d8 100644
--- a/go.sum
+++ b/go.sum
@@ -21,6 +21,8 @@ github.com/mailru/easyjson v0.7.0 h1:aizVhC/NAAcKWb+5QsU1iNOZb4Yws5UO2I+aIprQITM
 github.com/mailru/easyjson v0.7.0/go.mod h1:KAzv3t3aY1NaHWoQz1+4F1ccyAH66Jk7yos7ldAVICs=
 github.com/mailru/easyjson v0.7.1 h1:mdxE1MF9o53iCb2Ghj1VfWvh7ZOwHpnVG/xwXrV90U8=
 github.com/mailru/easyjson v0.7.1/go.mod h1:KAzv3t3aY1NaHWoQz1+4F1ccyAH66Jk7yos7ldAVICs=
+github.com/rakyll/statik v0.1.7 h1:OF3QCZUuyPxuGEP7B4ypUa7sB/iHtqOTDYZXGM8KOdQ=
+github.com/rakyll/statik v0.1.7/go.mod h1:AlZONWzMtEnMs7W4e/1LURLiI49pIMmp6V9Unghqrcc=
 golang.org/x/sys v0.0.0-20200116001909-b77594299b42 h1:vEOn+mP2zCOVzKckCZy6YsCtDblrpj/w7B9nxGNELpg=
 golang.org/x/sys v0.0.0-20200116001909-b77594299b42/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200420163511-1957bb5e6d1f h1:gWF768j/LaZugp8dyS4UwsslYCYz9XgFxvlgsn0n9H8=
diff --git a/wrp.go b/wrp.go
index 5444ab7..360d8f6 100644
--- a/wrp.go
+++ b/wrp.go
@@ -5,6 +5,8 @@
 // Copyright (c) 2019-2020 Google LLC
 //
 
+//go:generate statik -f -src=. -include=wrp.html
+
 package main
 
 import (
@@ -16,6 +18,7 @@ import (
 	"image"
 	"image/gif"
 	"image/png"
+	"io/ioutil"
 	"log"
 	"math"
 	"math/rand"
@@ -34,6 +37,8 @@ import (
 	"github.com/chromedp/cdproto/page"
 	"github.com/chromedp/chromedp"
 	"github.com/ericpauley/go-quantize/quantize"
+	"github.com/rakyll/statik/fs"
+	_ "github.com/tenox7/wrp/statik"
 )
 
 var (
@@ -393,9 +398,40 @@ func (w wrpReq) capture() {
 	log.Printf("%s Done with capture for %s\n", w.req.RemoteAddr, w.url)
 }
 
+func tmpl(t string) string {
+	var tmpl []byte
+	fh, err := os.Open(t)
+	if err != nil {
+		goto statik
+	}
+	tmpl, err = ioutil.ReadAll(fh)
+	if err != nil {
+		goto statik
+	}
+	log.Printf("Got UI template from %v file\n", t)
+	return string(tmpl)
+
+statik:
+	sfs, err := fs.New()
+	if err != nil {
+		log.Fatal(err)
+	}
+	fhs, err := sfs.Open("/wrp.html")
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	tmpl, err = ioutil.ReadAll(fhs)
+	if err != nil {
+		log.Fatal(err)
+	}
+	log.Printf("Got UI template from built-in\n")
+	return string(tmpl)
+}
+
 // Main...
 func main() {
-	var addr, fgeom string
+	var addr, fgeom, tHtml string
 	var headless bool
 	var debug bool
 	var err error
@@ -405,6 +441,7 @@ func main() {
 	flag.BoolVar(&nodel, "n", false, "Do not free maps and images after use")
 	flag.StringVar(&deftype, "t", "gif", "Image type: gif|png")
 	flag.StringVar(&fgeom, "g", "1152x600x256", "Geometry: width x height x colors, height can be 0 for unlimited")
+	flag.StringVar(&tHtml, "ui", "wrp.html", "HTML template file for the UI")
 	flag.Parse()
 	if len(os.Getenv("PORT")) > 0 {
 		addr = ":" + os.Getenv(("PORT"))
@@ -440,13 +477,13 @@ func main() {
 	http.HandleFunc("/img/", imgServer)
 	http.HandleFunc("/shutdown/", haltServer)
 	http.HandleFunc("/favicon.ico", http.NotFound)
-	htmlTmpl, err = template.New("wrp.html").ParseFiles("wrp.html")
-	if err != nil {
-		log.Fatal(err)
-	}
 	log.Printf("Web Rendering Proxy Version %s\n", version)
 	log.Printf("Args: %q", os.Args)
 	log.Printf("Default Img Type: %v, Geometry: %+v", deftype, defgeom)
+	htmlTmpl, err = template.New("wrp.html").Parse(tmpl(tHtml))
+	if err != nil {
+		log.Fatal(err)
+	}
 	log.Printf("Starting WRP http server on %s\n", addr)
 	srv.Addr = addr
 	err = srv.ListenAndServe()