allowing injecting custom JS, CSS, and other files into the gRPC UI w…

…eb server (#164)

cmd/grpcui/grpcui.go

@@ -17,6 +17,7 @@ import (
 	"net/http/httptest"
 	"net/http/httputil"
 	"os"
+	"path/filepath"
 	"sort"
 	"strconv"
 	"strings"
@@ -148,6 +149,10 @@ var (
 		Example: "/debug/grpcui".`))
 	services multiString
 	methods  multiString
+
+	extraJS     multiString
+	extraCSS    multiString
+	otherAssets multiString
 )
 
 func init() {
@@ -231,6 +236,24 @@ func init() {
 		-use-reflection is used in combination with a -proto or -protoset flag,
 		the provided descriptor sources will be used in addition to server
 		reflection to resolve messages and extensions.`))
+	flags.Var(&extraJS, "extra-js", prettify(`
+		Indicates the name of a JavaScript file to load from the web form. This
+		allows injecting custom behavior into the page. Multiple files can be
+		added by specifying multiple -extra-js flags.`))
+	flags.Var(&extraCSS, "extra-css", prettify(`
+		Indicates the name of a CSS file to load from the web form. This allows
+		injecting custom styles into the page, to customize the look. Multiple
+		files can be added by specifying multiple -extra-css flags.`))
+	flags.Var(&otherAssets, "also-serve", prettify(`
+		Indicates the name of an additional file or folder that the gRPC UI web
+		server can serve. This can be useful for serving other assets, like
+		images, when a custom CSS is used via -extra-css flags. Multiple assets
+		can be added to the server by specifying multiple -also-serve flags. The
+		named file will be available at a URI of "/<base-name>", where
+		<base-name> is the name of the file, excluding its path. If the given
+		name is a folder, the folder's contents are available at URIs that are
+		under "/<base-name>/". It is an error to specify multiple files or
+		folders that have the same base name.`))
 }
 
 type multiString []string
@@ -376,6 +399,11 @@ func main() {
 		fail(nil, `The -base-path must begin with a slash ("/")`)
 	}
 
+	assetNames := map[string]string{}
+	checkAssetNames(assetNames, extraJS, true)
+	checkAssetNames(assetNames, extraCSS, true)
+	checkAssetNames(assetNames, otherAssets, false)
+
 	// Protoset or protofiles provided and -use-reflection unset
 	if !reflection.set && (len(protoset) > 0 || len(protoFiles) > 0) {
 		reflection.val = false
@@ -563,6 +591,10 @@ func main() {
 	if debug.set {
 		handlerOpts = append(handlerOpts, standalone.WithClientDebug(debug.val))
 	}
+	handlerOpts = append(handlerOpts, configureJSandCSS(extraJS, standalone.AddJSFile)...)
+	handlerOpts = append(handlerOpts, configureJSandCSS(extraCSS, standalone.AddCSSFile)...)
+	handlerOpts = append(handlerOpts, configureAssets(otherAssets)...)
+
 	handler := standalone.Handler(cc, target, methods, allFiles, handlerOpts...)
 	if *maxTime > 0 {
 		timeout := time.Duration(*maxTime * float64(time.Second))
@@ -709,6 +741,71 @@ func fail(err error, msg string, args ...interface{}) {
 	}
 }
 
+func checkAssetNames(soFar map[string]string, names []string, requireFile bool) {
+	for _, n := range names {
+		st, err := os.Stat(n)
+		if err != nil {
+			if os.IsNotExist(err) {
+				fail(nil, "File %q does not exist", n)
+			} else {
+				fail(err, "Failed to check existence of file %q", n)
+			}
+		}
+		if requireFile && st.IsDir() {
+			fail(nil, "Path %q is a folder, not a file", n)
+		}
+
+		base := filepath.Base(n)
+		if existing, ok := soFar[base]; ok {
+			fail(nil, "Multiple assets with the same base name specified: %s and %s", existing, n)
+		}
+		soFar[base] = n
+	}
+}
+
+func configureJSandCSS(names []string, fn func(string, func() (io.ReadCloser, error)) standalone.HandlerOption) []standalone.HandlerOption {
+	opts := make([]standalone.HandlerOption, len(names))
+	for i := range names {
+		name := names[i] // no loop variable so that we don't close over loop var in lambda below
+		open := func() (io.ReadCloser, error) {
+			return os.Open(name)
+		}
+		opts[i] = fn(filepath.Base(name), open)
+	}
+	return opts
+}
+
+func configureAssets(names []string) []standalone.HandlerOption {
+	opts := make([]standalone.HandlerOption, len(names))
+	for i := range names {
+		name := names[i] // no loop variable so that we don't close over loop var in lambdas below
+		st, err := os.Stat(name)
+		if err != nil {
+			fail(err, "failed to inspect file %q", name)
+		}
+		if st.IsDir() {
+			open := func(p string) (io.ReadCloser, error) {
+				path := filepath.Join(name, p)
+				st, err := os.Stat(path)
+				if err == nil && st.IsDir() {
+					// Strangely, os.Open does not return an error if given a directory
+					// and instead returns an empty reader :(
+					// So check that first and return a 404 if user indicates directory name
+					return nil, os.ErrNotExist
+				}
+				return os.Open(path)
+			}
+			opts[i] = standalone.ServeAssetDirectory(filepath.Base(name), open)
+		} else {
+			open := func() (io.ReadCloser, error) {
+				return os.Open(name)
+			}
+			opts[i] = standalone.ServeAssetFile(filepath.Base(name), open)
+		}
+	}
+	return opts
+}
+
 type svcConfig struct {
 	includeService bool
 	includeMethods map[string]struct{}

standalone/opts.go

@@ -2,6 +2,7 @@ package standalone
 
 import (
 	"html/template"
+	"io"
 	"path"
 )
 
@@ -55,6 +56,16 @@ func AddJS(filename string, js []byte) HandlerOption {
 	})
 }
 
+// AddJSFile is like AddJS except that the contents are provided in the form of
+// a function that is used to "open" the file to read. This means that the
+// contents of the file need not be eagerly loaded into memory. Each time a
+// request is received for this file, the function is called.
+func AddJSFile(filename string, open func() (io.ReadCloser, error)) HandlerOption {
+	return optFunc(func(opts *handlerOptions) {
+		opts.tmplResources = append(opts.tmplResources, newDeferredResource(path.Join("/s", filename), open, "text/javascript; charset=utf-8"))
+	})
+}
+
 // AddCSS adds a CSS file to Handler, serving the supplied contents at the URI
 // "/s/<filename>" with a Content-Type of "text/css; charset=UTF-8". It
 // will also be added to the AddlResources field of the WebFormContainerTemplateData
@@ -68,6 +79,16 @@ func AddCSS(filename string, css []byte) HandlerOption {
 	})
 }
 
