Repository: a-h/templ (Go)

Issue:
Bug: hot reloader adds a wrong reload script in nested body responses
I'll be much more clear if I give an example:

if the server responds this:

```html
<!doctype html>
<html lang="en">
   <head>
      <meta charset="UTF-8">
      <meta name="viewport" content="width=device-width, initial-scale=1.0">
      <title>HTML Playground</title>
      <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/codemirror/lib/codemirror.css">
      <style>
         body { display: flex; height: 100vh; }
         #editors { flex: 1; display: flex; flex-direction: column; }
         .editor { flex: 1; }
         iframe { flex: 1; border: 1px solid #ccc; }
      </style>
   </head>
   <body>
      <div id="editors"><textarea id="html-editor" class="editor"></textarea> <textarea id="css-editor" class="editor"></textarea> <textarea id="js-editor" class="editor"></textarea></div>
      <iframe id="preview"></iframe><script src="https://cdn.jsdelivr.net/npm/codemirror/lib/codemirror.js"></script><script src="https://cdn.jsdelivr.net/npm/codemirror/mode/htmlmixed/htmlmixed.js"></script><script src="https://cdn.jsdelivr.net/npm/codemirror/mode/css/css.js"></script><script src="https://cdn.jsdelivr.net/npm/codemirror/mode/javascript/javascript.js"></script><script>
         document.addEventListener('DOMContentLoaded', () => {
           const htmlEditor = CodeMirror.fromTextArea(document.getElementById('html-editor'), { mode: 'htmlmixed', lineNumbers: true });
           const cssEditor = CodeMirror.fromTextArea(document.getElementById('css-editor'), { mode: 'css', lineNumbers: true });
           const jsEditor = CodeMirror.fromTextArea(document.getElementById('js-editor'), { mode: 'javascript', lineNumbers: true });
         
           const previewFrame = document.getElementById('preview');
           const updatePreview = () => {
         	const html = htmlEditor.getValue();
         	const css = cssEditor.getValue();
         	const js = jsEditor.getValue();
         	const previewDoc = previewFrame.contentDocument || previewFrame.contentWindow.document;
         	previewDoc.open();
         	previewDoc.write(`
         	  <html>
         		<head><style>${css}</style></head>
         		<body>${html}<script>${js}<\/script></body>
         	  </html>
         	`);
         	previewDoc.close();
           };
         
           htmlEditor.on('change', updatePreview);
           cssEditor.on('change', updatePreview);
           jsEditor.on('change', updatePreview);
         
           // Initial Preview
           updatePreview();
         });
          
      </script>
   </body>
</html>
```

The proxy will rewrite it as this:

```html
<!doctype html>
<html lang="en">
   <head>
      <meta charset="UTF-8">
      <meta name="viewport" content="width=device-width, initial-scale=1.0">
      <title>HTML Playground</title>
      <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/codemirror/lib/codemirror.css">
      <style>
         body { display: flex; height: 100vh; }
         #editors { flex: 1; display: flex; flex-direction: column; }
         .editor { flex: 1; }
         iframe { flex: 1; border: 1px solid #ccc; }
      </style>
   </head>
   <body>
      <div id="editors"><textarea id="html-editor" class="editor"></textarea> <textarea id="css-editor" class="editor"></textarea> <textarea id="js-editor" class="editor"></textarea></div>
      <iframe id="preview"></iframe><script src="https://cdn.jsdelivr.net/npm/codemirror/lib/codemirror.js"></script><script src="https://cdn.jsdelivr.net/npm/codemirror/mode/htmlmixed/htmlmixed.js"></script><script src="https://cdn.jsdelivr.net/npm/codemirror/mode/css/css.js"></script><script src="https://cdn.jsdelivr.net/npm/codemirror/mode/javascript/javascript.js"></script><script>
         document.addEventListener('DOMContentLoaded', () => {
           const htmlEditor = CodeMirror.fromTextArea(document.getElementById('html-editor'), { mode: 'htmlmixed', lineNumbers: true });
           const cssEditor = CodeMirror.fromTextArea(document.getElementById('css-editor'), { mode: 'css', lineNumbers: true });
           const jsEditor = CodeMirror.fromTextArea(document.getElementById('js-editor'), { mode: 'javascript', lineNumbers: true });
         
           const previewFrame = document.getElementById('preview');
           const updatePreview = () => {
         	const html = htmlEditor.getValue();
         	const css = cssEditor.getValue();
         	const js = jsEditor.getValue();
         	const previewDoc = previewFrame.contentDocument || previewFrame.contentWindow.document;
         	previewDoc.open();
         	previewDoc.write(`
         	  <html>
         		<head><style>${css}</style></head>
         		<body>${html}<script>${js}<\/script><script src="/_templ/reload/script.js">
      </script>
   </body>
