[assetserver] Inject runtime/IPC into all index html files (#2203)

It will also be injected into all html files returned when requesting a folder path.

v2/pkg/assetserver/assetserver.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"math/rand"
 	"net/http"
-	"net/http/httptest"
 	"strings"
 
 	"golang.org/x/net/html"
@@ -111,23 +110,60 @@ func (d *AssetServer) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	header := rw.Header()
 	if d.servingFromDisk {
-		header.Add(HeaderCacheControl, "no-cache")
+		rw.Header().Add(HeaderCacheControl, "no-cache")
+	}
+
+	handler := d.handler
+	if req.Method != http.MethodGet {
+		handler.ServeHTTP(rw, req)
+		return
 	}
 
 	path := req.URL.Path
-	switch path {
-	case "", "/", "/index.html":
-		recorder := httptest.NewRecorder()
-		d.handler.ServeHTTP(recorder, req)
-		for k, v := range recorder.HeaderMap {
-			header[k] = v
+	if path == runtimeJSPath {
+		d.writeBlob(rw, path, d.runtimeJS)
+
+	} else if path == runtimePath && d.runtimeHandler != nil {
+		d.runtimeHandler.HandleRuntimeCall(rw, req)
+
+	} else if path == ipcJSPath {
+		content := d.runtime.DesktopIPC()
+		if d.ipcJS != nil {
+			content = d.ipcJS(req)
+		}
+		d.writeBlob(rw, path, content)
+
+	} else if script, ok := d.pluginScripts[path]; ok {
+		d.writeBlob(rw, path, []byte(script))
+
+	} else if d.isRuntimeInjectionMatch(path) {
+		recorder := &bodyRecorder{
+			ResponseWriter: rw,
+			doRecord: func(code int, h http.Header) bool {
+				if code == http.StatusNotFound {
+					return true
+				}
+
+				if code != http.StatusOK {
+					return false
+				}
+
+				return strings.Contains(h.Get(HeaderContentType), "text/html")
+			}}
+
+		handler.ServeHTTP(recorder, req)
+
+		body := recorder.Body()
+		if body == nil {
+			// The body has been streamed and not recorded, we are finished
+			return
 		}
 
-		switch recorder.Code {
+		code := recorder.Code()
+		switch code {
 		case http.StatusOK:
-			content, err := d.processIndexHTML(recorder.Body.Bytes())
+			content, err := d.processIndexHTML(body.Bytes())
 			if err != nil {
 				d.serveError(rw, err, "Unable to processIndexHTML")
 				return
@@ -138,34 +174,12 @@ func (d *AssetServer) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 			d.writeBlob(rw, indexHTML, defaultHTML)
 
 		default:
-			rw.WriteHeader(recorder.Code)
+			rw.WriteHeader(code)
 
 		}
 
-	case runtimeJSPath:
-		d.writeBlob(rw, path, d.runtimeJS)
-
-	case runtimePath:
-		if d.runtimeHandler != nil {
-			d.runtimeHandler.HandleRuntimeCall(rw, req)
-		} else {
-			d.handler.ServeHTTP(rw, req)
-		}
-
-	case ipcJSPath:
-		content := d.runtime.DesktopIPC()
-		if d.ipcJS != nil {
-			content = d.ipcJS(req)
-		}
-		d.writeBlob(rw, path, content)
-
-	default:
-		// Check if this is a plugin script
-		if script, ok := d.pluginScripts[path]; ok {
-			d.writeBlob(rw, path, []byte(script))
-			return
-		}
-		d.handler.ServeHTTP(rw, req)
+	} else {
+		handler.ServeHTTP(rw, req)
 	}
 }
 
@@ -229,3 +243,12 @@ func (d *AssetServer) logError(message string, args ...interface{}) {
 		d.logger.Error("[AssetServer] "+message, args...)
 	}
 }
+
+func (AssetServer) isRuntimeInjectionMatch(path string) bool {
+	if path == "" {
+		path = "/"
+	}
+
+	return strings.HasSuffix(path, "/") ||
+		strings.HasSuffix(path, "/"+indexHTML)
+}

v2/pkg/assetserver/body_recorder.go

@@ -0,0 +1,61 @@
+package assetserver
+
+import (
+	"bytes"
+	"net/http"
+)
+
+type bodyRecorder struct {
+	http.ResponseWriter
+	doRecord func(code int, header http.Header) bool
+
+	body        *bytes.Buffer
+	code        int
+	wroteHeader bool
+}
+
+func (rw *bodyRecorder) Write(buf []byte) (int, error) {
+	rw.writeHeader(buf, http.StatusOK)
+	if rw.body != nil {
+		return rw.body.Write(buf)
+	}
+	return rw.ResponseWriter.Write(buf)
+}
+
+func (rw *bodyRecorder) WriteHeader(code int) {
+	rw.writeHeader(nil, code)
+}
+
+func (rw *bodyRecorder) Code() int {
+	return rw.code
+}
+
+func (rw *bodyRecorder) Body() *bytes.Buffer {
+	return rw.body
+}
+
+func (rw *bodyRecorder) writeHeader(buf []byte, code int) {
+	if rw.wroteHeader {
+		return
+	}
+
+	if rw.doRecord != nil {
+		header := rw.Header()
+		if len(buf) != 0 {
+			if _, hasType := header[HeaderContentType]; !hasType {
+				header.Set(HeaderContentType, http.DetectContentType(buf))
+			}
+		}
+
+		if rw.doRecord(code, header) {
+			rw.body = bytes.NewBuffer(nil)
+		}
+	}
+
+	if rw.body == nil {
+		rw.ResponseWriter.WriteHeader(code)
+	}
+
+	rw.code = code
+	rw.wroteHeader = true
+}

website/src/pages/changelog.mdx

@@ -25,6 +25,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - New task created for linting v2 `task v2:lint`. Workflow updated to run the task. Added by @mikeee in [PR](https://github.com/wailsapp/wails/pull/2957)
 - Added new community template wails-htmx-templ-chi-tailwind. Added by [@pylotlight](https://github.com/pylotlight) in [PR](https://github.com/wailsapp/wails/pull/2984)
 - Added CPU/GPU/Memory detection for `wails doctor`. Added by @leaanthony in #d51268b8d0680430f3a614775b13e6cd2b906d1c
+- The [AssetServer](/docs/reference/options#assetserver) now injects the runtime/IPC into all index html files and into all html files returned when requesting a folder path. Added by @stffabi in [PR](https://github.com/wailsapp/wails/pull/2203)
 
 ### Changed