+// AddCSSFile is like AddCSS except that the contents are provided in the form
+// of a function that is used to "open" the file to read. This means that the
+// contents of the file need not be eagerly loaded into memory. Each time a
+// request is received for this file, the function is called.
+func AddCSSFile(filename string, open func() (io.ReadCloser, error)) HandlerOption {
+	return optFunc(func(opts *handlerOptions) {
+		opts.tmplResources = append(opts.tmplResources, newDeferredResource(path.Join("/s", filename), open, "text/css; charset=utf-8"))
+	})
+}
+
 // ServeAsset will add an additional file to Handler, serving the supplied contents
 // at the URI "/s/<filename>" with a Content-Type that is computed based on the given
 // filename's extension.
@@ -81,6 +102,33 @@ func ServeAsset(filename string, contents []byte) HandlerOption {
 	})
 }
 
+// ServeAssetFile is like ServeAsset except that the contents are provided in
+// the form of a function that is used to "open" the file to read. This means
+// that the contents of the file need not be eagerly loaded into memory. Each
+// time a request is received for this file, the function is called.
+func ServeAssetFile(filename string, open func() (io.ReadCloser, error)) HandlerOption {
+	return optFunc(func(opts *handlerOptions) {
+		opts.servedOnlyResources = append(opts.servedOnlyResources, newDeferredResource(path.Join("/s", filename), open, ""))
+	})
+}
+
+// ServeAssetDirectory is similar to ServeAssetFile except the give name is the
+// root of a subtree, which can be used to serve a directory of assets. When a
+// request is received, the remaining relative path is provided to the open
+// function, to indicate which path in the subtree to open. For example, if the
+// given name is "foo/bar" and a request is made for "foo/bar/baz/buzz", then
+// the open function will be called with "baz/buzz" as the argument.
+//
+// If a given path does not exist or is a directory, not a file, the open function
+// should return an error that can be classified via os.IsNotExist, so that the
+// server can return a "404 Not Found" status. Any other error will result in the
+// server sending a "500 Internal Server Error" status.
+func ServeAssetDirectory(dirname string, open func(filename string) (io.ReadCloser, error)) HandlerOption {
+	return optFunc(func(opts *handlerOptions) {
+		opts.servedOnlyResources = append(opts.servedOnlyResources, newDeferredResourceFolder(path.Join("/s", dirname), open))
+	})
+}
+
 // WithDefaultMetadata sets the default metadata in the web form to the given
 // values. Each string should be in the form "name: value".
 func WithDefaultMetadata(headers []string) HandlerOption {

standalone/standalone.go

@@ -8,9 +8,14 @@ import (
 	"encoding/base64"
 	"fmt"
 	"html/template"
+	"io"
+	"io/ioutil"
 	"mime"
 	"net/http"
+	"os"
 	"path"
+	"path/filepath"
+	"strconv"
 	"strings"
 
 	"github.com/jhump/protoreflect/desc"
@@ -148,7 +153,8 @@ func getIndexContents(tmpl *template.Template, target string, webFormHTML []byte
 
 type resource struct {
 	Path        string
-	Data        []byte
+	Len         int
+	Open        func(string) (io.ReadCloser, error)
 	ContentType string
 	ETag        string
 	Public      bool
@@ -159,32 +165,90 @@ func newResource(uriPath string, data []byte, contentType string, public bool) *
 		contentType = mime.TypeByExtension(path.Ext(uriPath))
 	}
 	return &resource{
-		Path:        uriPath,
-		Data:        data,
+		Path: uriPath,
+		Open: func(_ string) (io.ReadCloser, error) {
+			return ioutil.NopCloser(bytes.NewReader(data)), nil
+		},
+		Len:         len(data),
 		ContentType: contentType,
 		ETag:        computeETag(data),
 		Public:      public,
 	}
 }
 
+func newDeferredResource(uriPath string, open func() (io.ReadCloser, error), contentType string) *resource {
+	if contentType == "" {
+		contentType = mime.TypeByExtension(path.Ext(uriPath))
+	}
+	return &resource{
+		Path: uriPath,
+		Open: func(_ string) (io.ReadCloser, error) {
+			return open()
+		},
+		ContentType: contentType,
+	}
+}
+
+func newDeferredResourceFolder(uriPath string, open func(string) (io.ReadCloser, error)) *resource {
+	return &resource{
+		Path: uriPath + "/",
+		Open: func(filename string) (io.ReadCloser, error) {
+			return open(filename)
+		},
+	}
+}
+
 func handle(mux *http.ServeMux, res *resource) {
 	mux.Handle(res.Path, res)
+	if withoutSlash := strings.TrimSuffix(res.Path, "/"); withoutSlash != res.Path {
+		// if res.Path is a folder, return a 404 if the base directory is
+		// requested (default behavior is a redirect to URI with trailing slash)
+		mux.Handle(withoutSlash, http.NotFoundHandler())
+	}
 }
 
 func (res *resource) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	name, err := filepath.Rel(res.Path, r.URL.Path)
+	var reader io.ReadCloser
+	if err == nil {
+		reader, err = res.Open(name)
+	}
+	if err != nil {
+		if os.IsNotExist(err) {
+			http.NotFound(w, r)
+		} else {
+			http.Error(w, fmt.Sprintf("failed to open file %q: %v", r.URL.Path, err), http.StatusInternalServerError)
+		}
+		return
+	}
+	defer func() {
+		_ = reader.Close()
+	}()
+
 	etag := r.Header.Get("If-None-Match")
 	if etag != "" && etag == res.ETag {
 		w.WriteHeader(http.StatusNotModified)
 		return
 	}
-	w.Header().Set("Content-Type", res.ContentType)
+	ct := res.ContentType
+	if ct == "" {
+		ct = mime.TypeByExtension(path.Ext(r.URL.Path))
+	}
+	if ct != "" {
+		w.Header().Set("Content-Type", ct)
+	}
 	if res.Public {
 		w.Header().Set("Cache-Control", "public, max-age=3600")
 	} else {
 		w.Header().Set("Cache-Control", "private, max-age=3600")
 	}
-	w.Header().Set("ETag", res.ETag)
-	_, _ = w.Write(res.Data)
+	if res.ETag != "" {
+		w.Header().Set("ETag", res.ETag)
+	}
+	if res.Len > 0 {
+		w.Header().Set("Content-Length", strconv.Itoa(res.Len))
+	}
+	_, _ = io.Copy(w, reader)
 }
 
 // AsHTMLTag returns an HTML string corresponding to a tag that would load this resource (by inspecting ContentType).