</html>
`);
previewDoc.close();
};
htmlEditor.on('change', updatePreview);
cssEditor.on('change', updatePreview);
jsEditor.on('change', updatePreview);
// Initial Preview
updatePreview();
});
</script><script src="/_templ/reload/script.js"></script></body></html>
```

To solve this I have to manually break the inner body tag, which is not ideal.



Discussion:
Yes, that makes sense. The proxy may have to be updated to parse the HTML in order to find the correct location to place the script. At the moment it just looks for `</body>` in the HTML.



Patch:
diff --git a/cmd/templ/generatecmd/proxy/proxy.go b/cmd/templ/generatecmd/proxy/proxy.go
index 9b942be18..a129888c9 100644
--- a/cmd/templ/generatecmd/proxy/proxy.go
+++ b/cmd/templ/generatecmd/proxy/proxy.go
@@ -16,6 +16,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/PuerkitoBio/goquery"
 	"github.com/a-h/templ/cmd/templ/generatecmd/sse"
 	"github.com/andybalholm/brotli"
 
@@ -36,7 +37,16 @@ type Handler struct {
 }
 
 func insertScriptTagIntoBody(body string) (updated string) {
-	return strings.Replace(body, "</body>", scriptTag+"</body>", -1)
+	doc, err := goquery.NewDocumentFromReader(strings.NewReader(body))
+	if err != nil {
+		return strings.Replace(body, "</body>", scriptTag+"</body>", -1)
+	}
+	doc.Find("body").AppendHtml(scriptTag)
+	r, err := doc.Html()
+	if err != nil {
+		return strings.Replace(body, "</body>", scriptTag+"</body>", -1)
+	}
+	return r
 }
 
 type passthroughWriteCloser struct {


Test patch:
diff --git a/cmd/templ/generatecmd/proxy/proxy_test.go b/cmd/templ/generatecmd/proxy/proxy_test.go
index b33478fca..d2bc91517 100644
--- a/cmd/templ/generatecmd/proxy/proxy_test.go
+++ b/cmd/templ/generatecmd/proxy/proxy_test.go
@@ -161,6 +161,49 @@ func TestProxy(t *testing.T) {
 			t.Errorf("unexpected response body (-got +want):\n%s", diff)
 		}
 	})
+	t.Run("plain: body tags get the script inserted ignoring js with body tags", func(t *testing.T) {
+		// Arrange
+		r := &http.Response{
+			Body:   io.NopCloser(strings.NewReader(`<html><body><script>console.log("<body></body>")</script></body></html>`)),
+			Header: make(http.Header),
+			Request: &http.Request{
+				URL: &url.URL{
+					Scheme: "http",
+					Host:   "example.com",
+				},
+			},
+		}
+		r.Header.Set("Content-Type", "text/html, charset=utf-8")
+		r.Header.Set("Content-Length", "26")
+
+		expectedString := insertScriptTagIntoBody(`<html><body><script>console.log("<body></body>")</script></body></html>`)
+		if !strings.Contains(expectedString, scriptTag) {
+			t.Fatalf("expected the script tag to be inserted, but it wasn't: %q", expectedString)
+		}
+		if !strings.Contains(expectedString, `console.log("<body></body>")`) {
+			t.Fatalf("expected the script tag to be inserted, but mangled the html: %q", expectedString)
+		}
+
+		// Act
+		log := slog.New(slog.NewJSONHandler(io.Discard, nil))
+		h := New(log, "127.0.0.1", 7474, &url.URL{Scheme: "http", Host: "example.com"})
+		err := h.modifyResponse(r)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+
+		// Assert
+		if r.Header.Get("Content-Length") != fmt.Sprintf("%d", len(expectedString)) {
+			t.Errorf("expected content length to be %d, got %v", len(expectedString), r.Header.Get("Content-Length"))
+		}
+		actualBody, err := io.ReadAll(r.Body)
+		if err != nil {
+			t.Fatalf("unexpected error reading response: %v", err)
+		}
+		if diff := cmp.Diff(expectedString, string(actualBody)); diff != "" {
+			t.Errorf("unexpected response body (-got +want):\n%s", diff)
+		}
+	})
 	t.Run("gzip: non-html content is not modified", func(t *testing.T) {
 		// Arrange
 		r := &http.Response{