package middleware
import (
"net/http"
"strconv"
"strings"
"testing"
expect "github.com/yusing/go-proxy/internal/utils/testing"
)
func TestInjectCSS(t *testing.T) {
opts := OptionsRaw{
"target": "head",
"html": "",
}
result, err := newMiddlewareTest(ModifyHTML, &testArgs{
middlewareOpt: opts,
respHeaders: http.Header{
"Content-Type": []string{"text/html; charset=utf-8"},
},
respBody: []byte(`
Unclosed div
Some content`),
})
expect.NoError(t, err)
// Should handle malformed HTML gracefully
expect.True(t, strings.Contains(string(result.Data), "Valid injection"), "Should inject content even with malformed HTML")
}
func TestInjectHTML_ContentTypes(t *testing.T) {
testCases := []struct {
name string
contentType string
shouldModify bool
}{
{"HTML with charset", "text/html; charset=utf-8", true},
{"Plain HTML", "text/html", true},
{"XHTML", "application/xhtml+xml", true},
{"JSON", "application/json", false},
{"Plain text", "text/plain", false},
{"JavaScript", "application/javascript", false},
{"CSS", "text/css", false},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
opts := OptionsRaw{
"target": "body",
"html": "
Test injection
",
}
originalBody := []byte(`Original content`)
result, err := newMiddlewareTest(ModifyHTML, &testArgs{
middlewareOpt: opts,
respHeaders: http.Header{
"Content-Type": []string{tc.contentType},
},
respBody: originalBody,
})
expect.NoError(t, err)
if tc.shouldModify {
expect.True(t, strings.Contains(string(result.Data), "Test injection"),
"Should modify HTML content for content-type: %s", tc.contentType)
} else {
expect.Equal(t, string(result.Data), string(originalBody),
"Should not modify non-HTML content for content-type: %s", tc.contentType)
}
})
}
}
func TestInjectHTML_ReplaceTrue(t *testing.T) {
opts := OptionsRaw{
"target": "body",
"html": "
Replacement content
",
"replace": true,
}
result, err := newMiddlewareTest(ModifyHTML, &testArgs{
middlewareOpt: opts,
respHeaders: http.Header{
"Content-Type": []string{"text/html"},
},
respBody: []byte(`
Test
Original content
More original content
`),
})
expect.NoError(t, err)
expect.Equal(t, removeTabsAndNewlines(result.Data), removeTabsAndNewlines(`
Test
Replacement content
`))
contentLength, _ := strconv.Atoi(result.ResponseHeaders.Get("Content-Length"))
expect.Equal(t, contentLength, len(result.Data), "Content-Length should be updated")
}
func TestInjectHTML_ReplaceVsAppend(t *testing.T) {
originalBody := []byte(`
Original content
`)
// Test append behavior (default)
appendOpts := OptionsRaw{
"target": ".target",
"html": "
Added content",
}
appendResult, err := newMiddlewareTest(ModifyHTML, &testArgs{
middlewareOpt: appendOpts,
respHeaders: http.Header{
"Content-Type": []string{"text/html"},
},
respBody: originalBody,
})
expect.NoError(t, err)
expect.Equal(t, removeTabsAndNewlines(appendResult.Data), removeTabsAndNewlines(`
Original contentAdded content
`))
// Test replace behavior
replaceOpts := OptionsRaw{
"target": ".target",
"html": "
Replacement content",
"replace": true,
}
replaceResult, err := newMiddlewareTest(ModifyHTML, &testArgs{
middlewareOpt: replaceOpts,
respHeaders: http.Header{
"Content-Type": []string{"text/html"},
},
respBody: originalBody,
})
expect.NoError(t, err)
expect.Equal(t, removeTabsAndNewlines(replaceResult.Data), removeTabsAndNewlines(`
Replacement content
`))
}
func TestInjectHTML_ReplaceWithDifferentSelectors(t *testing.T) {
testCases := []struct {
name string
selector string
original string
html string
expected string
}{
{
name: "ID selector replace",
selector: "#main",
original: `
Original content
`,
html: "
Replaced by ID",
expected: `
Replaced by ID`,
},
{
name: "Class selector replace",
selector: ".highlight",
original: `
Original content
`,
html: "
Replaced by class",
expected: `
Replaced by class`,
},
{
name: "Element selector replace",
selector: "body",
original: `Original content`,
html: "
Replaced body content",
expected: `
Replaced body content`,
},
{
name: "Attribute selector replace",
selector: "[data-test='target']",
original: `
Original content
`,
html: "
Replaced by attribute",
expected: `
Replaced by attribute`,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
opts := OptionsRaw{
"target": tc.selector,
"html": tc.html,
"replace": true,
}
result, err := newMiddlewareTest(ModifyHTML, &testArgs{
middlewareOpt: opts,
respHeaders: http.Header{
"Content-Type": []string{"text/html"},
},
respBody: []byte(`` + tc.original + ``),
})
expect.NoError(t, err)
expect.Equal(t, removeTabsAndNewlines(result.Data), removeTabsAndNewlines(tc.expected))
})
}
}
func TestInjectHTML_ReplaceWithEmpty(t *testing.T) {
opts := OptionsRaw{
"target": ".content",
"html": "",
"replace": true,
}
result, err := newMiddlewareTest(ModifyHTML, &testArgs{
middlewareOpt: opts,
respHeaders: http.Header{
"Content-Type": []string{"text/html"},
},
respBody: []byte(`
Content to be cleared
`),
})
expect.NoError(t, err)
expect.Equal(t, removeTabsAndNewlines(result.Data), removeTabsAndNewlines(`
`))
}
func TestInjectHTML_ReplaceMultipleTargets(t *testing.T) {
opts := OptionsRaw{
"target": ".container",
"html": "
Replaced content
",
"replace": true,
}
result, err := newMiddlewareTest(ModifyHTML, &testArgs{
middlewareOpt: opts,
respHeaders: http.Header{
"Content-Type": []string{"text/html"},
},
respBody: []byte(`
First container content
Second container content
`),
})
expect.NoError(t, err)
// Should only replace the first matching element
expectedContent := `
Replaced content
Replaced content
`
expect.Equal(t, removeTabsAndNewlines(result.Data), removeTabsAndNewlines(expectedContent))
}
func TestInjectHTML_ReplaceComplexHTML(t *testing.T) {
opts := OptionsRaw{
"target": "main",
"html": `
New Section
This replaces the entire main content.
`,
"replace": true,
}
result, err := newMiddlewareTest(ModifyHTML, &testArgs{
middlewareOpt: opts,
respHeaders: http.Header{
"Content-Type": []string{"text/html; charset=utf-8"},
},
respBody: []byte(`
Complex Page
Original Title
Original content that will be replaced.
More original content
`),
})
expect.NoError(t, err)
resultStr := removeTabsAndNewlines(result.Data)
expect.Equal(t, resultStr, removeTabsAndNewlines(`
Complex Page
New Section
This replaces the entire main content.
`))
contentLength, _ := strconv.Atoi(result.ResponseHeaders.Get("Content-Length"))
expect.Equal(t, contentLength, len(result.Data), "Content-Length should be updated correctly")
}
func removeTabsAndNewlines[T string | []byte](s T) string {
replacer := strings.NewReplacer(
"\n", "",
"\r", "",
"\t", "",
)
return replacer.Replace(string(s))
}