package proxy
import (
"bytes"
"compress/flate"
"compress/gzip"
"io"
"net/http"
"strings"
"testing"
"github.com/andybalholm/brotli"
"github.com/klauspost/compress/zstd"
)
func TestModifyResponse_GzipCompression(t *testing.T) {
ps := &ProxyServer{
ListenAddr: ":8080",
}
htmlContent := []byte("<html><head><title>Test</title></head><body>Hello World</body></html>")
// Compress the HTML content
var buf bytes.Buffer
gzWriter := gzip.NewWriter(&buf)
_, err := gzWriter.Write(htmlContent)
if err != nil {
t.Fatalf("failed to compress: %v", err)
}
gzWriter.Close()
// Create a response with gzip-compressed content
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/html"},
"Content-Encoding": []string{"gzip"},
},
Body: io.NopCloser(&buf),
}
// Modify the response
err = ps.modifyResponse(resp)
if err != nil {
t.Fatalf("modifyResponse failed: %v", err)
}
// Read the modified response
modifiedBody, err := io.ReadAll(resp.Body)
if err != nil {
t.Fatalf("failed to read modified body: %v", err)
}
// Verify it's no longer compressed (Content-Encoding should be removed)
if resp.Header.Get("Content-Encoding") != "" {
t.Errorf("Content-Encoding header should be removed, got: %s", resp.Header.Get("Content-Encoding"))
}
// Verify the content is valid HTML (not binary garbage)
bodyStr := string(modifiedBody)
if !strings.Contains(bodyStr, "<html>") {
t.Errorf("modified body doesn't contain <html> tag")
}
if !strings.Contains(bodyStr, "Hello World") {
t.Errorf("modified body doesn't contain original content")
}
// Verify instrumentation was injected
if !strings.Contains(bodyStr, "__devtool") {
t.Errorf("instrumentation script was not injected")
}
if !strings.Contains(bodyStr, "WebSocket") {
t.Errorf("WebSocket code was not injected")
}
}
func TestModifyResponse_DeflateCompression(t *testing.T) {
ps := &ProxyServer{
ListenAddr: ":8080",
}
htmlContent := []byte("<html><head><title>Test</title></head><body>Hello World</body></html>")
// Compress the HTML content with deflate
var buf bytes.Buffer
flateWriter, err := flate.NewWriter(&buf, flate.DefaultCompression)
if err != nil {
t.Fatalf("failed to create deflate writer: %v", err)
}
_, err = flateWriter.Write(htmlContent)
if err != nil {
t.Fatalf("failed to compress: %v", err)
}
flateWriter.Close()
// Create a response with deflate-compressed content
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/html"},
"Content-Encoding": []string{"deflate"},
},
Body: io.NopCloser(&buf),
}
// Modify the response
err = ps.modifyResponse(resp)
if err != nil {
t.Fatalf("modifyResponse failed: %v", err)
}
// Read the modified response
modifiedBody, err := io.ReadAll(resp.Body)
if err != nil {
t.Fatalf("failed to read modified body: %v", err)
}
// Verify it's no longer compressed
if resp.Header.Get("Content-Encoding") != "" {
t.Errorf("Content-Encoding header should be removed, got: %s", resp.Header.Get("Content-Encoding"))
}
// Verify the content is valid HTML
bodyStr := string(modifiedBody)
if !strings.Contains(bodyStr, "<html>") {
t.Errorf("modified body doesn't contain <html> tag")
}
if !strings.Contains(bodyStr, "Hello World") {
t.Errorf("modified body doesn't contain original content")
}
// Verify instrumentation was injected
if !strings.Contains(bodyStr, "__devtool") {
t.Errorf("instrumentation script was not injected")
}
}
func TestModifyResponse_NoCompression(t *testing.T) {
ps := &ProxyServer{
ListenAddr: ":8080",
}
htmlContent := []byte("<html><head><title>Test</title></head><body>Hello World</body></html>")
// Create a response with uncompressed content
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/html"},
},
Body: io.NopCloser(bytes.NewReader(htmlContent)),
}
// Modify the response
err := ps.modifyResponse(resp)
if err != nil {
t.Fatalf("modifyResponse failed: %v", err)
}
// Read the modified response
modifiedBody, err := io.ReadAll(resp.Body)
if err != nil {
t.Fatalf("failed to read modified body: %v", err)
}
// Verify the content is valid HTML
bodyStr := string(modifiedBody)
if !strings.Contains(bodyStr, "<html>") {
t.Errorf("modified body doesn't contain <html> tag")
}
if !strings.Contains(bodyStr, "Hello World") {
t.Errorf("modified body doesn't contain original content")
}
// Verify instrumentation was injected
if !strings.Contains(bodyStr, "__devtool") {
t.Errorf("instrumentation script was not injected")
}
}
func TestModifyResponse_NonHTMLSkipped(t *testing.T) {
ps := &ProxyServer{
ListenAddr: ":8080",
}
jsonContent := []byte(`{"message": "Hello World"}`)
// Create a response with JSON content (should not be modified)
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"application/json"},
},
Body: io.NopCloser(bytes.NewReader(jsonContent)),
}
// Modify the response
err := ps.modifyResponse(resp)
if err != nil {
t.Fatalf("modifyResponse failed: %v", err)
}
// Read the response
modifiedBody, err := io.ReadAll(resp.Body)
if err != nil {
t.Fatalf("failed to read modified body: %v", err)
}
// Verify JSON content is unchanged
if !bytes.Equal(modifiedBody, jsonContent) {
t.Errorf("JSON content should be unchanged, got: %s", string(modifiedBody))
}
// Verify instrumentation was NOT injected
if strings.Contains(string(modifiedBody), "__devtool") {
t.Errorf("instrumentation should not be injected in non-HTML responses")
}
}
func TestModifyResponse_CorruptGzipData(t *testing.T) {
ps := &ProxyServer{
ListenAddr: ":8080",
}
// Create corrupt gzip data
corruptData := []byte("this is not valid gzip data")
// Create a response with corrupt gzip content
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/html"},
"Content-Encoding": []string{"gzip"},
},
Body: io.NopCloser(bytes.NewReader(corruptData)),
}
// Modify the response (should not return error, just skip injection)
err := ps.modifyResponse(resp)
if err != nil {
t.Errorf("modifyResponse should handle corrupt gzip gracefully, got error: %v", err)
}
}
// Benchmark decompression overhead
func BenchmarkModifyResponse_Gzip(b *testing.B) {
ps := &ProxyServer{
ListenAddr: ":8080",
}
htmlContent := []byte("<html><head><title>Test</title></head><body>" + strings.Repeat("Hello World ", 1000) + "</body></html>")
// Pre-compress content
var buf bytes.Buffer
gzWriter := gzip.NewWriter(&buf)
gzWriter.Write(htmlContent)
gzWriter.Close()
compressedData := buf.Bytes()
b.ResetTimer()
for i := 0; i < b.N; i++ {
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/html"},
"Content-Encoding": []string{"gzip"},
},
Body: io.NopCloser(bytes.NewReader(compressedData)),
}
ps.modifyResponse(resp)
io.ReadAll(resp.Body)
}
}
func BenchmarkModifyResponse_NoCompression(b *testing.B) {
ps := &ProxyServer{
ListenAddr: ":8080",
}
htmlContent := []byte("<html><head><title>Test</title></head><body>" + strings.Repeat("Hello World ", 1000) + "</body></html>")
b.ResetTimer()
for i := 0; i < b.N; i++ {
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/html"},
},
Body: io.NopCloser(bytes.NewReader(htmlContent)),
}
ps.modifyResponse(resp)
io.ReadAll(resp.Body)
}
}
func TestModifyResponse_BrotliCompression(t *testing.T) {
ps := &ProxyServer{
ListenAddr: ":8080",
}
htmlContent := []byte("<html><head><title>Test</title></head><body>Hello World</body></html>")
// Compress the HTML content with Brotli
var buf bytes.Buffer
brWriter := brotli.NewWriter(&buf)
_, err := brWriter.Write(htmlContent)
if err != nil {
t.Fatalf("failed to compress: %v", err)
}
brWriter.Close()
// Create a response with brotli-compressed content
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/html"},
"Content-Encoding": []string{"br"},
},
Body: io.NopCloser(&buf),
}
// Modify the response
err = ps.modifyResponse(resp)
if err != nil {
t.Fatalf("modifyResponse failed: %v", err)
}
// Read the modified response
modifiedBody, err := io.ReadAll(resp.Body)
if err != nil {
t.Fatalf("failed to read modified body: %v", err)
}
// Verify it's no longer compressed
if resp.Header.Get("Content-Encoding") != "" {
t.Errorf("Content-Encoding header should be removed, got: %s", resp.Header.Get("Content-Encoding"))
}
// Verify the content is valid HTML
bodyStr := string(modifiedBody)
if !strings.Contains(bodyStr, "<html>") {
t.Errorf("modified body doesn't contain <html> tag")
}
if !strings.Contains(bodyStr, "Hello World") {
t.Errorf("modified body doesn't contain original content")
}
// Verify instrumentation was injected
if !strings.Contains(bodyStr, "__devtool") {
t.Errorf("instrumentation script was not injected")
}
}
func TestModifyResponse_ZstdCompression(t *testing.T) {
ps := &ProxyServer{
ListenAddr: ":8080",
}
htmlContent := []byte("<html><head><title>Test</title></head><body>Hello World</body></html>")
// Compress the HTML content with Zstandard
var buf bytes.Buffer
zstdWriter, err := zstd.NewWriter(&buf)
if err != nil {
t.Fatalf("failed to create zstd writer: %v", err)
}
_, err = zstdWriter.Write(htmlContent)
if err != nil {
t.Fatalf("failed to compress: %v", err)
}
zstdWriter.Close()
// Create a response with zstd-compressed content
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/html"},
"Content-Encoding": []string{"zstd"},
},
Body: io.NopCloser(&buf),
}
// Modify the response
err = ps.modifyResponse(resp)
if err != nil {
t.Fatalf("modifyResponse failed: %v", err)
}
// Read the modified response
modifiedBody, err := io.ReadAll(resp.Body)
if err != nil {
t.Fatalf("failed to read modified body: %v", err)
}
// Verify it's no longer compressed
if resp.Header.Get("Content-Encoding") != "" {
t.Errorf("Content-Encoding header should be removed, got: %s", resp.Header.Get("Content-Encoding"))
}
// Verify the content is valid HTML
bodyStr := string(modifiedBody)
if !strings.Contains(bodyStr, "<html>") {
t.Errorf("modified body doesn't contain <html> tag")
}
if !strings.Contains(bodyStr, "Hello World") {
t.Errorf("modified body doesn't contain original content")
}
// Verify instrumentation was injected
if !strings.Contains(bodyStr, "__devtool") {
t.Errorf("instrumentation script was not injected")
}
}
func TestModifyResponse_UnsupportedCompression(t *testing.T) {
ps := &ProxyServer{
ListenAddr: ":8080",
}
// Create a response with an unsupported compression format
htmlContent := []byte("<html><head><title>Test</title></head><body>Hello World</body></html>")
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/html"},
"Content-Encoding": []string{"lzma"},
},
Body: io.NopCloser(bytes.NewReader(htmlContent)),
}
// Modify the response (should pass through without modification)
err := ps.modifyResponse(resp)
if err != nil {
t.Fatalf("modifyResponse should handle unsupported encoding gracefully, got error: %v", err)
}
// Read the response
modifiedBody, err := io.ReadAll(resp.Body)
if err != nil {
t.Fatalf("failed to read body: %v", err)
}
// Verify content is unchanged (no decompression attempted)
if !bytes.Equal(modifiedBody, htmlContent) {
t.Errorf("body should be unchanged for unsupported encoding")
}
// Verify instrumentation was NOT injected (since we can't decode)
if strings.Contains(string(modifiedBody), "__devtool") {
t.Errorf("instrumentation should not be injected for unsupported encoding")
}
}
func BenchmarkModifyResponse_Brotli(b *testing.B) {
ps := &ProxyServer{
ListenAddr: ":8080",
}
htmlContent := []byte("<html><head><title>Test</title></head><body>" + strings.Repeat("Hello World ", 1000) + "</body></html>")
// Pre-compress content with Brotli
var buf bytes.Buffer
brWriter := brotli.NewWriter(&buf)
brWriter.Write(htmlContent)
brWriter.Close()
compressedData := buf.Bytes()
b.ResetTimer()
for i := 0; i < b.N; i++ {
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/html"},
"Content-Encoding": []string{"br"},
},
Body: io.NopCloser(bytes.NewReader(compressedData)),
}
ps.modifyResponse(resp)
io.ReadAll(resp.Body)
}
}
func BenchmarkModifyResponse_Zstd(b *testing.B) {
ps := &ProxyServer{
ListenAddr: ":8080",
}
htmlContent := []byte("<html><head><title>Test</title></head><body>" + strings.Repeat("Hello World ", 1000) + "</body></html>")
// Pre-compress content with Zstandard
var buf bytes.Buffer
zstdWriter, _ := zstd.NewWriter(&buf)
zstdWriter.Write(htmlContent)
zstdWriter.Close()
compressedData := buf.Bytes()
b.ResetTimer()
for i := 0; i < b.N; i++ {
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"text/html"},
"Content-Encoding": []string{"zstd"},
},
Body: io.NopCloser(bytes.NewReader(compressedData)),
}
ps.modifyResponse(resp)
io.ReadAll(resp.Body)
}
}
