package main import ( "bytes" "compress/gzip" "io" "net/http" "strings" "testing" "github.com/klauspost/compress/zstd" "ublproxy/internal/blocklist" ) func TestBootstrapScriptTagWithSession(t *testing.T) { sm := newSessionMap() sm.Set("192.168.1.10", sessionEntry{Token: "test-token-abc", CredentialID: "cred-1"}) p := &proxyHandler{ sessions: sm, portalOrigin: "https://192.168.1.1:8443", } tag := p.bootstrapScriptTag("192.168.1.10", "example.com") if !strings.Contains(tag, "` + `` + `` + `

Content

` + `` resp := &http.Response{ StatusCode: 200, Header: http.Header{"Content-Type": []string{"text/html"}}, Body: io.NopCloser(strings.NewReader(htmlBody)), } _, stats := p.applyElementHiding(resp, "page.example.com", "127.0.0.1", false) if !stats.Modified { t.Fatal("expected modification") } if stats.Stripped != 2 { t.Errorf("Stripped = %d, want 2", stats.Stripped) } } func TestElementHidingStatsCombined(t *testing.T) { rs := blocklist.NewRuleSet() rs.AddLine("##.ad-banner") rs.AddLine("||ads.example.com^") p := &proxyHandler{sessions: newSessionMap()} p.baselineRules.Store(rs) htmlBody := `` + `` + `` + `
Ad
` + `

Content

` + `` resp := &http.Response{ StatusCode: 200, Header: http.Header{"Content-Type": []string{"text/html"}}, Body: io.NopCloser(strings.NewReader(htmlBody)), } _, stats := p.applyElementHiding(resp, "page.example.com", "127.0.0.1", false) if !stats.Modified { t.Fatal("expected modification") } if stats.Hidden != 1 { t.Errorf("Hidden = %d, want 1", stats.Hidden) } if stats.Stripped != 1 { t.Errorf("Stripped = %d, want 1", stats.Stripped) } want := "hidden=1; stripped=1" if got := stats.header(); got != want { t.Errorf("header() = %q, want %q", got, want) } } func TestElementHidingStatsNoModification(t *testing.T) { p := &proxyHandler{sessions: newSessionMap()} resp := &http.Response{ StatusCode: 200, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"ok":true}`)), } _, stats := p.applyElementHiding(resp, "example.com", "127.0.0.1", false) if stats.Modified { t.Error("should not modify non-HTML response") } if stats.Hidden != 0 || stats.Stripped != 0 { t.Errorf("stats should be zero for unmodified response, got hidden=%d stripped=%d", stats.Hidden, stats.Stripped) } } func TestElementHidingFiltersUnmatchedSelectors(t *testing.T) { rs := blocklist.NewRuleSet() rs.AddLine("##.ad-banner") // matches rs.AddLine("##.tracking-pixel") // no match — class not in HTML rs.AddLine("##.sponsored") // no match rs.AddLine("##.augl") // matches rs.AddLine("###slot-668") // matches rs.AddLine("##.nonexistent") // no match p := &proxyHandler{sessions: newSessionMap()} p.baselineRules.Store(rs) htmlBody := `` + `
Ad
` + `` + `

Content

` + `` resp := &http.Response{ StatusCode: 200, Header: http.Header{"Content-Type": []string{"text/html"}}, Body: io.NopCloser(strings.NewReader(htmlBody)), } modified, stats := p.applyElementHiding(resp, "example.com", "127.0.0.1", false) if !stats.Modified { t.Fatal("expected modification") } // Only 3 selectors match the HTML (ad-banner, augl, #slot-668) if stats.Hidden != 3 { t.Errorf("Hidden = %d, want 3", stats.Hidden) } body := string(modified) // Matching selectors should be in the CSS if !strings.Contains(body, ".ad-banner") { t.Error("CSS should contain .ad-banner") } if !strings.Contains(body, ".augl") { t.Error("CSS should contain .augl") } if !strings.Contains(body, "#slot-668") { t.Error("CSS should contain #slot-668") } // Non-matching selectors should NOT be in the CSS if strings.Contains(body, ".tracking-pixel") { t.Error("CSS should NOT contain .tracking-pixel (not in HTML)") } if strings.Contains(body, ".sponsored") { t.Error("CSS should NOT contain .sponsored (not in HTML)") } if strings.Contains(body, ".nonexistent") { t.Error("CSS should NOT contain .nonexistent (not in HTML)") } } func TestScriptletInjection(t *testing.T) { rs := blocklist.NewRuleSet() rs.AddLine("example.com##+js(nowebrtc)") rs.AddLine("example.com##+js(set-constant, ads.enabled, true)") p := &proxyHandler{sessions: newSessionMap()} p.baselineRules.Store(rs) htmlBody := `Test

Content

` resp := &http.Response{ StatusCode: 200, Header: http.Header{"Content-Type": []string{"text/html"}}, Body: io.NopCloser(strings.NewReader(htmlBody)), } modified, stats := p.applyElementHiding(resp, "example.com", "127.0.0.1", false) if !stats.Modified { t.Fatal("expected modification for scriptlet injection") } body := string(modified) // Both scriptlets should be injected as which must be escaped rs.AddLine(`example.com##+js(set-constant, x, true)`) p := &proxyHandler{sessions: newSessionMap()} p.baselineRules.Store(rs) htmlBody := `

Content

` resp := &http.Response{ StatusCode: 200, Header: http.Header{"Content-Type": []string{"text/html"}}, Body: io.NopCloser(strings.NewReader(htmlBody)), } modified, _ := p.applyElementHiding(resp, "example.com", "127.0.0.1", false) body := string(modified) // Must not contain raw inside the scriptlet injection // (the closing of the wrapper tag is OK, but not inside the JS) scriptStart := strings.Index(body, "") if scriptStart >= 0 && scriptEnd > scriptStart { jsContent := body[scriptStart+len("") } } } func TestScriptletInjectionWithCSSRules(t *testing.T) { rs := blocklist.NewRuleSet() rs.AddLine("example.com##+js(nowebrtc)") rs.AddLine("##.ad-banner") p := &proxyHandler{sessions: newSessionMap()} p.baselineRules.Store(rs) htmlBody := `
Ad

Content

` resp := &http.Response{ StatusCode: 200, Header: http.Header{"Content-Type": []string{"text/html"}}, Body: io.NopCloser(strings.NewReader(htmlBody)), } modified, stats := p.applyElementHiding(resp, "example.com", "127.0.0.1", false) if !stats.Modified { t.Fatal("expected modification") } body := string(modified) // Both CSS and scriptlet should be present if !strings.Contains(body, "