feat: enhance HTML injection with body fallback and self-removing polyfill
- Modify InjectHTML to inject script into body if head is unavailable - Update polyfill.js to remove its script tag after initialization for DOM cleanup
lemon-mint committed
Nov 1, 2025 at 02:30 UTC
c1a8e56ab7ac50d95f27093d6286692f5f745556
2 files changed
+38
-17
cmd/webclient/inject.go
+29
-17
@@ -19,13 +19,17 @@ func InjectHTML(body []byte) []byte {
19
return body
20
}
21
22
- // Find the head element
22
+ // Find the head or body element
23
var head *html.Node
24
+ var bodyNode *html.Node
25
var crawler func(*html.Node)
26
crawler = func(node *html.Node) {
26
- if node.Type == html.ElementNode && node.Data == "head" {
27
- head = node
28
- return
27
+ if node.Type == html.ElementNode {
28
+ if node.Data == "head" {
29
+ head = node
30
+ } else if node.Data == "body" {
31
+ bodyNode = node
32
+ }
33
}
34
for child := node.FirstChild; child != nil; child = child.NextSibling {
35
crawler(child)
@@ -33,27 +37,35 @@ func InjectHTML(body []byte) []byte {
37
}
38
crawler(doc)
39
36
- if head != nil {
37
- // Create script element
38
- script := &html.Node{
39
- Type: html.ElementNode,
40
- Data: "script",
41
- Attr: []html.Attribute{},
42
- }
40
+ // Create script element
41
+ script := &html.Node{
42
+ Type: html.ElementNode,
43
+ Data: "script",
44
+ Attr: []html.Attribute{},
45
+ }
46
44
- // Add the script content
45
- scriptContent := &html.Node{
46
- Type: html.TextNode,
47
- Data: string(polyfillJS),
48
- }
49
- script.AppendChild(scriptContent)
47
+ // Add the script content
48
+ scriptContent := &html.Node{
49
+ Type: html.TextNode,
50
+ Data: string(polyfillJS),
51
+ }
52
+ script.AppendChild(scriptContent)
53
54
+ // Inject into head if available, otherwise into body
55
+ if head != nil {
56
// Insert as the first child of head
57
if head.FirstChild != nil {
58
head.InsertBefore(script, head.FirstChild)
59
} else {
60
head.AppendChild(script)
61
}
62
+ } else if bodyNode != nil {
63
+ // Insert as the first child of body if head doesn't exist
64
+ if bodyNode.FirstChild != nil {
65
+ bodyNode.InsertBefore(script, bodyNode.FirstChild)
66
+ } else {
67
+ bodyNode.AppendChild(script)
68
+ }
69
}
70
71
// Convert back to bytes
cmd/webclient/polyfill.js
+9
@@ -1,6 +1,9 @@
1
(function() {
2
'use strict';
3
4
+ // Capture reference to current script for later removal
5
+ const currentScript = document.currentScript;
6
+
7
// Save original WebSocket
8
const NativeWebSocket = window.WebSocket;
9
@@ -343,4 +346,10 @@
346
window.WebSocket.CLOSED = NativeWebSocket.CLOSED;
347
348
console.log('[WebSocket Polyfill] Initialized');
349
+
350
+ // Remove the polyfill script tag after initialization
351
+ if (currentScript && currentScript.parentNode) {
352
+ currentScript.parentNode.removeChild(currentScript);
353
+ console.log('[WebSocket Polyfill] Script tag removed');
354
+ }
355
})();
\ No newline at end of file