Harden scan robustness and tests

2026-09-24 12:11:52 +08:00 · 2026-06-14 22:23:48 +08:00
parent 5ad914a1bb
commit c49c23c7f0
100 changed files with 4483 additions and 412 deletions
@@ -0,0 +1,53 @@
+package web
+
+import (
+	"context"
+	"testing"
+
+	"github.com/shadow1ng/fscan/common"
+)
+
+func TestMatchCDNorWAF(t *testing.T) {
+	tests := []struct {
+		name         string
+		fingerprints []string
+		want         string
+	}{
+		{name: "empty", fingerprints: nil, want: ""},
+		{name: "no match", fingerprints: []string{"nginx", "wordpress"}, want: ""},
+		{name: "case insensitive cdn", fingerprints: []string{"site behind cloudflare"}, want: "CloudFlare"},
+		{name: "chinese waf", fingerprints: []string{"命中安全狗防护"}, want: "安全狗"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := matchCDNorWAF(tt.fingerprints); got != tt.want {
+				t.Fatalf("matchCDNorWAF(%v) = %q, want %q", tt.fingerprints, got, tt.want)
+			}
+		})
+	}
+}
+
+func TestWebPocEarlyReturnBranches(t *testing.T) {
+	plugin := NewWebPocPlugin()
+	if plugin == nil || plugin.Name() != "webpoc" {
+		t.Fatalf("unexpected plugin: %#v", plugin)
+	}
+
+	cfg := common.NewConfig()
+	session := common.NewScanSession(cfg, common.NewState(), &common.FlagVars{})
+	info := &common.HostInfo{Host: "example.com", Port: 80}
+
+	cfg.POC.Disabled = true
+	disabled := plugin.Scan(context.Background(), info, session)
+	if disabled.Success || disabled.Error == nil {
+		t.Fatalf("disabled scan = %#v, want failed result with error", disabled)
+	}
+
+	cfg.POC.Disabled = false
+	cfg.POC.Full = false
+	skipped := plugin.Scan(context.Background(), info, session)
+	if !skipped.Success || !skipped.Skipped {
+		t.Fatalf("non-full scan = %#v, want skipped success", skipped)
+	}
+}
@@ -24,6 +24,8 @@ import (
 	"github.com/shadow1ng/fscan/webscan/lib"
 )

