First Commit

2026-10-09 08:38:16 +08:00
commit 0335d572de
756 changed files with 201663 additions and 0 deletions
@@ -0,0 +1,405 @@
+package db
+
+import (
+	"fmt"
+	"testing"
+	"time"
+)
+
+// TestParseScopeLine covers classification + guardrails without a DB.
+func TestParseScopeLine(t *testing.T) {
+	cases := []struct {
+		in      string
+		kind    string
+		wantErr bool
+	}{
+		{"example.com", "domain", false},
+		{"https://sub.example.com/path", "domain", false},
+		{"1.2.3.4", "ip", false},
+		{"10.0.0.0/8", "", true}, // over-broad IPv4 (< /16)
+		{"198.51.100.0/24", "cidr", false},
+		{"co.uk", "", true}, // bare public suffix
+		{"not a host", "", true},
+		{"1.2.3.1-1.2.3.9", "", true}, // ranges must be CIDR
+	}
+	for _, c := range cases {
+		r, err := ParseScopeLine(c.in)
+		if c.wantErr {
+			if err == nil {
+				t.Errorf("ParseScopeLine(%q) want error, got %+v", c.in, r)
+			}
+			continue
+		}
+		if err != nil {
+			t.Errorf("ParseScopeLine(%q) unexpected error: %v", c.in, err)
+			continue
+		}
+		if r.Kind != c.kind {
+			t.Errorf("ParseScopeLine(%q) kind=%q want %q", c.in, r.Kind, c.kind)
+		}
+	}
+}
+
+func TestParseAutoScopeLine(t *testing.T) {
+	cases := []struct {
+		name       string
+		input      string
+		kind       string
+		normalized string
+		wantErr    bool
+	}{
+		{name: "domain", input: "example.com", kind: "domain", normalized: "example.com"},
+		{name: "url", input: "https://sub.example.com/path", kind: "domain", normalized: "sub.example.com"},
+		{name: "url query", input: "https://example.com/path?source=x", kind: "domain", normalized: "example.com"},
+		{name: "url credentials", input: "https://user:pass@example.com/path", kind: "domain", normalized: "example.com"},
+		{name: "url ip", input: "http://203.0.113.10/path", kind: "ip", normalized: "203.0.113.10/32"},
+		{name: "ipv4", input: "203.0.113.10", kind: "ip", normalized: "203.0.113.10/32"},
+		{name: "ipv6", input: "2001:db8::10", kind: "ip", normalized: "2001:db8::10/128"},
+		{name: "cidr", input: "198.51.100.0/24", kind: "cidr", normalized: "198.51.100.0/24"},
+		{name: "icp latin", input: "京 ICP备 123号", kind: "icp", normalized: "京icp备123号"},
+		{name: "icp chinese", input: "沪网备案 9988", kind: "icp", normalized: "沪网备案9988"},
+		{name: "icp domain", input: "icp.example.com", kind: "domain", normalized: "icp.example.com"},
+		{name: "icp url query", input: "https://example.com/path?icp=1", kind: "domain", normalized: "example.com"},
+		// 备案号不含点号：掺了域名/版本号的描述性文字归关键词，否则会存成一条
+		// 永远匹配不上的死 ICP 规则。
+		{name: "icp with domain text", input: "备案 www.example.com", kind: "keyword", normalized: "备案 www.example.com"},
+		{name: "icp with version text", input: "某公司 ICP v1.0", kind: "keyword", normalized: "某公司 icp v1.0"},
+		{name: "icp fullwidth dot", input: "备案 例．com", kind: "keyword", normalized: "备案 例．com"},
+		{name: "keyword", input: "  ACME   Security  ", kind: "keyword", normalized: "acme security"},
+		{name: "colon keyword", input: "ACME: Cloud: Security", kind: "keyword", normalized: "acme: cloud: security"},
+		{name: "empty", input: "  ", wantErr: true},
+		{name: "invalid cidr", input: "10.0.0.0/not-a-prefix", wantErr: true},
+		{name: "invalid ipv4", input: "999.0.0.1", wantErr: true},
+		{name: "invalid ipv6", input: "2001:db8::zz", wantErr: true},
+		{name: "invalid ipv6 cidr", input: "2001:db8::zz/64", wantErr: true},
+		{name: "overbroad cidr", input: "10.0.0.0/8", wantErr: true},
+		{name: "bare suffix", input: "co.uk", wantErr: true},
+		{name: "empty domain label", input: "foo..example.com", wantErr: true},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			rule, err := ParseAutoScopeLine(tc.input)
+			if tc.wantErr {
+				if err == nil {
+					t.Fatalf("ParseAutoScopeLine(%q) = %+v, want error", tc.input, rule)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("ParseAutoScopeLine(%q): %v", tc.input, err)
+			}
+			if rule.Kind != tc.kind {
+				t.Fatalf("kind=%q want %q", rule.Kind, tc.kind)
+			}
+			got := rule.Value
+			if rule.Kind == "domain" {
+				got = rule.Domain
+			} else if rule.Kind == "ip" || rule.Kind == "cidr" {
+				got = rule.Net
+			}
+			if got != tc.normalized {
+				t.Fatalf("normalized=%q want %q", got, tc.normalized)
+			}
+		})
+	}
+}
+
+func TestExplicitCompanyAttributionSurvivesScopeRebuild(t *testing.T) {
+	d, as, cs := testSetup(t)
+	defer d.Close()
+
+	stamp := time.Now().UnixNano()
+	explicitCompany, _, err := cs.UpsertCompany(fmt.Sprintf("Explicit Attribution %d", stamp), "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	autoCompany, _, err := cs.UpsertCompany(fmt.Sprintf("Automatic Attribution %d", stamp), "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cleanupCompany(d, explicitCompany)
+	defer cleanupCompany(d, autoCompany)
+
+	domain := fmt.Sprintf("explicit-%d.invalid", stamp)
+	icp := fmt.Sprintf("ICP-%d", stamp)
+	network := fmt.Sprintf("2001:db8:%x::/64", uint64(stamp)&0xffff)
+	ip := fmt.Sprintf("2001:db8:%x::10", uint64(stamp)&0xffff)
+
+	// A pre-existing row with company_id and no provenance value represents old
+	// installations. The schema default conservatively treats it as explicit.
+	var assetID int64
+	if err := d.QueryRow(`INSERT INTO assets(type,domain,root_domain,company_id)
+		VALUES ('root_domain',$1,$1,$2) RETURNING id`, domain, explicitCompany).Scan(&assetID); err != nil {
+		t.Fatal(err)
+	}
+	defer deleteAsset(d, assetID)
+	appID, err := as.UpsertApp(UpsertAppReq{
+		Name: fmt.Sprintf("explicit-app-%d", stamp), ICP: icp, CompanyID: &explicitCompany,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer deleteAsset(d, appID)
+	autoAssetID, err := as.UpsertIP(UpsertIPReq{IP: ip})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer deleteAsset(d, autoAssetID)
+
+	rules := []ScopeInput{
+		{Kind: "domain", Value: domain},
+		{Kind: "icp", Value: icp},
+		{Kind: "cidr", Value: network},
+	}
+	if added, _, invalid, errs := cs.AddScopeInputs(autoCompany, rules, "test"); added != len(rules) || invalid != 0 {
+		t.Fatalf("AddScopeInputs: added=%d invalid=%d errors=%v", added, invalid, errs)
+	}
+
+	assertCompany := func(id, want int64, source string) {
+		t.Helper()
+		var got *int64
+		var gotSource string
+		if err := d.QueryRow(`SELECT company_id,company_source FROM assets WHERE id=$1`, id).Scan(&got, &gotSource); err != nil {
+			t.Fatal(err)
+		}
+		if got == nil || *got != want || gotSource != source {
+			t.Fatalf("asset %d company=%v source=%q, want %d/%q", id, got, gotSource, want, source)
+		}
+	}
+	assertCompany(assetID, explicitCompany, "explicit")
+	assertCompany(appID, explicitCompany, "explicit")
+	assertCompany(autoAssetID, autoCompany, "scope")
+
+	// Replacing every matching rule detaches only the automatically-owned row.
+	if _, invalid, errs := cs.UpdateScopeInputs(autoCompany, []ScopeInput{
+		{Kind: "domain", Value: fmt.Sprintf("replacement-%d.invalid", stamp)},
+	}, "test"); invalid != 0 || len(errs) != 0 {
+		t.Fatalf("UpdateScopeInputs: invalid=%d errors=%v", invalid, errs)
+	}
+	assertCompany(assetID, explicitCompany, "explicit")
+	assertCompany(appID, explicitCompany, "explicit")
+	var autoCompanyID *int64
+	var autoSource string
+	if err := d.QueryRow(`SELECT company_id,company_source FROM assets WHERE id=$1`, autoAssetID).Scan(&autoCompanyID, &autoSource); err != nil {
+		t.Fatal(err)
+	}
+	if autoCompanyID != nil || autoSource != "scope" {
+		t.Fatalf("automatic asset was not detached: company=%v source=%q", autoCompanyID, autoSource)
+	}
+
+	if err := cs.RecomputeAttribution(); err != nil {
+		t.Fatal(err)
+	}
+	assertCompany(assetID, explicitCompany, "explicit")
+	assertCompany(appID, explicitCompany, "explicit")
+
+	// Deleting the explicitly selected company detaches through the FK, then the
+	// transactional rebuild may adopt the assets into a still-valid scope.
+	if _, invalid, errs := cs.UpdateScopeInputs(autoCompany, []ScopeInput{
+		{Kind: "domain", Value: domain},
+		{Kind: "icp", Value: icp},
+	}, "test"); invalid != 0 || len(errs) != 0 {
+		t.Fatalf("restore fallback scope: invalid=%d errors=%v", invalid, errs)
+	}
+	if err := cs.DeleteCompany(explicitCompany); err != nil {
+		t.Fatal(err)
+	}
+	assertCompany(assetID, autoCompany, "scope")
+	assertCompany(appID, autoCompany, "scope")
+}
+
+func TestParseStructuredCompanyScope(t *testing.T) {
+	icp, err := ParseScopeInput(ScopeInput{Kind: "ICP", Value: " 京ICP 备 123号-1\t"})
+	if err != nil {
+		t.Fatalf("parse ICP: %v", err)
+	}
+	if icp.Kind != "icp" || icp.Value != "京icp备123号-1" {
+		t.Fatalf("unexpected normalized ICP: %+v", icp)
+	}
+	keyword, err := ParseScopeInput(ScopeInput{Kind: "keyword", Value: "  ACME   Security  "})
+	if err != nil {
+		t.Fatalf("parse keyword: %v", err)
+	}
+	if keyword.Value != "acme security" {
+		t.Fatalf("unexpected normalized keyword: %+v", keyword)
+	}
+	if _, err := ParseScopeInput(ScopeInput{Kind: "ip", Value: "example.com"}); err == nil {
+		t.Fatal("typed IP accepted a domain")
+	}
+}
+
+func TestCompanyICPAttribution(t *testing.T) {
+	d, err := Open(testDSN(t))
+	if err != nil {
+		t.Skipf("postgres unavailable (%v) — skipping", err)
+	}
+	// 关连接必须走 t.Cleanup 且**注册在清理之前**：t.Cleanup 是后进先出，
+	// 先注册关闭 → 关闭最后执行，下面的数据清理才连得上库。
+	// 原先这里是 `defer d.Close()`：defer 在函数返回时先跑，t.Cleanup 在那之后
+	// 才执行，于是清理语句全落在**已关闭的连接**上、错误又被 `_, _ =` 丢弃，
+	// 资产与公司就永久残留在库里。残留本身不会立刻报错，但本用例用
+	// `MAX(companies.id)+1` 当假 TaskID 给资产打标（见下方 suffix），
+	// 一旦这个数字与别的用例的任务 id 撞上，那个用例按「恰好 N 个资产」的断言
+	// 就会莫名失败——排查成本极高。
+	t.Cleanup(func() { d.Close() })
+
+	var suffix int64
+	if err := d.QueryRow(`SELECT COALESCE(MAX(id),0)+1 FROM companies`).Scan(&suffix); err != nil {
+		t.Fatal(err)
+	}
+	cs := d.Companies()
+	as := d.Assets()
+	companyID, _, err := cs.UpsertCompany(fmt.Sprintf("ICP Scope Co %d", suffix), "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		// 不吞错误：清理失败会污染后续用例，必须让它在本次运行里显形。
+		if _, err := d.Exec(`DELETE FROM assets WHERE task_ids @> ARRAY[$1]::bigint[]`, suffix); err != nil {
+			t.Errorf("清理测试资产失败: %v", err)
+		}
+		if _, err := d.Exec(`DELETE FROM companies WHERE id=$1`, companyID); err != nil {
+			t.Errorf("清理测试公司失败: %v", err)
+		}
+	})
+
+	// A keyword can guide an Agent, but must never claim an asset by its name.
+	added, _, invalid, errs := cs.AddScopeInputs(companyID, []ScopeInput{
+		{Kind: "icp", Value: "京 ICP备 998877号"},
+		{Kind: "keyword", Value: "ICP Scope"},
+	}, "unit test")
+	if added != 2 || invalid != 0 || len(errs) != 0 {
+		t.Fatalf("add structured scope: added=%d invalid=%d errors=%v", added, invalid, errs)
+	}
+
+	rootID, err := as.UpsertRootDomain(UpsertRootDomainReq{
+		Domain: fmt.Sprintf("icp-scope-%d.example", suffix), ICP: "京icp备998877号", TaskID: suffix,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	appID, err := as.UpsertApp(UpsertAppReq{
+		Name: fmt.Sprintf("ICP Scope Keyword Only %d", suffix), TaskID: suffix,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	icpAppID, err := as.UpsertApp(UpsertAppReq{
+		Name: fmt.Sprintf("ICP Matched App %d", suffix), ICP: " 京 ICP备 998877号 ", TaskID: suffix,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assertCompany := func(assetID int64, want *int64) {
+		t.Helper()
+		var got *int64
+		if err := d.QueryRow(`SELECT company_id FROM assets WHERE id=$1`, assetID).Scan(&got); err != nil {
+			t.Fatal(err)
+		}
+		if want == nil && got != nil {
+			t.Fatalf("asset %d attributed by keyword: %d", assetID, *got)
+		}
+		if want != nil && (got == nil || *got != *want) {
+			t.Fatalf("asset %d company=%v want %d", assetID, got, *want)
+		}
+	}
+	assertCompany(rootID, &companyID)
+	assertCompany(appID, nil)
+	assertCompany(icpAppID, &companyID)
+}
+
+// TestCompanyScopeAttribution exercises the full loop against dev PG: create
+// company (unique name), add scope, and verify auto-attribution at insert time,
+// backfill of a pre-existing asset, CIDR + domain-suffix matching, and that
+// out-of-scope assets stay unattributed.
+func TestCompanyScopeAttribution(t *testing.T) {
+	d, err := Open(testDSN(t))
+	if err != nil {
+		t.Skipf("postgres unavailable (%v) — skipping", err)
+	}
+	as := d.Assets()
+	cs := d.Companies()
+
+	var startMax int64
+	if err := d.QueryRow(`SELECT COALESCE(MAX(id),0) FROM assets`).Scan(&startMax); err != nil {
+		d.Close()
+		t.Fatalf("startMax: %v", err)
+	}
+	t.Cleanup(func() {
+		_, _ = d.Exec(`DELETE FROM assets WHERE id > $1`, startMax)
+		d.Close()
+	})
+
+	uniq := startMax + 1
+	root := fmt.Sprintf("scopetest%d.com", uniq)
+	sub := "api." + root
+	ipIn := "198.51.100.9"
+	ipOut := "203.0.113.9"
+	outDomain := fmt.Sprintf("other%d.net", uniq)
+
+	cid, _, err := cs.UpsertCompany(fmt.Sprintf("ScopeCo %d", uniq), "")
+	if err != nil {
+		t.Fatalf("UpsertCompany: %v", err)
+	}
+
+	// a pre-existing asset (inserted BEFORE any scope) — must be back-filled.
+	preID, err := as.UpsertSubdomain(UpsertSubdomainReq{Domain: sub})
+	if err != nil {
+		t.Fatalf("pre upsert: %v", err)
+	}
+	var preCompanyID *int64
+	d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, preID).Scan(&preCompanyID)
+	if preCompanyID != nil {
+		t.Fatalf("pre-scope asset should be unattributed, got %v", *preCompanyID)
+	}
+
+	cs.AddScope(cid, []string{root, "198.51.100.0/24"}, "unit test")
+
+	mustCid := func(id int64, want int64, label string) {
+		var cID *int64
+		d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, id).Scan(&cID)
+		if cID == nil {
+			t.Fatalf("%s company_id = nil, want %d", label, want)
+		}
+		if *cID != want {
+			t.Fatalf("%s company_id = %d, want %d", label, *cID, want)
+		}
+	}
+	mustNil := func(id int64, label string) {
+		var cID *int64
+		d.QueryRow(`SELECT company_id FROM assets WHERE id = $1`, id).Scan(&cID)
+		if cID != nil {
+			t.Fatalf("%s should be unattributed, got %d", label, *cID)
+		}
+	}
+
+	// backfill attributed the pre-existing subdomain (domain suffix match).
+	mustCid(preID, cid, "pre-existing subdomain (backfill)")
+
+	// insert-time attribution: ip in CIDR, another subdomain.
+	ipInID, err := as.UpsertIP(UpsertIPReq{IP: ipIn})
+	if err != nil {
+		t.Fatalf("UpsertIP in: %v", err)
+	}
+	mustCid(ipInID, cid, "in-CIDR ip (insert-time)")
+
+	sub2ID, err := as.UpsertSubdomain(UpsertSubdomainReq{Domain: "www." + root})
+	if err != nil {
+		t.Fatalf("UpsertSubdomain: %v", err)
+	}
+	mustCid(sub2ID, cid, "new subdomain (insert-time)")
+
+	// out of scope stays unattributed.
+	outID, err := as.UpsertRootDomain(UpsertRootDomainReq{Domain: outDomain})
+	if err != nil {
+		t.Fatalf("UpsertRootDomain out: %v", err)
+	}
+	mustNil(outID, "out-of-scope domain")
+
+	ipOutID, err := as.UpsertIP(UpsertIPReq{IP: ipOut})
+	if err != nil {
+		t.Fatalf("UpsertIP out: %v", err)
+	}
+	mustNil(ipOutID, "out-of-CIDR ip")
+}