+const maxWebTitleBodyBytes = 2 << 20
+
 // 预编译正则表达式
 var (
 	titleRegex      = regexp.MustCompile(`(?i)<title[^>]*>([^<]+)</title>`)
@@ -136,10 +138,7 @@ func (p *WebTitlePlugin) getWebTitle(ctx context.Context, info *common.HostInfo,
 	baseURL := webTitleURL(urlScheme, info.Host, info.Port)

 	// 选择对应的 HTTP 客户端
-	clientNR, clientR := lib.ClientNoRedirect, lib.Client
-	if isGM {
-		clientNR, clientR = lib.ClientNoRedirectGM, lib.ClientGM
-	}
+	clientNR, clientR := webTitleHTTPClients(isGM)

 	// 构建显示用URL（隐藏标准端口）
 	var displayURL string
@@ -164,7 +163,7 @@ func (p *WebTitlePlugin) getWebTitle(ctx context.Context, info *common.HostInfo,
 		return "", 0, 0, "", nil, displayURL, err
 	}

-	body, err := io.ReadAll(resp.Body)
+	body, err := readWebTitleBody(resp.Body)
 	_ = resp.Body.Close()
 	contentLen := len(body)
 	if err != nil {
@@ -196,7 +195,7 @@ func (p *WebTitlePlugin) getWebTitle(ctx context.Context, info *common.HostInfo,
 					reqRedirect.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36")
 					respRedirect, err := clientR.Do(reqRedirect)
 					if err == nil {
-						bodyRedirect, err := io.ReadAll(respRedirect.Body)
+						bodyRedirect, err := readWebTitleBody(respRedirect.Body)
 						_ = respRedirect.Body.Close()
 						if err == nil && len(bodyRedirect) > 0 {
 							// 添加跳转后页面的指纹数据
@@ -299,6 +298,34 @@ func (p *WebTitlePlugin) triggerPocScan(ctx context.Context, info *common.HostIn
 	WebScan.WebScan(ctx, info, config, session)
 }

+func webTitleHTTPClients(isGM bool) (*http.Client, *http.Client) {
+	if isGM {
+		return firstHTTPClient(lib.ClientNoRedirectGM, defaultNoRedirectClient()), firstHTTPClient(lib.ClientGM, http.DefaultClient)
+	}
+	return firstHTTPClient(lib.ClientNoRedirect, defaultNoRedirectClient()), firstHTTPClient(lib.Client, http.DefaultClient)
+}
+
+func firstHTTPClient(clients ...*http.Client) *http.Client {
+	for _, client := range clients {
+		if client != nil {
+			return client
+		}
+	}
+	return http.DefaultClient
+}
+
+func defaultNoRedirectClient() *http.Client {
+	return &http.Client{
+		CheckRedirect: func(*http.Request, []*http.Request) error {
+			return http.ErrUseLastResponse
+		},
+	}
+}
+
+func readWebTitleBody(r io.Reader) ([]byte, error) {
+	return io.ReadAll(io.LimitReader(r, maxWebTitleBodyBytes))
+}
+
 // formatHeaders 将 HTTP Header 格式化为字符串
 func (p *WebTitlePlugin) formatHeaders(headers http.Header) string {
 	var builder strings.Builder
@@ -361,9 +388,7 @@ func (p *WebTitlePlugin) extractTitle(html string) string {
 		title := strings.TrimSpace(matches[1])
 		title = whitespaceRegex.ReplaceAllString(title, " ")

-		if len(title) > 100 {
-			title = title[:100] + "..."
-		}
+		title = truncateRunes(title, 100)

 		if utf8.ValidString(title) {
 			return title
@@ -373,6 +398,19 @@ func (p *WebTitlePlugin) extractTitle(html string) string {
 	return ""
 }

+func truncateRunes(s string, maxRunes int) string {
+	if maxRunes < 0 {
+		return s
+	}
+	for i := range s {
+		if maxRunes == 0 {
+			return s[:i] + "..."
+		}
+		maxRunes--
+	}
+	return s
+}
+
 // fetchFaviconHash 下载 favicon.ico 并计算 hash
 func (p *WebTitlePlugin) fetchFaviconHash(ctx context.Context, baseURL string) fingerprint.FaviconHashes {
 	// 构造 favicon URL
@@ -3,7 +3,9 @@ package web
 import (
 	"context"
 	"net/http"
+	"strings"
 	"testing"
+	"unicode/utf8"

 	"github.com/shadow1ng/fscan/webscan/lib"
 )
@@ -12,6 +14,14 @@ type faviconRoundTripper struct {
 	called bool
 }

+func TestExtractTitleTruncatesByRune(t *testing.T) {
+	title := strings.Repeat("界", 105)
+	got := NewWebTitlePlugin().extractTitle("<html><title>" + title + "</title></html>")
+	if !utf8.ValidString(got) || len([]rune(got)) != 103 || !strings.HasSuffix(got, "...") {
+		t.Fatalf("extractTitle() = %q", got)
+	}
+}
+
 func (rt *faviconRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
 	rt.called = true
 	<-req.Context().Done()
@@ -56,3 +66,38 @@ func TestWebTitleURLUsesJoinHostPort(t *testing.T) {
 		})
 	}
 }
+
+func TestWebTitleHTTPClientsFallbackWhenGlobalsNil(t *testing.T) {
+	previousClient, previousNoRedirect := lib.Client, lib.ClientNoRedirect
+	previousGM, previousNoRedirectGM := lib.ClientGM, lib.ClientNoRedirectGM
+	lib.Client, lib.ClientNoRedirect = nil, nil
+	lib.ClientGM, lib.ClientNoRedirectGM = nil, nil
+	defer func() {
+		lib.Client, lib.ClientNoRedirect = previousClient, previousNoRedirect
+		lib.ClientGM, lib.ClientNoRedirectGM = previousGM, previousNoRedirectGM
+	}()
+
+	clientNR, clientR := webTitleHTTPClients(false)
+	if clientNR == nil || clientR == nil {
+		t.Fatal("webTitleHTTPClients returned nil fallback client")
+	}
+
+	req, err := http.NewRequest(http.MethodGet, "http://example.com", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := clientNR.CheckRedirect(req, []*http.Request{req}); err != http.ErrUseLastResponse {
+		t.Fatalf("no-redirect fallback error = %v, want http.ErrUseLastResponse", err)
+	}
+}
+
+func TestReadWebTitleBodyIsBounded(t *testing.T) {
+	body := strings.NewReader(strings.Repeat("a", maxWebTitleBodyBytes+1024))
+	got, err := readWebTitleBody(body)
+	if err != nil {
+		t.Fatalf("readWebTitleBody error = %v", err)
+	}
+	if len(got) != maxWebTitleBodyBytes {
+		t.Fatalf("body len = %d, want %d", len(got), maxWebTitleBodyBytes)
+	}
+}