admin/ARTEX
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(assets): add editable company and task scopes

Browse Source
This commit is contained in:
pphua committed 2026-08-24 19:29:14 +08:00
1 parent 4392b4fee1
commit 31dc9e7446
20 files changed
+4186 -783
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

db/asset_dsl.go
+16 -1
View File
@@ -457,6 +457,18 @@ func isIdentChar(c byte) bool {
// ── QueryDSL ─────────────────────────────────────────────────────────────────
// ValidateDSL parses and compiles a DSL expression without touching the
// database. HTTP callers use it to distinguish client syntax errors from query
// failures, which must remain server errors.
func ValidateDSL(dsl string) error {
node, err := ParseDSL(dsl)
if err != nil {
return err
}
_, _, err = buildDSLWhere(node)
return err
}
// CountDSL returns the total number of assets matching a DSL expression (and optional
// type), for server-side pagination — same WHERE as QueryDSL, without LIMIT/OFFSET.
func (s *AssetStore) CountDSL(dsl, typ string) (int, error) {
@@ -483,6 +495,9 @@ func (s *AssetStore) QueryDSL(dsl, typ string, limit, offset int) ([]*Asset, err
if limit <= 0 {
limit = 50
}
if offset < 0 {
offset = 0
}
node, err := ParseDSL(dsl)
if err != nil {
return nil, err
@@ -497,7 +512,7 @@ func (s *AssetStore) QueryDSL(dsl, typ string, limit, offset int) ([]*Asset, err
}
args = append(args, limit, offset)
q := assetSelectCols + " WHERE " + where +
fmt.Sprintf(" ORDER BY last_seen DESC LIMIT $%d OFFSET $%d", len(args)-1, len(args))
fmt.Sprintf(" ORDER BY last_seen DESC, id DESC LIMIT $%d OFFSET $%d", len(args)-1, len(args))
rows, err := s.db.Query(q, args...)
if err != nil {
return nil, err
db/assets.go
+408 -59
View File
@@ -6,6 +6,7 @@ import (
"fmt"
"net"
"net/url"
"sort"
"strings"
)
@@ -51,8 +52,11 @@ type Asset struct {
Method string `json:"method,omitempty"`
Params []map[string]any `json:"params,omitempty"`
// meta
Extra map[string]any `json:"extra,omitempty"`
LastSeen string `json:"last_seen"`
Extra map[string]any `json:"extra,omitempty"`
LastSeen string `json:"last_seen"`
TaskSource string `json:"task_source,omitempty"`
TaskSourceSummary string `json:"task_source_summary,omitempty"`
TaskSourceNodeID *int64 `json:"task_source_node_id,omitempty"`
}
// AuthItem is one entry in the auth array.
@@ -82,6 +86,7 @@ type PortService struct {
type AssetStore struct {
db *DB
company *CompanyStore
tx *sql.Tx
}
// Assets returns the asset store.
@@ -92,6 +97,39 @@ func (d *DB) Assets() *AssetStore {
// Companies returns the company store associated with this asset store.
func (s *AssetStore) Companies() *CompanyStore { return s.company }
// withCompanyScopeMutation serializes scope resolution and every asset write
// that consumes its result in one transaction. Nested asset side effects reuse
// the same transaction through the scoped store.
func (s *AssetStore) withCompanyScopeMutation(fn func(*AssetStore) (int64, error)) (int64, error) {
if s.tx != nil {
return fn(s)
}
tx, err := s.db.Begin()
if err != nil {
return 0, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, err
}
scoped := &AssetStore{db: s.db, company: s.company, tx: tx}
id, err := fn(scoped)
if err != nil {
return 0, err
}
if err := tx.Commit(); err != nil {
return 0, err
}
return id, nil
}
func (s *AssetStore) resolveCompanyWithICP(rootDomain, ipStr, icp string) (*int64, error) {
if s.tx != nil {
return resolveCompanyWithICP(s.tx, rootDomain, ipStr, icp)
}
return s.company.ResolveCompanyWithICP(rootDomain, ipStr, icp)
}
// =====================================================================
// Helpers
// =====================================================================
@@ -218,7 +256,12 @@ func (s *AssetStore) UpsertRootDomain(req UpsertRootDomainReq) (int64, error) {
if domain == "" {
return 0, fmt.Errorf("domain is required")
}
companyID, err := s.company.ResolveCompany(domain, "")
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertRootDomain(req)
})
}
companyID, err := s.resolveCompanyWithICP(domain, "", req.ICP)
if err != nil {
return 0, err
}
@@ -236,12 +279,20 @@ func (s *AssetStore) UpsertRootDomain(req UpsertRootDomainReq) (int64, error) {
}
var id int64
err = s.db.QueryRow(`
INSERT INTO assets(type, domain, root_domain, icp, company_id, task_ids)
VALUES ('root_domain', $1, $1, $2, $3, $4::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, domain, root_domain, icp, company_id, company_source, task_ids)
VALUES ('root_domain', $1, $1, $2, $3, 'scope', $4::bigint[])
ON CONFLICT (domain) WHERE type = 'root_domain' DO UPDATE SET
icp = COALESCE(EXCLUDED.icp, assets.icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
extra = assets.extra || EXCLUDED.extra,
last_seen = now()
@@ -266,8 +317,13 @@ func (s *AssetStore) UpsertIP(req UpsertIPReq) (int64, error) {
if req.IP == "" {
return 0, fmt.Errorf("ip is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertIP(req)
})
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany("", req.IP)
companyID, err := s.resolveCompanyWithICP("", req.IP, "")
if err != nil {
return 0, err
}
@@ -291,9 +347,9 @@ func (s *AssetStore) UpsertIP(req UpsertIPReq) (int64, error) {
}
var id int64
err = s.db.QueryRow(`
INSERT INTO assets(type, ip, c_segment, bound_domains, open_ports, company_id, task_ids)
VALUES ('ip', $1, $2::cidr, $3::text[], $4::jsonb[], $5, $6::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, ip, c_segment, bound_domains, open_ports, company_id, company_source, task_ids)
VALUES ('ip', $1, $2::cidr, $3::text[], $4::jsonb[], $5, 'scope', $6::bigint[])
ON CONFLICT (ip) WHERE type = 'ip' DO UPDATE SET
bound_domains = (SELECT ARRAY(SELECT DISTINCT unnest(assets.bound_domains || EXCLUDED.bound_domains))),
open_ports = (
@@ -305,7 +361,15 @@ ON CONFLICT (ip) WHERE type = 'ip' DO UPDATE SET
)
),
c_segment = COALESCE(assets.c_segment, EXCLUDED.c_segment),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
last_seen = now()
RETURNING id`, req.IP, csegVal, boundDomains, openPortsJSON, companyID, taskIDs).Scan(&id)
@@ -366,14 +430,19 @@ func (s *AssetStore) UpsertSubdomain(req UpsertSubdomainReq) (id int64, err erro
if domain == "" {
return 0, fmt.Errorf("domain is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertSubdomain(req)
})
}
rootDomain, _ := RootDomain(domain)
if rootDomain == "" {
rootDomain = domain
}
// resolve company by root domain
companyID, err := s.company.ResolveCompany(rootDomain, "")
// Resolve by root domain first, then the asset's exact normalized ICP.
companyID, err := s.resolveCompanyWithICP(rootDomain, "", req.ICP)
if err != nil {
return 0, err
}
@@ -408,7 +477,7 @@ func (s *AssetStore) UpsertSubdomain(req UpsertSubdomainReq) (id int64, err erro
cseg := calcCSegment(ipStr)
if companyID == nil && ipStr != "" {
companyID, _ = s.company.ResolveCompany("", ipStr)
companyID, _ = s.resolveCompanyWithICP("", ipStr, "")
}
var taskIDs string
@@ -431,14 +500,22 @@ func (s *AssetStore) UpsertSubdomain(req UpsertSubdomainReq) (id int64, err erro
recordType := req.RecordType
recordValueArr := marshalStringArray(req.RecordValue)
err = s.db.QueryRow(`
INSERT INTO assets(type, domain, root_domain, record_type, record_value, ip, c_segment, icp, company_id, task_ids)
VALUES ('subdomain', $1, $2, $3, $4::text[], $5, $6::cidr, $7, $8, $9::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, domain, root_domain, record_type, record_value, ip, c_segment, icp, company_id, company_source, task_ids)
VALUES ('subdomain', $1, $2, $3, $4::text[], $5, $6::cidr, $7, $8, 'scope', $9::bigint[])
ON CONFLICT (domain, COALESCE(record_type,'')) WHERE type = 'subdomain' DO UPDATE SET
ip = COALESCE(EXCLUDED.ip, assets.ip),
c_segment = COALESCE(EXCLUDED.c_segment, assets.c_segment),
icp = COALESCE(EXCLUDED.icp, assets.icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
record_value = (SELECT ARRAY(SELECT DISTINCT unnest(assets.record_value || EXCLUDED.record_value))),
extra = assets.extra || EXCLUDED.extra,
@@ -458,7 +535,7 @@ type UpsertAppReq struct {
Category string
Description string
ICP string
CompanyID *int64 // explicit override; nil = no auto attribution for apps
CompanyID *int64 // explicit override; nil = exact ICP auto-attribution when available
TaskID int64
}
@@ -467,6 +544,11 @@ func (s *AssetStore) UpsertApp(req UpsertAppReq) (int64, error) {
if req.Name == "" {
return 0, fmt.Errorf("app name is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertApp(req)
})
}
var taskIDs string
if req.TaskID > 0 {
@@ -476,6 +558,7 @@ func (s *AssetStore) UpsertApp(req UpsertAppReq) (int64, error) {
}
var bundleVal, catVal, descVal, icpVal, companyIDVal any
companySource := "scope"
if req.BundleID != "" {
bundleVal = req.BundleID
}
@@ -490,36 +573,65 @@ func (s *AssetStore) UpsertApp(req UpsertAppReq) (int64, error) {
}
if req.CompanyID != nil {
companyIDVal = *req.CompanyID
companySource = "explicit"
} else {
companyID, err := s.resolveCompanyWithICP("", "", req.ICP)
if err != nil {
return 0, err
}
if companyID != nil {
companyIDVal = *companyID
}
}
var id int64
var err error
if req.BundleID != "" {
err = s.db.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, task_ids)
VALUES ('app', $1, $2, $3, $4, $5, $6, $7::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, company_source, task_ids)
VALUES ('app', $1, $2, $3, $4, $5, $6, $7, $8::bigint[])
ON CONFLICT (bundle_id) WHERE type = 'app' AND bundle_id IS NOT NULL DO UPDATE SET
app_name = COALESCE(EXCLUDED.app_name, assets.app_name),
category = COALESCE(EXCLUDED.category, assets.category),
app_description = COALESCE(EXCLUDED.app_description, assets.app_description),
app_icp = COALESCE(EXCLUDED.app_icp, assets.app_icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN EXCLUDED.company_id
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN 'explicit'
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
last_seen = now()
RETURNING id`, bundleVal, req.Name, catVal, descVal, icpVal, companyIDVal, taskIDs).Scan(&id)
RETURNING id`, bundleVal, req.Name, catVal, descVal, icpVal, companyIDVal, companySource, taskIDs).Scan(&id)
} else {
err = s.db.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, task_ids)
VALUES ('app', NULL, $1, $2, $3, $4, $5, $6::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, company_source, task_ids)
VALUES ('app', NULL, $1, $2, $3, $4, $5, $6, $7::bigint[])
ON CONFLICT (app_name) WHERE type = 'app' AND bundle_id IS NULL DO UPDATE SET
category = COALESCE(EXCLUDED.category, assets.category),
app_description = COALESCE(EXCLUDED.app_description, assets.app_description),
app_icp = COALESCE(EXCLUDED.app_icp, assets.app_icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN EXCLUDED.company_id
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN 'explicit'
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
last_seen = now()
RETURNING id`, req.Name, catVal, descVal, icpVal, companyIDVal, taskIDs).Scan(&id)
RETURNING id`, req.Name, catVal, descVal, icpVal, companyIDVal, companySource, taskIDs).Scan(&id)
}
return id, err
}
@@ -547,6 +659,11 @@ func (s *AssetStore) UpsertHTTPService(req UpsertHTTPServiceReq) (int64, error)
if req.URL == "" {
return 0, fmt.Errorf("url is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertHTTPService(req)
})
}
normURL := normalizeURL(req.URL)
domain, port, serviceName := parseURL(normURL)
rootDomain, _ := RootDomain(domain)
@@ -555,7 +672,7 @@ func (s *AssetStore) UpsertHTTPService(req UpsertHTTPServiceReq) (int64, error)
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany(rootDomain, req.IP)
companyID, err := s.resolveCompanyWithICP(rootDomain, req.IP, "")
if err != nil {
return 0, err
}
@@ -591,16 +708,16 @@ func (s *AssetStore) UpsertHTTPService(req UpsertHTTPServiceReq) (int64, error)
}
var id int64
err = s.db.QueryRow(`
err = s.tx.QueryRow(`
INSERT INTO assets(
type, url, service_type, service_name, domain, ip, port, root_domain,
c_segment, favicon_mmh3, technologies, status_code, content_length, page_title,
auth, company_id, task_ids
auth, company_id, company_source, task_ids
)
VALUES (
'service', $1, 'http', $2, $3, $4, $5, $6,
$7::cidr, $8, $9::text[], $10, $11, $12,
$13::jsonb[], $14, $15::bigint[]
$13::jsonb[], $14, 'scope', $15::bigint[]
)
ON CONFLICT (url) WHERE type = 'service' AND service_type = 'http' DO UPDATE SET
status_code = COALESCE(EXCLUDED.status_code, assets.status_code),
@@ -617,7 +734,15 @@ ON CONFLICT (url) WHERE type = 'service' AND service_type = 'http' DO UPDATE SET
)
),
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
last_seen = now()
RETURNING id`,
normURL, serviceName, domainVal, ipVal, nullableInt(port), rootDomain,
@@ -674,6 +799,11 @@ func (s *AssetStore) UpsertOtherService(req UpsertOtherServiceReq) (int64, error
if req.ServiceName == "" {
return 0, fmt.Errorf("service_name is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertOtherService(req)
})
}
// normalize service_name (lowercase) so the (domain,ip,port,service_name)
// dedup key doesn't split "SSH" and "ssh" into separate rows.
serviceName := strings.ToLower(strings.TrimSpace(req.ServiceName))
@@ -688,7 +818,7 @@ func (s *AssetStore) UpsertOtherService(req UpsertOtherServiceReq) (int64, error
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany(rootDomain, req.IP)
companyID, err := s.resolveCompanyWithICP(rootDomain, req.IP, "")
if err != nil {
return 0, err
}
@@ -720,12 +850,12 @@ func (s *AssetStore) UpsertOtherService(req UpsertOtherServiceReq) (int64, error
}
var id int64
err = s.db.QueryRow(`
err = s.tx.QueryRow(`
INSERT INTO assets(
type, service_type, service_name, domain, ip, port, root_domain,
c_segment, auth, company_id, task_ids
c_segment, auth, company_id, company_source, task_ids
)
VALUES ('service', 'other', $1, $2, $3, $4, $5, $6::cidr, $7::jsonb[], $8, $9::bigint[])
VALUES ('service', 'other', $1, $2, $3, $4, $5, $6::cidr, $7::jsonb[], $8, 'scope', $9::bigint[])
ON CONFLICT (COALESCE(domain,''), COALESCE(ip,''), port, service_name) WHERE type = 'service' AND service_type = 'other' DO UPDATE SET
domain = COALESCE(EXCLUDED.domain, assets.domain),
ip = COALESCE(EXCLUDED.ip, assets.ip),
@@ -737,7 +867,15 @@ ON CONFLICT (COALESCE(domain,''), COALESCE(ip,''), port, service_name) WHERE typ
)
),
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
last_seen = now()
RETURNING id`,
serviceName, domainVal, ipVal, req.Port, rootDomainVal,
@@ -798,6 +936,11 @@ func (s *AssetStore) UpsertEndpoint(req UpsertEndpointReq) (int64, error) {
if req.Method == "" {
return 0, fmt.Errorf("method is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertEndpoint(req)
})
}
method := strings.ToUpper(req.Method)
normURL := normalizeURL(req.URL)
domain, port, _ := parseURL(normURL)
@@ -807,7 +950,7 @@ func (s *AssetStore) UpsertEndpoint(req UpsertEndpointReq) (int64, error) {
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany(rootDomain, req.IP)
companyID, err := s.resolveCompanyWithICP(rootDomain, req.IP, "")
if err != nil {
return 0, err
}
@@ -839,12 +982,12 @@ func (s *AssetStore) UpsertEndpoint(req UpsertEndpointReq) (int64, error) {
}
var id int64
err = s.db.QueryRow(`
err = s.tx.QueryRow(`
INSERT INTO assets(
type, url, method, domain, ip, port, root_domain,
params, company_id, task_ids, c_segment
params, company_id, company_source, task_ids, c_segment
)
VALUES ('endpoint', $1, $2, $3, $4, $5, $6, $7::jsonb[], $8, $9::bigint[], $10::cidr)
VALUES ('endpoint', $1, $2, $3, $4, $5, $6, $7::jsonb[], $8, 'scope', $9::bigint[], $10::cidr)
ON CONFLICT (url, method) WHERE type = 'endpoint' DO UPDATE SET
params = (
SELECT ARRAY(
@@ -856,7 +999,15 @@ ON CONFLICT (url, method) WHERE type = 'endpoint' DO UPDATE SET
ip = COALESCE(EXCLUDED.ip, assets.ip),
c_segment = COALESCE(EXCLUDED.c_segment, assets.c_segment),
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
last_seen = now()
RETURNING id`,
normURL, method, domainVal, ipVal, nullableInt(port), rootDomainVal,
@@ -883,6 +1034,9 @@ func (s *AssetStore) QueryByType(typ string, limit, offset int) ([]*Asset, error
if limit <= 0 {
limit = 50
}
if offset < 0 {
offset = 0
}
rows, err := s.db.Query(`
SELECT id, type, company_id, array_to_json(task_ids)::text,
COALESCE(domain,''), COALESCE(root_domain,''), COALESCE(ip,''),
@@ -896,7 +1050,7 @@ SELECT id, type, company_id, array_to_json(task_ids)::text,
COALESCE(method,''), array_to_json(params)::text, extra, last_seen::text
FROM assets
WHERE type = $1
ORDER BY last_seen DESC
ORDER BY last_seen DESC, id DESC
LIMIT $2 OFFSET $3`, typ, limit, offset)
if err != nil {
return nil, err
@@ -913,10 +1067,8 @@ func (s *AssetStore) CountByType(typ string) (int, error) {
}
// QueryByCompany returns assets for a company, optionally filtered by type.
func (s *AssetStore) QueryByCompany(companyID int64, typ string, limit int) ([]*Asset, error) {
if limit <= 0 {
limit = 200
}
// limit <= 0 means no limit.
func (s *AssetStore) QueryByCompany(companyID int64, typ string, limit, offset int) ([]*Asset, error) {
q := `SELECT id, type, company_id, array_to_json(task_ids)::text,
COALESCE(domain,''), COALESCE(root_domain,''), COALESCE(ip,''),
COALESCE(c_segment::text,''), port,
@@ -933,8 +1085,7 @@ FROM assets WHERE company_id = $1`
args = append(args, typ)
q += fmt.Sprintf(` AND type = $%d`, len(args))
}
args = append(args, limit)
q += fmt.Sprintf(` ORDER BY last_seen DESC LIMIT $%d`, len(args))
q += pageClause(&args, limit, offset)
rows, err := s.db.Query(q, args...)
if err != nil {
return nil, err
@@ -943,11 +1094,33 @@ FROM assets WHERE company_id = $1`
return scanAssets(rows)
}
// QueryByTask returns assets rows that have a given task_id in task_ids.
func (s *AssetStore) QueryByTask(taskID int64, typ string, limit int) ([]*Asset, error) {
if limit <= 0 {
limit = 200
func (s *AssetStore) CountByCompany(companyID int64, typ string) (int, error) {
q := `SELECT count(*) FROM assets WHERE company_id = $1`
args := []any{companyID}
if typ != "" {
args = append(args, typ)
q += fmt.Sprintf(` AND type = $%d`, len(args))
}
var n int
err := s.db.QueryRow(q, args...).Scan(&n)
return n, err
}
func pageClause(args *[]any, limit, offset int) string {
q := ` ORDER BY last_seen DESC, id DESC`
if limit > 0 {
*args = append(*args, limit)
q += fmt.Sprintf(` LIMIT $%d`, len(*args))
}
if offset > 0 {
*args = append(*args, offset)
q += fmt.Sprintf(` OFFSET $%d`, len(*args))
}
return q
}
// QueryByTask returns assets rows that have a given task_id in task_ids.
func (s *AssetStore) QueryByTask(taskID int64, typ string, limit, offset int) ([]*Asset, error) {
q := `SELECT id, type, company_id, array_to_json(task_ids)::text,
COALESCE(domain,''), COALESCE(root_domain,''), COALESCE(ip,''),
COALESCE(c_segment::text,''), port,
@@ -964,14 +1137,186 @@ FROM assets WHERE $1 = ANY(task_ids)`
args = append(args, typ)
q += fmt.Sprintf(` AND type = $%d`, len(args))
}
args = append(args, limit)
q += fmt.Sprintf(` ORDER BY last_seen DESC LIMIT $%d`, len(args))
q += pageClause(&args, limit, offset)
rows, err := s.db.Query(q, args...)
if err != nil {
return nil, err
}
defer rows.Close()
return scanAssets(rows)
assets, err := scanAssets(rows)
if err != nil {
return nil, err
}
if err := s.hydrateTaskAssetSources(taskID, assets); err != nil {
return nil, err
}
return assets, nil
}
func (s *AssetStore) CountByTask(taskID int64, typ string) (int, error) {
q := `SELECT count(*) FROM assets WHERE $1 = ANY(task_ids)`
args := []any{taskID}
if typ != "" {
args = append(args, typ)
q += fmt.Sprintf(` AND type = $%d`, len(args))
}
var n int
err := s.db.QueryRow(q, args...).Scan(&n)
return n, err
}
func (s *AssetStore) CountsByTypeForTask(taskID int64) (map[string]int, error) {
rows, err := s.db.Query(`SELECT type, COUNT(*) FROM assets WHERE $1 = ANY(task_ids) GROUP BY type`, taskID)
if err != nil {
return nil, err
}
defer rows.Close()
return scanTypeCounts(rows)
}
// DeleteByTaskID removes assets owned only by taskID and detaches taskID from
// assets shared with other tasks. Full task deletion uses the coordinated
// transaction in DeleteTaskCascadePrepared; this method remains for callers
// that explicitly manage only asset associations.
func (s *AssetStore) DeleteByTaskID(taskID int64) (int64, error) {
res, err := s.db.Exec(`DELETE FROM assets WHERE task_ids = ARRAY[$1]::bigint[]`, taskID)
if err != nil {
return 0, err
}
deleted, _ := res.RowsAffected()
if _, err := s.db.Exec(`UPDATE assets SET task_ids = array_remove(task_ids, $1) WHERE $1 = ANY(task_ids)`, taskID); err != nil {
return deleted, err
}
return deleted, nil
}
// HostsByTask returns the exact HTTP host candidates attached to a task's
// assets. Domain/IP columns cover root domains, subdomains and non-HTTP
// services; URL covers HTTP services and endpoints.
func (s *AssetStore) HostsByTask(taskID int64) ([]string, error) {
rows, err := s.db.Query(`
SELECT COALESCE(domain,''), COALESCE(ip,''), COALESCE(url,'')
FROM assets WHERE $1 = ANY(task_ids)`, taskID)
if err != nil {
return nil, err
}
defer rows.Close()
hosts := make(map[string]struct{})
add := func(host string) {
host = strings.TrimSpace(strings.ToLower(host))
if host != "" {
hosts[host] = struct{}{}
}
}
for rows.Next() {
var domain, ip, rawURL string
if err := rows.Scan(&domain, &ip, &rawURL); err != nil {
return nil, err
}
add(domain)
add(ip)
if u, err := url.Parse(rawURL); err == nil {
add(u.Hostname())
}
}
if err := rows.Err(); err != nil {
return nil, err
}
out := make([]string, 0, len(hosts))
for host := range hosts {
out = append(out, host)
}
sort.Strings(out)
return out, nil
}
// HostsForTaskDeletion returns hosts that belong to the task being deleted and
// are not referenced by any other live task. Current-task candidates include
// both task_ids ownership and exploration anchors so legacy seeded assets (which
// were anchor-only) are covered. Protection is host-wide: if another live task
// references any asset for a candidate host, that host's global traffic remains.
func (s *AssetStore) HostsForTaskDeletion(taskID, explorationID int64) ([]string, error) {
return hostsForTaskDeletion(s.db, taskID, explorationID)
}
type rowsQuerier interface {
Query(query string, args ...any) (*sql.Rows, error)
}
// hostsForTaskDeletion is shared by the read-only AssetStore API and the task
// deletion transaction. Coordinated traffic deletion must call it through the
// transaction path so asset/anchor writes stay locked until the task delete is
// committed.
func hostsForTaskDeletion(q rowsQuerier, taskID, explorationID int64) ([]string, error) {
rows, err := q.Query(`
WITH current_assets AS (
SELECT id FROM assets WHERE $1 = ANY(task_ids)
UNION
SELECT ea.asset_id
FROM exploration_anchors ea
JOIN exploration_nodes n ON n.id=ea.node_id
WHERE n.exploration_id=$2
),
other_assets AS (
SELECT DISTINCT a.id
FROM assets a
WHERE EXISTS (
SELECT 1 FROM tasks t
WHERE t.id<>$1 AND t.deleted_at IS NULL AND t.id=ANY(a.task_ids)
) OR EXISTS (
SELECT 1
FROM exploration_anchors ea
JOIN exploration_nodes n ON n.id=ea.node_id
JOIN tasks t ON t.exploration_id=n.exploration_id
WHERE ea.asset_id=a.id AND t.id<>$1 AND t.deleted_at IS NULL
)
)
SELECT COALESCE(a.domain,''), COALESCE(a.ip,''), COALESCE(a.url,''), true
FROM assets a JOIN current_assets c ON c.id=a.id
UNION ALL
SELECT COALESCE(a.domain,''), COALESCE(a.ip,''), COALESCE(a.url,''), false
FROM assets a JOIN other_assets o ON o.id=a.id`, taskID, explorationID)
if err != nil {
return nil, err
}
defer rows.Close()
candidates := make(map[string]struct{})
protected := make(map[string]struct{})
add := func(dst map[string]struct{}, raw string) {
raw = strings.TrimSpace(strings.ToLower(raw))
if raw != "" {
dst[raw] = struct{}{}
}
}
for rows.Next() {
var domain, ip, rawURL string
var candidate bool
if err := rows.Scan(&domain, &ip, &rawURL, &candidate); err != nil {
return nil, err
}
dst := protected
if candidate {
dst = candidates
}
add(dst, domain)
add(dst, ip)
if u, err := url.Parse(rawURL); err == nil {
add(dst, u.Hostname())
}
}
if err := rows.Err(); err != nil {
return nil, err
}
out := make([]string, 0, len(candidates))
for host := range candidates {
if _, shared := protected[host]; !shared {
out = append(out, host)
}
}
sort.Strings(out)
return out, nil
}
const assetSelectCols = `SELECT id, type, company_id, array_to_json(task_ids)::text,
@@ -1072,6 +1417,10 @@ func (s *AssetStore) CountsByType() (map[string]int, error) {
return nil, err
}
defer rows.Close()
return scanTypeCounts(rows)
}
func scanTypeCounts(rows *sql.Rows) (map[string]int, error) {
out := map[string]int{}
for rows.Next() {
var typ string
db/assets_test.go
+250 -7
View File
@@ -1,7 +1,10 @@
package db
import (
"fmt"
"slices"
"testing"
"time"
)
// testSetup opens a DB and returns both stores. Skips if no PG.
@@ -288,12 +291,12 @@ func TestUpsertHTTPService(t *testing.T) {
sc := 200
cl := int64(1024)
id1, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{
URL: "https://www.httptest.example.com/",
Technologies: []string{"nginx", "vue"},
StatusCode: &sc,
URL: "https://www.httptest.example.com/",
Technologies: []string{"nginx", "vue"},
StatusCode: &sc,
ContentLength: &cl,
PageTitle: "Test Site",
TaskID: 1,
PageTitle: "Test Site",
TaskID: 1,
})
if err != nil {
t.Fatal(err)
@@ -534,6 +537,59 @@ func TestQueryByType(t *testing.T) {
}
}
// TestDeleteByTaskID: 独有资产被删,与其他任务共享的资产仅解除关联(保留),host 反查正确。
func TestDeleteByTaskID(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
const taskA = int64(90001)
const taskB = int64(90002)
// solo:仅属 taskA
solo, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "solo-del.test", TaskID: taskA})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, solo)
// shared:先 taskA 再 taskB → task_ids={A,B}
shared, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "shared-del.test", TaskID: taskA})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, shared)
if _, err := av2.UpsertRootDomain(UpsertRootDomainReq{Domain: "shared-del.test", TaskID: taskB}); err != nil {
t.Fatal(err)
}
// host 反查(删资产前):应含两个域名
hosts, err := av2.HostsByTask(taskA)
if err != nil {
t.Fatal(err)
}
if !slices.Contains(hosts, "solo-del.test") || !slices.Contains(hosts, "shared-del.test") {
t.Fatalf("HostsByTask 缺 host: %v", hosts)
}
n, err := av2.DeleteByTaskID(taskA)
if err != nil {
t.Fatal(err)
}
if n != 1 {
t.Fatalf("DeleteByTaskID: 应删 1 个独有资产,实删 %d", n)
}
// solo 已删
if a, _ := av2.GetByIDs([]int64{solo}); len(a) != 0 {
t.Fatalf("solo 资产应被删除")
}
// shared 保留,且 task_ids 只剩 taskB
sa, _ := av2.GetByIDs([]int64{shared})
if len(sa) != 1 {
t.Fatalf("shared 资产应保留")
}
if slices.Contains(sa[0].TaskIDs, taskA) || !slices.Contains(sa[0].TaskIDs, taskB) {
t.Fatalf("shared task_ids 应解除 A 保留 B,得 %v", sa[0].TaskIDs)
}
}
func TestQueryByTask(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
@@ -545,7 +601,7 @@ func TestQueryByTask(t *testing.T) {
}
defer deleteAsset(d, id)
assets, err := av2.QueryByTask(taskID, "", 10)
assets, err := av2.QueryByTask(taskID, "", 10, 0)
if err != nil {
t.Fatal(err)
}
@@ -558,6 +614,73 @@ func TestQueryByTask(t *testing.T) {
if !found {
t.Error("QueryByTask: inserted asset not found")
}
n, err := av2.CountByTask(taskID, "")
if err != nil {
t.Fatal(err)
}
if n < 1 {
t.Errorf("CountByTask: got %d, want >= 1", n)
}
if counts, err := av2.CountsByTypeForTask(taskID); err != nil {
t.Fatal(err)
} else if counts["root_domain"] < 1 {
t.Errorf("CountsByTypeForTask: root_domain = %d, want >= 1", counts["root_domain"])
}
// offset past the end must return nothing, not fall back to page 1
rest, err := av2.QueryByTask(taskID, "", 10, n)
if err != nil {
t.Fatal(err)
}
if len(rest) != 0 {
t.Errorf("QueryByTask offset=%d: got %d rows, want 0", n, len(rest))
}
}
// 任务资产列表按页取,不再被固定条数截断:60 条资产用 25/页要能完整翻出来。
func TestQueryByTaskPaging(t *testing.T) {
d, av2, _ := testSetup(t)
defer d.Close()
const taskID = int64(99998)
const n = 60
for i := 0; i < n; i++ {
id, err := av2.UpsertRootDomain(UpsertRootDomainReq{
Domain: fmt.Sprintf("paging-%02d.querybytask.test", i),
TaskID: taskID,
})
if err != nil {
t.Fatal(err)
}
defer deleteAsset(d, id)
}
total, err := av2.CountByTask(taskID, "root_domain")
if err != nil {
t.Fatal(err)
}
if total != n {
t.Fatalf("CountByTask: got %d, want %d", total, n)
}
const size = 25
seen := map[int64]bool{}
for offset := 0; offset < total; offset += size {
page, err := av2.QueryByTask(taskID, "root_domain", size, offset)
if err != nil {
t.Fatal(err)
}
for _, a := range page {
if seen[a.ID] {
t.Errorf("offset=%d: asset %d returned twice", offset, a.ID)
}
seen[a.ID] = true
}
}
if len(seen) != n {
t.Errorf("paged through %d assets, want %d", len(seen), n)
}
}
func TestQueryByCompany(t *testing.T) {
@@ -582,7 +705,7 @@ func TestQueryByCompany(t *testing.T) {
t.Fatal(err)
}
assets, err := av2.QueryByCompany(companyID, "root_domain", 10)
assets, err := av2.QueryByCompany(companyID, "root_domain", 10, 0)
if err != nil {
t.Fatal(err)
}
@@ -595,6 +718,77 @@ func TestQueryByCompany(t *testing.T) {
if !found {
t.Error("QueryByCompany: attributed asset not found")
}
n, err := av2.CountByCompany(companyID, "root_domain")
if err != nil {
t.Fatal(err)
}
if n < 1 {
t.Errorf("CountByCompany: got %d, want >= 1", n)
}
// offset past the end must return nothing, not fall back to page 1
rest, err := av2.QueryByCompany(companyID, "root_domain", 10, n)
if err != nil {
t.Fatal(err)
}
if len(rest) != 0 {
t.Errorf("QueryByCompany offset=%d: got %d rows, want 0", n, len(rest))
}
}
// 企业资产列表同样按页取,不被固定条数截断。
func TestQueryByCompanyPaging(t *testing.T) {
d, av2, cs := testSetup(t)
defer d.Close()
companyID, _, err := cs.UpsertCompany("QueryByCompanyPagingCorp", "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, companyID)
if added, _, _, errs := cs.AddScope(companyID, []string{"qbc-paging.io"}, "test"); added != 1 {
t.Fatalf("AddScope: added=%d, errors=%v", added, errs)
}
// UpsertSubdomain 会顺带建根域名资产,一并清掉
defer d.Exec(`DELETE FROM assets WHERE root_domain = 'qbc-paging.io'`)
const n = 60
for i := 0; i < n; i++ {
if _, err := av2.UpsertSubdomain(UpsertSubdomainReq{
Domain: fmt.Sprintf("paging-%02d.qbc-paging.io", i),
}); err != nil {
t.Fatal(err)
}
}
total, err := av2.CountByCompany(companyID, "subdomain")
if err != nil {
t.Fatal(err)
}
if total != n {
t.Fatalf("CountByCompany: got %d, want %d", total, n)
}
const size = 25
seen := map[int64]bool{}
for offset := 0; offset < total; offset += size {
page, err := av2.QueryByCompany(companyID, "subdomain", size, offset)
if err != nil {
t.Fatal(err)
}
for _, a := range page {
if seen[a.ID] {
t.Errorf("offset=%d: asset %d returned twice", offset, a.ID)
}
seen[a.ID] = true
}
}
if len(seen) != n {
t.Errorf("paged through %d assets, want %d", len(seen), n)
}
}
func TestCountsByType(t *testing.T) {
@@ -703,3 +897,52 @@ func TestNormalizeURL(t *testing.T) {
}
}
}
func TestAssetPaginationUsesStableIDTieBreaker(t *testing.T) {
d, assets, _ := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
taskID := stamp
marker := fmt.Sprintf("stable-page-%d", stamp)
sharedSeen := time.Date(2026, time.August, 21, 9, 0, 0, 0, time.UTC)
ids := make([]int64, 0, 5)
for i := 0; i < 5; i++ {
var id int64
domain := fmt.Sprintf("%s-%d.invalid", marker, i)
if err := d.QueryRow(`INSERT INTO assets(type,domain,root_domain,task_ids,last_seen)
VALUES ('root_domain',$1,$1,ARRAY[$2]::bigint[],$3) RETURNING id`,
domain, taskID, sharedSeen).Scan(&id); err != nil {
t.Fatal(err)
}
ids = append(ids, id)
}
defer func() { _, _ = d.Exec(`DELETE FROM assets WHERE id=ANY($1::bigint[])`, ids) }()
want := slices.Clone(ids)
slices.Reverse(want)
collect := func(query func(limit, offset int) ([]*Asset, error)) []int64 {
t.Helper()
var got []int64
for offset := 0; offset < len(ids); offset += 2 {
page, err := query(2, offset)
if err != nil {
t.Fatal(err)
}
for _, asset := range page {
got = append(got, asset.ID)
}
}
return got
}
if got := collect(func(limit, offset int) ([]*Asset, error) {
return assets.QueryByTask(taskID, "root_domain", limit, offset)
}); !slices.Equal(got, want) {
t.Fatalf("task pagination order=%v want=%v", got, want)
}
if got := collect(func(limit, offset int) ([]*Asset, error) {
return assets.QueryDSL(marker, "root_domain", limit, offset)
}); !slices.Equal(got, want) {
t.Fatalf("DSL pagination order=%v want=%v", got, want)
}
}
db/companies.go
+462 -109
View File
@@ -2,9 +2,11 @@ package db
import (
"database/sql"
"errors"
"fmt"
"net"
"strings"
"unicode/utf8"
)
// =====================================================================
@@ -25,7 +27,7 @@ type Company struct {
type CompanyWithScope struct {
Company
Scope []ScopeRule `json:"scope"`
AssetCount int `json:"asset_count"`
AssetCount int `json:"asset_count"`
}
// ScopeRule is one company_scope row.
@@ -35,6 +37,7 @@ type ScopeRule struct {
Kind string `json:"kind"`
Domain string `json:"domain,omitempty"`
Net string `json:"net,omitempty"`
Value string `json:"value,omitempty"`
Raw string `json:"raw"`
Reason string `json:"reason,omitempty"`
}
@@ -42,6 +45,51 @@ type ScopeRule struct {
// CompanyStore operates on the companies + company_scope tables.
type CompanyStore struct{ db *DB }
var (
ErrCompanyNameConflict = errors.New("company name already exists")
ErrCompanyNotFound = errors.New("company not found")
)
const (
// MaxCompanyScopeRules bounds both request processing and the total number
// of persisted rules for one company.
MaxCompanyScopeRules = 256
// Raw and normalized textual scope payloads are bounded by Unicode rune
// count so multi-byte input is treated consistently by the API and DB layer.
MaxCompanyScopeRawRunes = 1024
MaxCompanyScopeValueRunes = 1024
)
// CompanyScopeValidationError identifies a client-correctable scope error.
// Storage and transaction failures are returned as ordinary errors instead.
type CompanyScopeValidationError struct{ Message string }
func (e *CompanyScopeValidationError) Error() string { return e.Message }
// ValidateCompanyScopeInputBounds applies request-wide limits before parsing.
// Store methods call it again so non-HTTP callers cannot bypass the limits.
func ValidateCompanyScopeInputBounds(inputs []ScopeInput) error {
if len(inputs) > MaxCompanyScopeRules {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围规则过多: 最多 %d 条", MaxCompanyScopeRules,
)}
}
for i, input := range inputs {
if utf8.RuneCountInString(input.Value) > MaxCompanyScopeRawRunes {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围第 %d 条原始值过长: 最多 %d 个字符", i+1, MaxCompanyScopeRawRunes,
)}
}
}
return nil
}
// Scope writes rebuild derived asset ownership globally, so serialize them to
// ensure the committed attribution always reflects the latest committed rules.
// This key is reserved for company mutations; 7337741001 is the schema lock and
// 7337741002 is the cross-package test-suite lock.
const companyScopeMutationLock int64 = 7337741003
// Companies returns the company store.
func (d *DB) Companies() *CompanyStore { return &CompanyStore{db: d} }
@@ -69,6 +117,60 @@ RETURNING id, (xmax = 0)`, name, nkey, logoVal).Scan(&id, &created)
return
}
// CreateCompanyWithScope creates a company without updating an existing row.
// The company, its valid initial scope rules, and derived asset attribution are
// committed atomically. Invalid inputs retain the legacy partial-validation
// contract and are reported without preventing valid rules from being stored.
func (s *CompanyStore) CreateCompanyWithScope(name, logo string, inputs []ScopeInput, reason string) (
id int64, added, skipped, invalid int, validationErrors []string, err error,
) {
if err := ValidateCompanyScopeInputBounds(inputs); err != nil {
return 0, 0, 0, 0, nil, err
}
rules, invalid, validationErrors := parseScopeInputs(inputs)
if err := validateParsedScopeBounds(rules); err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
tx, err := s.db.Begin()
if err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
nkey := companyNKey(name)
var logoVal any
if logo != "" {
logoVal = logo
}
if err := tx.QueryRow(`
INSERT INTO companies(name, nkey, logo)
VALUES ($1, $2, $3)
ON CONFLICT (nkey) DO NOTHING
RETURNING id`, name, nkey, logoVal).Scan(&id); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return 0, 0, 0, invalid, validationErrors, ErrCompanyNameConflict
}
return 0, 0, 0, invalid, validationErrors, err
}
added, skipped, needsAttribution, err := insertScopeRulesTx(tx, id, rules, reason)
if err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
if needsAttribution {
if err := recomputeAttributionTx(tx); err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
}
if err := tx.Commit(); err != nil {
return 0, 0, 0, invalid, validationErrors, err
}
return id, added, skipped, invalid, validationErrors, nil
}
// GetCompany returns one company by id (nil if not found).
func (s *CompanyStore) GetCompany(id int64) (*Company, error) {
c := &Company{}
@@ -102,12 +204,55 @@ func (s *CompanyStore) UpsertByName(name string) (int64, error) {
return id, err
}
// DeleteCompany deletes a company (cascades to scope rules; assets.company_id → NULL via FK).
// DeleteCompany deletes a company and re-evaluates automatic ownership against
// the remaining companies in the same transaction. Explicitly-owned assets are
// detached by the FK and may then fall back to a remaining scope match.
func (s *CompanyStore) DeleteCompany(id int64) error {
_, err := s.db.Exec(`DELETE FROM companies WHERE id = $1`, id)
_, err := s.DeleteCompanyWithAssets(id, false)
return err
}
// DeleteCompanyWithAssets deletes a company and optionally all of its assets in
// one transaction, then re-evaluates ownership against the remaining companies.
func (s *CompanyStore) DeleteCompanyWithAssets(id int64, deleteAssets bool) (assetsDeleted int64, err error) {
tx, err := s.db.Begin()
if err != nil {
return 0, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, err
}
if deleteAssets {
res, err := tx.Exec(`DELETE FROM assets WHERE company_id = $1`, id)
if err != nil {
return 0, err
}
assetsDeleted, err = res.RowsAffected()
if err != nil {
return 0, err
}
}
res, err := tx.Exec(`DELETE FROM companies WHERE id = $1`, id)
if err != nil {
return 0, err
}
deleted, err := res.RowsAffected()
if err != nil {
return 0, err
}
if deleted == 0 {
return 0, ErrCompanyNotFound
}
if err := recomputeAttributionTx(tx); err != nil {
return 0, err
}
if err := tx.Commit(); err != nil {
return 0, err
}
return assetsDeleted, nil
}
// ListCompanies returns all companies with scope and asset count.
func (s *CompanyStore) ListCompanies() ([]*CompanyWithScope, error) {
rows, err := s.db.Query(`
@@ -149,7 +294,7 @@ ORDER BY c.name`)
func (s *CompanyStore) GetScope(companyID int64) ([]ScopeRule, error) {
rows, err := s.db.Query(`
SELECT id, company_id, kind,
COALESCE(domain,''), COALESCE(net::text,''), raw, COALESCE(reason,'')
COALESCE(domain,''), COALESCE(net::text,''), COALESCE(value,''), raw, COALESCE(reason,'')
FROM company_scope
WHERE company_id = $1
ORDER BY id`, companyID)
@@ -157,10 +302,10 @@ ORDER BY id`, companyID)
return nil, err
}
defer rows.Close()
var out []ScopeRule
out := make([]ScopeRule, 0)
for rows.Next() {
var r ScopeRule
if err := rows.Scan(&r.ID, &r.CompanyID, &r.Kind, &r.Domain, &r.Net, &r.Raw, &r.Reason); err != nil {
if err := rows.Scan(&r.ID, &r.CompanyID, &r.Kind, &r.Domain, &r.Net, &r.Value, &r.Raw, &r.Reason); err != nil {
return nil, err
}
out = append(out, r)
@@ -171,48 +316,172 @@ ORDER BY id`, companyID)
// AddScope parses and inserts scope lines for a company, then reattributes assets.
// Returns counts of added, skipped, and invalid lines.
func (s *CompanyStore) AddScope(companyID int64, lines []string, reason string) (added, skipped, invalid int, errors []string) {
inputs := make([]ScopeInput, 0, len(lines))
for _, line := range lines {
rule, err := ParseScopeLine(line)
if err != nil {
invalid++
errors = append(errors, fmt.Sprintf("%s: %v", line, err))
continue
}
inserted, err := s.insertScopeRule(companyID, rule, reason)
if err != nil {
invalid++
errors = append(errors, fmt.Sprintf("%s: %v", line, err))
continue
}
if inserted {
added++
} else {
skipped++
}
inputs = append(inputs, ScopeInput{Value: line})
}
// attribute newly-covered assets
if added > 0 {
_ = s.attributeByCompany(companyID)
}
return
return s.AddScopeInputs(companyID, inputs, reason)
}
// insertScopeRule inserts a scope rule. Returns (inserted, error):
// inserted=false means ON CONFLICT DO NOTHING triggered (duplicate), not an error.
func (s *CompanyStore) insertScopeRule(companyID int64, rule ParsedScope, reason string) (inserted bool, err error) {
// AddScopeInputs inserts structured scope rules. Empty kinds use the automatic
// CIDR/IP/ICP/domain/keyword classification used by AddScope.
func (s *CompanyStore) AddScopeInputs(companyID int64, inputs []ScopeInput, reason string) (added, skipped, invalid int, errors []string) {
added, skipped, invalid, validationErrors, err := s.AddScopeInputsChecked(companyID, inputs, reason)
if err != nil {
validationErrors = append(validationErrors, err.Error())
}
return added, skipped, invalid, validationErrors
}
// AddScopeInputsChecked inserts structured scope rules while keeping input
// validation separate from storage and transaction errors.
func (s *CompanyStore) AddScopeInputsChecked(companyID int64, inputs []ScopeInput, reason string) (
added, skipped, invalid int, validationErrors []string, err error,
) {
if err := ValidateCompanyScopeInputBounds(inputs); err != nil {
return 0, 0, 0, nil, err
}
rules, invalid, errors := parseScopeInputs(inputs)
if err := validateParsedScopeBounds(rules); err != nil {
return 0, 0, invalid, errors, err
}
tx, err := s.db.Begin()
if err != nil {
return 0, 0, invalid, errors, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, 0, invalid, errors, err
}
if err := ensureCompanyExistsTx(tx, companyID); err != nil {
return 0, 0, invalid, errors, err
}
if len(rules) == 0 {
if err := tx.Commit(); err != nil {
return 0, 0, invalid, errors, err
}
return 0, 0, invalid, errors, nil
}
added, skipped, needsAttribution, err := insertScopeRulesTx(tx, companyID, rules, reason)
if err != nil {
return 0, 0, invalid, errors, err
}
if err := ensureCompanyScopeCapacityTx(tx, companyID); err != nil {
return 0, 0, invalid, errors, err
}
if needsAttribution {
if err := recomputeAttributionTx(tx); err != nil {
return 0, 0, invalid, errors, fmt.Errorf("重新计算企业归属失败: %w", err)
}
}
if err := tx.Commit(); err != nil {
return 0, 0, invalid, errors, err
}
return added, skipped, invalid, errors, nil
}
func parseScopeInputs(inputs []ScopeInput) (rules []ParsedScope, invalid int, validationErrors []string) {
rules = make([]ParsedScope, 0, len(inputs))
for _, input := range inputs {
rule, err := ParseScopeInput(input)
if err != nil {
invalid++
validationErrors = append(validationErrors, fmt.Sprintf("%s: %v", input.Value, err))
continue
}
rules = append(rules, rule)
}
return rules, invalid, validationErrors
}
func validateParsedScopeBounds(rules []ParsedScope) error {
for i, rule := range rules {
if utf8.RuneCountInString(rule.Raw) > MaxCompanyScopeRawRunes {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围第 %d 条原始值过长: 最多 %d 个字符", i+1, MaxCompanyScopeRawRunes,
)}
}
if utf8.RuneCountInString(rule.Value) > MaxCompanyScopeValueRunes {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围第 %d 条规范化值过长: 最多 %d 个字符", i+1, MaxCompanyScopeValueRunes,
)}
}
}
return nil
}
func ensureCompanyExistsTx(tx *sql.Tx, companyID int64) error {
var exists bool
if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM companies WHERE id = $1)`, companyID).Scan(&exists); err != nil {
return err
}
if !exists {
return ErrCompanyNotFound
}
return nil
}
func ensureCompanyScopeCapacityTx(tx *sql.Tx, companyID int64) error {
var count int
if err := tx.QueryRow(`SELECT count(*) FROM company_scope WHERE company_id = $1`, companyID).Scan(&count); err != nil {
return err
}
if count > MaxCompanyScopeRules {
return &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围规则过多: 每个企业最多 %d 条", MaxCompanyScopeRules,
)}
}
return nil
}
func lockCompanyScopeMutation(tx *sql.Tx) error {
_, err := tx.Exec(`SELECT pg_advisory_xact_lock($1)`, companyScopeMutationLock)
return err
}
func insertScopeRulesTx(tx *sql.Tx, companyID int64, rules []ParsedScope, reason string) (
added, skipped int, needsAttribution bool, err error,
) {
for _, rule := range rules {
inserted, insertErr := insertScopeRuleTx(tx, companyID, rule, reason)
if insertErr != nil {
return 0, 0, false, insertErr
}
if !inserted {
skipped++
continue
}
added++
needsAttribution = needsAttribution || rule.Kind != "keyword"
}
return added, skipped, needsAttribution, nil
}
// insertScopeRuleTx inserts a scope rule. inserted=false means a duplicate was
// ignored by ON CONFLICT, not an error.
func insertScopeRuleTx(tx *sql.Tx, companyID int64, rule ParsedScope, reason string) (inserted bool, err error) {
var res interface{ RowsAffected() (int64, error) }
if rule.Kind == "domain" {
res, err = s.db.Exec(`
switch rule.Kind {
case "domain":
res, err = tx.Exec(`
INSERT INTO company_scope(company_id, kind, domain, raw, reason)
VALUES ($1, 'domain', $2, $3, $4)
ON CONFLICT ON CONSTRAINT uq_sv2_domain DO NOTHING`,
companyID, rule.Domain, rule.Raw, reason)
} else {
res, err = s.db.Exec(`
case "ip", "cidr":
res, err = tx.Exec(`
INSERT INTO company_scope(company_id, kind, net, raw, reason)
VALUES ($1, $2, $3::cidr, $4, $5)
ON CONFLICT ON CONSTRAINT uq_sv2_net DO NOTHING`,
companyID, rule.Kind, rule.Net, rule.Raw, reason)
case "icp", "keyword":
res, err = tx.Exec(`
INSERT INTO company_scope(company_id, kind, value, raw, reason)
VALUES ($1, $2, $3, $4, $5)
ON CONFLICT (company_id, kind, value) WHERE kind IN ('icp','keyword') DO NOTHING`,
companyID, rule.Kind, rule.Value, rule.Raw, reason)
default:
return false, fmt.Errorf("unsupported company scope kind %q", rule.Kind)
}
if err != nil {
return false, err
@@ -221,114 +490,184 @@ ON CONFLICT ON CONSTRAINT uq_sv2_net DO NOTHING`,
return n > 0, nil
}
// RecomputeAttribution clears all company attribution on assets and rebuilds
// from scope rules. domain rules use root_domain equality; ip/cidr use INET >>= .
// RecomputeAttribution rebuilds only scope-derived ownership. Explicit company
// links are immutable under scope edits. Precedence is domain, IP/CIDR, then
// normalized exact ICP; keyword rules never attribute assets.
func (s *CompanyStore) RecomputeAttribution() error {
tx, err := s.db.Begin()
if err != nil {
return err
}
defer tx.Rollback() //nolint:errcheck
// step 1: clear all attribution
if _, err := tx.Exec(`UPDATE assets SET company_id = NULL`); err != nil {
if err := lockCompanyScopeMutation(tx); err != nil {
return err
}
// step 2: domain-based attribution (root_domain exact match)
if _, err := tx.Exec(`
UPDATE assets a SET company_id = (
SELECT cs.company_id FROM company_scope cs
WHERE cs.kind = 'domain'
AND a.root_domain = cs.domain
ORDER BY length(cs.domain) DESC
LIMIT 1
)
WHERE type IN ('root_domain','subdomain','service','endpoint')
AND root_domain IS NOT NULL`); err != nil {
if err := recomputeAttributionTx(tx); err != nil {
return err
}
// step 3: ip/cidr attribution (only for assets not yet attributed)
if _, err := tx.Exec(`
UPDATE assets a SET company_id = (
SELECT cs.company_id FROM company_scope cs
WHERE cs.kind IN ('ip','cidr')
AND cs.net >>= a.ip::inet
ORDER BY masklen(cs.net) DESC
LIMIT 1
)
WHERE type IN ('ip','subdomain','service','endpoint')
AND ip IS NOT NULL
AND company_id IS NULL`); err != nil {
return err
}
return tx.Commit()
}
// attributeByCompany reattributes only assets matching a specific company's scope.
// Cheaper than full recompute; used when a company's scope is updated.
func (s *CompanyStore) attributeByCompany(companyID int64) error {
tx, err := s.db.Begin()
if err != nil {
return err
}
defer tx.Rollback() //nolint:errcheck
// clear this company's existing claims
if _, err := tx.Exec(`UPDATE assets SET company_id = NULL WHERE company_id = $1`, companyID); err != nil {
return err
}
// domain rules
func recomputeAttributionTx(tx *sql.Tx) error {
// Only derived rows are cleared. Historical rows migrated without provenance
// are marked explicit by schema.sql, which is the non-destructive default.
if _, err := tx.Exec(`
UPDATE assets a SET company_id = $1
FROM company_scope cs
WHERE cs.company_id = $1
AND cs.kind = 'domain'
AND a.root_domain = cs.domain
AND type IN ('root_domain','subdomain','service','endpoint')
AND a.company_id IS NULL`, companyID); err != nil {
UPDATE assets
SET company_id = NULL, company_source = 'scope'
WHERE company_source = 'scope'`); err != nil {
return err
}
// ip/cidr rules
// Domain-based attribution (root_domain exact match).
if _, err := tx.Exec(`
UPDATE assets a SET company_id = $1
FROM company_scope cs
WHERE cs.company_id = $1
AND cs.kind IN ('ip','cidr')
AND cs.net >>= a.ip::inet
AND type IN ('ip','subdomain','service','endpoint')
AND a.company_id IS NULL`, companyID); err != nil {
WITH matched AS (
SELECT DISTINCT ON (a.id) a.id AS asset_id, cs.company_id
FROM assets a
JOIN company_scope cs ON cs.kind = 'domain' AND a.root_domain = cs.domain
WHERE a.company_id IS NULL
AND a.type IN ('root_domain','subdomain','service','endpoint')
AND a.root_domain IS NOT NULL
ORDER BY a.id, length(cs.domain) DESC, cs.company_id
)
UPDATE assets a
SET company_id = matched.company_id, company_source = 'scope'
FROM matched
WHERE a.id = matched.asset_id`); err != nil {
return err
}
return tx.Commit()
// IP/CIDR attribution for still-unowned assets.
if _, err := tx.Exec(`
WITH matched AS (
SELECT DISTINCT ON (a.id) a.id AS asset_id, cs.company_id
FROM assets a
JOIN company_scope cs ON cs.kind IN ('ip','cidr') AND cs.net >>= a.ip::inet
WHERE a.company_id IS NULL
AND a.type IN ('ip','subdomain','service','endpoint')
AND a.ip IS NOT NULL
ORDER BY a.id, masklen(cs.net) DESC, cs.company_id
)
UPDATE assets a
SET company_id = matched.company_id, company_source = 'scope'
FROM matched
WHERE a.id = matched.asset_id`); err != nil {
return err
}
// Exact normalized ICP attribution after domain/network precedence.
if _, err := tx.Exec(`
WITH matched AS (
SELECT DISTINCT ON (a.id) a.id AS asset_id, cs.company_id
FROM assets a
JOIN company_scope cs ON cs.kind = 'icp'
AND (
lower(regexp_replace(COALESCE(a.icp,''), '[[:space:]]+', '', 'g')) = cs.value
OR lower(regexp_replace(COALESCE(a.app_icp,''), '[[:space:]]+', '', 'g')) = cs.value
)
WHERE a.company_id IS NULL
AND (COALESCE(a.icp,'') <> '' OR COALESCE(a.app_icp,'') <> '')
ORDER BY a.id, cs.company_id
)
UPDATE assets a
SET company_id = matched.company_id, company_source = 'scope'
FROM matched
WHERE a.id = matched.asset_id`); err != nil {
return err
}
return nil
}
// UpdateScope replaces all scope rules for a company and reattributes.
func (s *CompanyStore) UpdateScope(companyID int64, lines []string, reason string) (added, invalid int, errs []string) {
// delete existing rules
if _, err := s.db.Exec(`DELETE FROM company_scope WHERE company_id = $1`, companyID); err != nil {
errs = append(errs, err.Error())
return
inputs := make([]ScopeInput, 0, len(lines))
for _, line := range lines {
inputs = append(inputs, ScopeInput{Value: line})
}
added, _, invalid, errs = s.AddScope(companyID, lines, reason)
return
return s.UpdateScopeInputs(companyID, inputs, reason)
}
// UpdateScopeInputs replaces all rules with a structured set.
func (s *CompanyStore) UpdateScopeInputs(companyID int64, inputs []ScopeInput, reason string) (added, invalid int, errs []string) {
added, invalid, validationErrors, err := s.UpdateScopeInputsChecked(companyID, inputs, reason)
if err != nil {
validationErrors = append(validationErrors, err.Error())
}
return added, invalid, validationErrors
}
// UpdateScopeInputsChecked replaces all rules while separating validation
// feedback from storage and transaction failures.
func (s *CompanyStore) UpdateScopeInputsChecked(companyID int64, inputs []ScopeInput, reason string) (
added, invalid int, validationErrors []string, err error,
) {
if err := ValidateCompanyScopeInputBounds(inputs); err != nil {
return 0, 0, nil, err
}
rules, invalid, errs := parseScopeInputs(inputs)
if invalid > 0 {
return 0, invalid, errs, &CompanyScopeValidationError{Message: fmt.Sprintf(
"企业范围包含 %d 条无效规则，未覆盖原有范围", invalid,
)}
}
if err := validateParsedScopeBounds(rules); err != nil {
return 0, invalid, errs, err
}
tx, err := s.db.Begin()
if err != nil {
return 0, invalid, errs, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, invalid, errs, err
}
if err := ensureCompanyExistsTx(tx, companyID); err != nil {
return 0, invalid, errs, err
}
if _, err := tx.Exec(`DELETE FROM company_scope WHERE company_id = $1`, companyID); err != nil {
return 0, invalid, errs, err
}
added, _, _, err = insertScopeRulesTx(tx, companyID, rules, reason)
if err != nil {
return 0, invalid, errs, err
}
if err := ensureCompanyScopeCapacityTx(tx, companyID); err != nil {
return 0, invalid, errs, err
}
// Rebuild even for an empty replacement because removing the old rules may
// detach scope-derived assets or expose a lower-precedence company match.
if err := recomputeAttributionTx(tx); err != nil {
return 0, invalid, errs, fmt.Errorf("重新计算企业归属失败: %w", err)
}
if err := tx.Commit(); err != nil {
return 0, invalid, errs, err
}
return added, invalid, errs, nil
}
// ResolveCompany returns the company_id for a given root_domain and/or ip, or nil
// if no scope rule matches. Mirrors the attribution logic used at asset insert time.
func (s *CompanyStore) ResolveCompany(rootDomain, ipStr string) (*int64, error) {
return s.ResolveCompanyWithICP(rootDomain, ipStr, "")
}
// ResolveCompanyWithICP mirrors RecomputeAttribution for insert-time ownership.
// ICP is consulted only after domain and IP/CIDR fail to match.
func (s *CompanyStore) ResolveCompanyWithICP(rootDomain, ipStr, icp string) (*int64, error) {
return resolveCompanyWithICP(s.db, rootDomain, ipStr, icp)
}
type companyScopeQueryer interface {
QueryRow(query string, args ...any) *sql.Row
}
func resolveCompanyWithICP(q companyScopeQueryer, rootDomain, ipStr, icp string) (*int64, error) {
if rootDomain != "" {
var cid int64
err := s.db.QueryRow(`
err := q.QueryRow(`
SELECT company_id FROM company_scope
WHERE kind = 'domain'
AND domain = $1
ORDER BY length(domain) DESC
ORDER BY length(domain) DESC, company_id
LIMIT 1`, rootDomain).Scan(&cid)
if err == nil {
return &cid, nil
@@ -340,11 +679,11 @@ LIMIT 1`, rootDomain).Scan(&cid)
if ipStr != "" {
if net.ParseIP(ipStr) != nil {
var cid int64
err := s.db.QueryRow(`
err := q.QueryRow(`
SELECT company_id FROM company_scope
WHERE kind IN ('ip','cidr')
AND net >>= $1::inet
ORDER BY masklen(net) DESC
ORDER BY masklen(net) DESC, company_id
LIMIT 1`, ipStr).Scan(&cid)
if err == nil {
return &cid, nil
@@ -354,5 +693,19 @@ LIMIT 1`, ipStr).Scan(&cid)
}
}
}
if normalized := NormalizeICP(icp); normalized != "" {
var cid int64
err := q.QueryRow(`
SELECT company_id FROM company_scope
WHERE kind = 'icp' AND value = $1
ORDER BY company_id
LIMIT 1`, normalized).Scan(&cid)
if err == nil {
return &cid, nil
}
if err != sql.ErrNoRows {
return nil, err
}
}
return nil, nil
}
db/companies_test.go
+167
View File
@@ -1,7 +1,11 @@
package db
import (
"errors"
"fmt"
"strings"
"testing"
"time"
)
// cleanup helpers to remove test data
@@ -57,6 +61,70 @@ func TestCompanyUpsertAndGet(t *testing.T) {
}
}
func TestCreateCompanyWithScopeRejectsNormalizedDuplicate(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
stamp := time.Now().UnixNano()
name := fmt.Sprintf("Strict Company %d", stamp)
id, added, _, _, validationErrors, err := cs.CreateCompanyWithScope(name, "", []ScopeInput{
{Kind: "domain", Value: fmt.Sprintf("strict-%d.example", stamp)},
}, "test")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
if added != 1 || len(validationErrors) != 0 {
t.Fatalf("initial scope: added=%d errors=%v", added, validationErrors)
}
_, _, _, _, _, err = cs.CreateCompanyWithScope(
" "+strings.ToUpper(strings.ReplaceAll(name, " ", " "))+" ",
"",
[]ScopeInput{{Kind: "domain", Value: fmt.Sprintf("replacement-%d.example", stamp)}},
"test",
)
if !errors.Is(err, ErrCompanyNameConflict) {
t.Fatalf("duplicate create error=%v want ErrCompanyNameConflict", err)
}
scope, err := cs.GetScope(id)
if err != nil {
t.Fatal(err)
}
if len(scope) != 1 || scope[0].Domain != fmt.Sprintf("strict-%d.example", stamp) {
t.Fatalf("duplicate create changed existing scope: %+v", scope)
}
}
func TestCreateCompanyWithScopeRollsBackOnScopeWriteFailure(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
name := fmt.Sprintf("Atomic Create %d", time.Now().UnixNano())
_, _, _, _, _, err = cs.CreateCompanyWithScope(name, "", []ScopeInput{
{Kind: "keyword", Value: "invalid\x00postgres-text"},
}, "test")
if err == nil {
t.Fatal("expected scope database write to fail")
}
company, getErr := cs.GetCompanyByName(name)
if getErr != nil {
t.Fatal(getErr)
}
if company != nil {
defer cleanupCompany(d, company.ID)
t.Fatalf("company row survived failed initial scope transaction: %+v", company)
}
}
func TestCompanyScope(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
@@ -187,6 +255,58 @@ func TestUpdateScope(t *testing.T) {
if len(scope) != 1 || scope[0].Domain != "new-domain.com" {
t.Errorf("UpdateScope: expected new-domain.com only, got %+v", scope)
}
// Invalid replacement input must not turn a partial validation response into
// a destructive replacement of the existing rules.
added, invalid, validationErrors, err := cs.UpdateScopeInputsChecked(id, []ScopeInput{
{Kind: "domain", Value: "co.uk"},
}, "invalid replacement")
var validationErr *CompanyScopeValidationError
if added != 0 || invalid != 1 || len(validationErrors) != 1 || !errors.As(err, &validationErr) {
t.Fatalf("invalid replacement: added=%d invalid=%d validation=%v err=%v", added, invalid, validationErrors, err)
}
scope, err = cs.GetScope(id)
if err != nil {
t.Fatal(err)
}
if len(scope) != 1 || scope[0].Domain != "new-domain.com" {
t.Fatalf("invalid replacement changed existing scope: %+v", scope)
}
}
func TestUpdateScopeRollsBackOnInsertFailure(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
name := fmt.Sprintf("Atomic Scope Update %d", time.Now().UnixNano())
id, _, err := cs.UpsertCompany(name, "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
oldDomain := fmt.Sprintf("old-%d.example", time.Now().UnixNano())
added, _, invalid, addErrors := cs.AddScopeInputs(id, []ScopeInput{{Kind: "domain", Value: oldDomain}}, "initial")
if added != 1 || invalid != 0 || len(addErrors) != 0 {
t.Fatalf("seed scope: added=%d invalid=%d errors=%v", added, invalid, addErrors)
}
added, invalid, updateErrors := cs.UpdateScopeInputs(id, []ScopeInput{
{Kind: "keyword", Value: "invalid\x00postgres-text"},
}, "replacement")
if added != 0 || invalid != 0 || len(updateErrors) == 0 {
t.Fatalf("failed update result: added=%d invalid=%d errors=%v", added, invalid, updateErrors)
}
scope, err := cs.GetScope(id)
if err != nil {
t.Fatal(err)
}
if len(scope) != 1 || scope[0].Domain != oldDomain {
t.Fatalf("failed replacement did not preserve old scope: %+v", scope)
}
}
func TestDeleteCompany(t *testing.T) {
@@ -217,6 +337,53 @@ func TestDeleteCompany(t *testing.T) {
}
}
func TestDeleteCompanyWithAssetsDeletesBoth(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer d.Close()
cs := d.Companies()
stamp := time.Now().UnixNano()
id, _, err := cs.UpsertCompany(fmt.Sprintf("Delete Assets Company %d", stamp), "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, id)
var assetID int64
domain := fmt.Sprintf("delete-assets-%d.example", stamp)
if err := d.QueryRow(`
INSERT INTO assets(type, domain, root_domain, company_id, company_source)
VALUES ('root_domain', $1, $1, $2, 'explicit')
RETURNING id`, domain, id).Scan(&assetID); err != nil {
t.Fatal(err)
}
defer d.Exec(`DELETE FROM assets WHERE id = $1`, assetID) //nolint:errcheck
deleted, err := cs.DeleteCompanyWithAssets(id, true)
if err != nil {
t.Fatal(err)
}
if deleted != 1 {
t.Fatalf("assets deleted=%d want 1", deleted)
}
company, err := cs.GetCompany(id)
if err != nil {
t.Fatal(err)
}
if company != nil {
t.Fatalf("company still exists: %+v", company)
}
var assetsRemaining int
if err := d.QueryRow(`SELECT COUNT(*) FROM assets WHERE id = $1`, assetID).Scan(&assetsRemaining); err != nil {
t.Fatal(err)
}
if assetsRemaining != 0 {
t.Fatalf("asset %d survived company deletion", assetID)
}
}
func TestRecomputeAttribution(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
db/company_lock_test.go
+15
View File
@@ -0,0 +1,15 @@
package db
import "testing"
func TestCompanyScopeMutationLockDoesNotReuseInfrastructureLocks(t *testing.T) {
reserved := map[string]int64{
"schema migration": 7337741001,
"cross-package test suite": 7337741002,
}
for name, key := range reserved {
if companyScopeMutationLock == key {
t.Fatalf("company scope mutation lock reuses the %s advisory lock key %d", name, key)
}
}
}
db/company_scope.go
+26 -40
View File
@@ -3,22 +3,24 @@ package db
import (
"fmt"
"net"
"net/url"
"strings"
"unicode"
"golang.org/x/net/idna"
"golang.org/x/net/publicsuffix"
)
// ParsedScope is one parsed asset-scope entry. Exactly one of Domain / Net is set.
// Used internally by ParseScopeLine / ParseScopeLines (passed to CompanyStore).
// ParsedScope is one parsed asset-scope entry. Its kind selects Domain, Net, or Value.
// Used internally by the company scope parsers and CompanyStore.
type ParsedScope struct {
Kind string // "domain" | "ip" | "cidr"
Kind string // "domain" | "ip" | "cidr" | "icp" | "keyword"
Domain string // normalized registrable/root domain (kind=domain)
Net string // normalized CIDR, single IP as /32 or /128 (kind=ip|cidr)
Value string // normalized text (kind=icp|keyword)
Raw string // original input line
}
<<<<<<< Updated upstream
=======
// ScopeInput is the structured API form for a company scope rule. Empty Kind
// uses the same automatic classification as the single-textarea UI.
type ScopeInput struct {
@@ -210,7 +212,6 @@ func scopeHostname(raw string) (string, error) {
return host, nil
}
>>>>>>> Stashed changes
// ParseScopeLine classifies and validates one scope line (root domain / IP /
// CIDR). Guardrails reject bare TLDs and over-broad networks so a rule can never
// swallow the internet. IP ranges must be expressed as CIDR.
@@ -220,12 +221,8 @@ func ParseScopeLine(line string) (ParsedScope, error) {
if raw == "" {
return r, fmt.Errorf("空行")
}
s := raw
if i := strings.Index(s, "://"); i >= 0 { // strip scheme if a URL was pasted
s = s[i+3:]
}
// CIDR first — it contains '/', which the path-strip below would remove.
if _, ipnet, err := net.ParseCIDR(s); err == nil {
// CIDR first because URL parsing treats its slash as a path separator.
if _, ipnet, err := net.ParseCIDR(raw); err == nil {
ones, bits := ipnet.Mask.Size()
if bits == 32 && ones < 16 {
return r, fmt.Errorf("网段过宽(IPv4 需 >= /16): %s", raw)
@@ -236,12 +233,8 @@ func ParseScopeLine(line string) (ParsedScope, error) {
r.Kind, r.Net = "cidr", ipnet.String()
return r, nil
}
// Strip any path and :port remnants for host-like inputs.
if i := strings.IndexByte(s, '/'); i >= 0 {
s = s[:i]
}
// Single IP.
if ip := net.ParseIP(strings.TrimSpace(s)); ip != nil {
if ip := net.ParseIP(raw); ip != nil {
r.Kind = "ip"
if ip.To4() != nil {
r.Net = ip.String() + "/32"
@@ -250,37 +243,30 @@ func ParseScopeLine(line string) (ParsedScope, error) {
}
return r, nil
}
if host, _, ok := strings.Cut(s, ":"); ok { // host:port → host
s = host
host, err := scopeHostname(raw)
if err != nil {
return r, fmt.Errorf("无法识别为有效域名/IP/CIDR: %s", raw)
}
if ip := net.ParseIP(host); ip != nil {
r.Kind = "ip"
if ip.To4() != nil {
r.Net = ip.String() + "/32"
} else {
r.Net = ip.String() + "/128"
}
return r, nil
}
if looksLikeIPAddress(host) {
return r, fmt.Errorf("无效 IP: %s", raw)
}
if strings.Contains(raw, "-") && strings.Count(raw, ".") >= 6 {
return r, fmt.Errorf("IP 段请用 CIDR 表示(如 1.2.3.0/24): %s", raw)
}
// Domain (registrable). Reject bare TLDs / public suffixes.
d := DomainKey(s)
if d == "" || !strings.Contains(d, ".") {
return r, fmt.Errorf("无法识别为根域名/IP/CIDR: %s", raw)
}
d := DomainKey(host)
if suf, icann := publicsuffix.PublicSuffix(d); icann && suf == d {
return r, fmt.Errorf("不能用裸 TLD 作为范围: %s", raw)
}
r.Kind, r.Domain = "domain", d
return r, nil
}
// ParseScopeLines parses a whole text block (one entry per line), returning the
// valid rules and the invalid lines (with reasons) so the caller can report both.
func ParseScopeLines(text string) (rules []ParsedScope, invalid []string) {
for _, ln := range strings.Split(text, "\n") {
if strings.TrimSpace(ln) == "" {
continue
}
r, err := ParseScopeLine(ln)
if err != nil {
invalid = append(invalid, err.Error())
continue
}
rules = append(rules, r)
}
return rules, invalid
}
db/company_scope_consistency_test.go
+248
View File
@@ -0,0 +1,248 @@
package db
import (
"errors"
"fmt"
"strings"
"testing"
"time"
)
func TestAssetUpsertWaitsForCompanyScopeMutation(t *testing.T) {
d, assets, companies := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
domain := fmt.Sprintf("scope-lock-%d.example", stamp)
oldCompany, _, err := companies.UpsertCompany(fmt.Sprintf("Scope Lock Old %d", stamp), "")
if err != nil {
t.Fatal(err)
}
newCompany, _, err := companies.UpsertCompany(fmt.Sprintf("Scope Lock New %d", stamp), "")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM assets WHERE domain = $1`, domain)
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, oldCompany, newCompany)
})
if added, _, invalid, validationErrors, err := companies.AddScopeInputsChecked(oldCompany, []ScopeInput{
{Kind: "domain", Value: domain},
}, "test"); err != nil || added != 1 || invalid != 0 || len(validationErrors) != 0 {
t.Fatalf("seed old scope: added=%d invalid=%d validation=%v err=%v", added, invalid, validationErrors, err)
}
mutation, err := d.Begin()
if err != nil {
t.Fatal(err)
}
defer mutation.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(mutation); err != nil {
t.Fatal(err)
}
if _, err := mutation.Exec(`DELETE FROM company_scope WHERE company_id=$1 AND kind='domain' AND domain=$2`, oldCompany, domain); err != nil {
t.Fatal(err)
}
if _, err := insertScopeRuleTx(mutation, newCompany, ParsedScope{
Kind: "domain", Domain: domain, Raw: domain,
}, "test"); err != nil {
t.Fatal(err)
}
started := make(chan struct{})
type result struct {
id int64
err error
}
resultCh := make(chan result, 1)
go func() {
close(started)
id, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: domain})
resultCh <- result{id: id, err: err}
}()
<-started
select {
case got := <-resultCh:
t.Fatalf("asset upsert escaped scope mutation lock: %+v", got)
case <-time.After(100 * time.Millisecond):
}
if err := mutation.Commit(); err != nil {
t.Fatal(err)
}
var got result
select {
case got = <-resultCh:
case <-time.After(3 * time.Second):
t.Fatal("asset upsert did not resume after scope mutation committed")
}
if got.err != nil {
t.Fatal(got.err)
}
var companyID *int64
if err := d.QueryRow(`SELECT company_id FROM assets WHERE id=$1`, got.id).Scan(&companyID); err != nil {
t.Fatal(err)
}
if companyID == nil || *companyID != newCompany {
t.Fatalf("asset retained stale company: got=%v want=%d", companyID, newCompany)
}
}
func TestResolveAndRecomputeUseStableCompanyIDTieBreak(t *testing.T) {
d, assets, companies := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
first, _, err := companies.UpsertCompany(fmt.Sprintf("Tie Break First %d", stamp), "")
if err != nil {
t.Fatal(err)
}
second, _, err := companies.UpsertCompany(fmt.Sprintf("Tie Break Second %d", stamp), "")
if err != nil {
t.Fatal(err)
}
want := min(first, second)
domain := fmt.Sprintf("tie-break-%d.example", stamp)
segment := uint64(stamp) & 0xffff
network := fmt.Sprintf("2001:db8:%x:1::/64", segment)
ip := fmt.Sprintf("2001:db8:%x:1::42", segment)
icp := fmt.Sprintf("ICP-TIE-%d", stamp)
appName := fmt.Sprintf("Tie Break App %d", stamp)
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM assets WHERE domain=$1 OR ip=$2 OR app_name=$3`, domain, ip, appName)
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, first, second)
})
rules := []ScopeInput{
{Kind: "domain", Value: domain},
{Kind: "cidr", Value: network},
{Kind: "icp", Value: icp},
}
for _, companyID := range []int64{second, first} {
if added, _, invalid, validationErrors, err := companies.AddScopeInputsChecked(companyID, rules, "test"); err != nil || added != len(rules) || invalid != 0 || len(validationErrors) != 0 {
t.Fatalf("add company %d rules: added=%d invalid=%d validation=%v err=%v", companyID, added, invalid, validationErrors, err)
}
}
assertResolved := func(label string, got *int64, err error) {
t.Helper()
if err != nil {
t.Fatalf("%s resolve: %v", label, err)
}
if got == nil || *got != want {
t.Fatalf("%s resolve=%v want lowest company id %d", label, got, want)
}
}
got, err := companies.ResolveCompany(domain, "")
assertResolved("domain", got, err)
got, err = companies.ResolveCompany("", ip)
assertResolved("network", got, err)
got, err = companies.ResolveCompanyWithICP("", "", icp)
assertResolved("icp", got, err)
rootID, err := assets.UpsertRootDomain(UpsertRootDomainReq{Domain: domain})
if err != nil {
t.Fatal(err)
}
ipID, err := assets.UpsertIP(UpsertIPReq{IP: ip})
if err != nil {
t.Fatal(err)
}
appID, err := assets.UpsertApp(UpsertAppReq{Name: appName, ICP: icp})
if err != nil {
t.Fatal(err)
}
assertAssets := func(stage string) {
t.Helper()
for _, id := range []int64{rootID, ipID, appID} {
var got *int64
if err := d.QueryRow(`SELECT company_id FROM assets WHERE id=$1`, id).Scan(&got); err != nil {
t.Fatal(err)
}
if got == nil || *got != want {
t.Fatalf("%s asset %d company=%v want=%d", stage, id, got, want)
}
}
}
assertAssets("live")
if err := companies.RecomputeAttribution(); err != nil {
t.Fatal(err)
}
assertAssets("recomputed")
}
func TestCompanyScopeLimitsAndCheckedErrors(t *testing.T) {
boundary := strings.Repeat("界", MaxCompanyScopeRawRunes)
if err := ValidateCompanyScopeInputBounds([]ScopeInput{{Kind: "keyword", Value: boundary}}); err != nil {
t.Fatalf("exact raw rune boundary rejected: %v", err)
}
var validationErr *CompanyScopeValidationError
if err := ValidateCompanyScopeInputBounds([]ScopeInput{{Kind: "keyword", Value: boundary + "界"}}); !errors.As(err, &validationErr) {
t.Fatalf("oversized raw value error=%v want CompanyScopeValidationError", err)
}
tooMany := make([]ScopeInput, MaxCompanyScopeRules+1)
if err := ValidateCompanyScopeInputBounds(tooMany); !errors.As(err, &validationErr) {
t.Fatalf("too many rules error=%v want CompanyScopeValidationError", err)
}
d, _, companies := testSetup(t)
defer d.Close()
stamp := time.Now().UnixNano()
companyID, _, err := companies.UpsertCompany(fmt.Sprintf("Scope Limits %d", stamp), "")
if err != nil {
t.Fatal(err)
}
defer cleanupCompany(d, companyID)
rules := make([]ScopeInput, MaxCompanyScopeRules)
for i := range rules {
rules[i] = ScopeInput{Kind: "keyword", Value: fmt.Sprintf("limit-%d-%d", stamp, i)}
}
if added, _, _, _, err := companies.AddScopeInputsChecked(companyID, rules, "test"); err != nil || added != MaxCompanyScopeRules {
t.Fatalf("fill scope capacity: added=%d err=%v", added, err)
}
if added, _, _, _, err := companies.AddScopeInputsChecked(companyID, []ScopeInput{
{Kind: "keyword", Value: fmt.Sprintf("limit-%d-overflow", stamp)},
}, "test"); added != 0 || !errors.As(err, &validationErr) {
t.Fatalf("scope overflow: added=%d err=%v", added, err)
}
var count int
if err := d.QueryRow(`SELECT count(*) FROM company_scope WHERE company_id=$1`, companyID).Scan(&count); err != nil {
t.Fatal(err)
}
if count != MaxCompanyScopeRules {
t.Fatalf("failed capacity mutation changed scope count=%d want=%d", count, MaxCompanyScopeRules)
}
missingID := companyID + 1_000_000_000
if _, _, _, _, err := companies.AddScopeInputsChecked(missingID, nil, "test"); !errors.Is(err, ErrCompanyNotFound) {
t.Fatalf("missing add error=%v want ErrCompanyNotFound", err)
}
if _, _, _, err := companies.UpdateScopeInputsChecked(missingID, nil, "test"); !errors.Is(err, ErrCompanyNotFound) {
t.Fatalf("missing update error=%v want ErrCompanyNotFound", err)
}
if _, err := companies.DeleteCompanyWithAssets(missingID, true); !errors.Is(err, ErrCompanyNotFound) {
t.Fatalf("missing delete error=%v want ErrCompanyNotFound", err)
}
}
func TestCompanyScopeCheckedReturnsSystemErrorSeparately(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
companies := d.Companies()
if err := d.Close(); err != nil {
t.Fatal(err)
}
_, _, _, validationErrors, err := companies.AddScopeInputsChecked(1, nil, "test")
if err == nil {
t.Fatal("closed database did not return a system error")
}
if len(validationErrors) != 0 {
t.Fatalf("system error leaked into validation errors: %v", validationErrors)
}
var validationErr *CompanyScopeValidationError
if errors.As(err, &validationErr) || errors.Is(err, ErrCompanyNotFound) {
t.Fatalf("system error misclassified: %v", err)
}
}
db/company_scope_test.go
+3 -12
View File
@@ -3,6 +3,7 @@ package db
import (
"fmt"
"testing"
"time"
)
// TestParseScopeLine covers classification + guardrails without a DB.
@@ -15,9 +16,9 @@ func TestParseScopeLine(t *testing.T) {
{"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)
{"10.0.0.0/8", "", true}, // over-broad IPv4 (< /16)
{"198.51.100.0/24", "cidr", false},
{"co.uk", "", true}, // bare public suffix
{"co.uk", "", true}, // bare public suffix
{"not a host", "", true},
{"1.2.3.1-1.2.3.9", "", true}, // ranges must be CIDR
}
@@ -39,8 +40,6 @@ func TestParseScopeLine(t *testing.T) {
}
}
<<<<<<< Updated upstream
=======
func TestParseAutoScopeLine(t *testing.T) {
cases := []struct {
name string
@@ -297,7 +296,6 @@ func TestCompanyICPAttribution(t *testing.T) {
assertCompany(icpAppID, &companyID)
}
>>>>>>> Stashed changes
// 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
@@ -343,13 +341,6 @@ func TestCompanyScopeAttribution(t *testing.T) {
t.Fatalf("pre-scope asset should be unattributed, got %v", *preCompanyID)
}
rules, invalid := ParseScopeLines(root + "\n198.51.100.0/24\nco.uk")
if len(rules) != 2 {
t.Fatalf("want 2 valid rules, got %d (%+v)", len(rules), rules)
}
if len(invalid) != 1 {
t.Fatalf("want 1 invalid line, got %v", invalid)
}
cs.AddScope(cid, []string{root, "198.51.100.0/24"}, "unit test")
mustCid := func(id int64, want int64, label string) {
db/task_assets.go
+421
View File
@@ -0,0 +1,421 @@
package db
import (
"database/sql"
"errors"
"fmt"
"net"
"strings"
"unicode/utf8"
)
const (
MaxTaskAssetMutationCount = 100
MaxTaskAssetSummaryRunes = 500
defaultTaskAssetSource = "system"
manualTaskScopeSummary = "用户在测试资产页手工新增"
)
var (
ErrTaskAssetInvalid = errors.New("invalid task asset association")
ErrTaskAssetTaskNotFound = errors.New("task not found")
ErrTaskAssetAssetNotFound = errors.New("asset not found")
)
// TaskAssetMutation summarizes one attach request. Attached counts newly added
// associations; Existing counts requested assets that were already on the task.
type TaskAssetMutation struct {
Requested int `json:"requested"`
Attached int `json:"attached"`
Existing int `json:"existing"`
}
// TaskAssetScopeMutation summarizes one free-form scope registration. Domain
// and IP entries create or reuse global assets; every entry also becomes an
// idempotent task_scope row.
type TaskAssetScopeMutation struct {
Requested int `json:"requested"`
AssetsLinked int `json:"assets_linked"`
AssetsExisting int `json:"assets_existing"`
ScopesAdded int `json:"scopes_added"`
ScopesExisting int `json:"scopes_existing"`
}
// IntentAsset describes an asset explicitly anchored to a worker intent.
type IntentAsset struct {
IntentID int64 `json:"intent_id"`
AssetID int64 `json:"asset_id"`
Type string `json:"type"`
Label string `json:"label"`
Source string `json:"source"`
SourceSummary string `json:"source_summary"`
SourceNodeID *int64 `json:"source_node_id,omitempty"`
SourceTaskID int64 `json:"source_task_id"`
Inherited bool `json:"inherited"`
}
func normalizeTaskAssetIDs(ids []int64) ([]int64, error) {
if len(ids) == 0 {
return nil, fmt.Errorf("%w: asset_ids is required", ErrTaskAssetInvalid)
}
seen := make(map[int64]struct{}, len(ids))
normalized := make([]int64, 0, len(ids))
for _, id := range ids {
if id <= 0 {
return nil, fmt.Errorf("%w: asset id must be positive", ErrTaskAssetInvalid)
}
if _, ok := seen[id]; ok {
continue
}
seen[id] = struct{}{}
normalized = append(normalized, id)
if len(normalized) > MaxTaskAssetMutationCount {
return nil, fmt.Errorf("%w: at most %d assets per request", ErrTaskAssetInvalid, MaxTaskAssetMutationCount)
}
}
return normalized, nil
}
func normalizeTaskAssetSource(source, summary string) (string, string, error) {
source = strings.TrimSpace(strings.ToLower(source))
if source == "" {
source = defaultTaskAssetSource
}
summary = strings.TrimSpace(summary)
if utf8.RuneCountInString(summary) > MaxTaskAssetSummaryRunes {
return "", "", fmt.Errorf("%w: source summary exceeds %d characters", ErrTaskAssetInvalid, MaxTaskAssetSummaryRunes)
}
return source, summary, nil
}
// SetTaskAssetSource improves the generic trigger-created provenance for one
// existing task association. It never creates or deletes an asset.
func (s *AssetStore) SetTaskAssetSource(taskID, assetID int64, source, summary string, sourceNodeID *int64) error {
if taskID <= 0 || assetID <= 0 {
return fmt.Errorf("%w: task and asset ids must be positive", ErrTaskAssetInvalid)
}
source, summary, err := normalizeTaskAssetSource(source, summary)
if err != nil {
return err
}
query := `
INSERT INTO task_asset_links(task_id, asset_id, source, source_summary, source_node_id)
SELECT task.id, asset.id, $3, $4, $5
FROM tasks task
JOIN assets asset ON asset.id=$2 AND task.id=ANY(asset.task_ids)
WHERE task.id=$1 AND task.deleted_at IS NULL
ON CONFLICT (task_id, asset_id) DO UPDATE
SET source=EXCLUDED.source,
source_summary=EXCLUDED.source_summary,
source_node_id=COALESCE(EXCLUDED.source_node_id, task_asset_links.source_node_id)`
var result sql.Result
if s.tx != nil {
result, err = s.tx.Exec(query, taskID, assetID, source, summary, sourceNodeID)
} else {
result, err = s.db.Exec(query, taskID, assetID, source, summary, sourceNodeID)
}
if err != nil {
return err
}
rows, err := result.RowsAffected()
if err != nil {
return err
}
if rows == 0 {
return fmt.Errorf("%w: task or asset association does not exist", ErrTaskAssetInvalid)
}
return nil
}
// RegisterTaskAssetScopes accepts the same structured scope rules as enterprise
// assets. The entire request is atomic: invalid input or any storage failure
// leaves both global assets and task scope unchanged.
func (s *AssetStore) RegisterTaskAssetScopes(taskID int64, inputs []ScopeInput) (TaskAssetScopeMutation, error) {
mutation := TaskAssetScopeMutation{Requested: len(inputs)}
if taskID <= 0 {
return mutation, fmt.Errorf("%w: task id must be positive", ErrTaskAssetInvalid)
}
if len(inputs) == 0 {
return mutation, fmt.Errorf("%w: scope is required", ErrTaskAssetInvalid)
}
if err := ValidateCompanyScopeInputBounds(inputs); err != nil {
return mutation, fmt.Errorf("%w: %v", ErrTaskAssetInvalid, err)
}
parsed := make([]ParsedScope, 0, len(inputs))
for index, input := range inputs {
rule, err := ParseScopeInput(input)
if err != nil {
return mutation, fmt.Errorf("%w: 第 %d 条范围无效: %v", ErrTaskAssetInvalid, index+1, err)
}
parsed = append(parsed, rule)
}
if err := validateParsedScopeBounds(parsed); err != nil {
return mutation, fmt.Errorf("%w: %v", ErrTaskAssetInvalid, err)
}
tx, err := s.db.Begin()
if err != nil {
return mutation, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return mutation, err
}
var taskExists bool
if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM tasks WHERE id=$1 AND deleted_at IS NULL)`, taskID).Scan(&taskExists); err != nil {
return mutation, err
}
if !taskExists {
return mutation, ErrTaskAssetTaskNotFound
}
scoped := &AssetStore{db: s.db, company: s.company, tx: tx}
for _, rule := range parsed {
taskScope := TaskScope{
TaskID: taskID,
Source: "manual",
Reason: manualTaskScopeSummary,
}
var assetID int64
switch rule.Kind {
case "domain":
taskScope.Kind = "root_domain"
taskScope.Domain = rule.Domain
var alreadyLinked bool
err := tx.QueryRow(`SELECT id, $2=ANY(task_ids) FROM assets WHERE type='root_domain' AND domain=$1`, rule.Domain, taskID).
Scan(&assetID, &alreadyLinked)
if err != nil && !errors.Is(err, sql.ErrNoRows) {
return mutation, err
}
assetID, err = scoped.UpsertRootDomain(UpsertRootDomainReq{Domain: rule.Domain, TaskID: taskID})
if err != nil {
return mutation, err
}
if alreadyLinked {
mutation.AssetsExisting++
} else {
mutation.AssetsLinked++
}
case "ip":
taskScope.Kind = "ip"
taskScope.Net = rule.Net
ip, _, parseErr := net.ParseCIDR(rule.Net)
if parseErr != nil {
return mutation, fmt.Errorf("%w: 无效 IP: %s", ErrTaskAssetInvalid, rule.Raw)
}
ipValue := ip.String()
var alreadyLinked bool
err := tx.QueryRow(`SELECT id, $2=ANY(task_ids) FROM assets WHERE type='ip' AND ip=$1`, ipValue, taskID).
Scan(&assetID, &alreadyLinked)
if err != nil && !errors.Is(err, sql.ErrNoRows) {
return mutation, err
}
assetID, err = scoped.UpsertIP(UpsertIPReq{IP: ipValue, TaskID: taskID})
if err != nil {
return mutation, err
}
if alreadyLinked {
mutation.AssetsExisting++
} else {
mutation.AssetsLinked++
}
case "cidr":
taskScope.Kind = "cidr"
taskScope.Net = rule.Net
case "icp", "keyword":
taskScope.Kind = rule.Kind
taskScope.Value = rule.Value
default:
return mutation, fmt.Errorf("%w: unsupported scope kind %q", ErrTaskAssetInvalid, rule.Kind)
}
if assetID > 0 {
if err := scoped.SetTaskAssetSource(taskID, assetID, "manual", manualTaskScopeSummary, nil); err != nil {
return mutation, err
}
}
inserted, err := scoped.upsertTaskScopeResult(taskScope)
if err != nil {
return mutation, err
}
if inserted {
mutation.ScopesAdded++
} else {
mutation.ScopesExisting++
}
}
if err := tx.Commit(); err != nil {
return mutation, err
}
return mutation, nil
}
// AttachAssetsToTask associates existing global assets with one live task and
// records an operator-authored source summary. Global asset rows are retained.
func (s *AssetStore) AttachAssetsToTask(taskID int64, assetIDs []int64, sourceSummary string) (TaskAssetMutation, error) {
var mutation TaskAssetMutation
assetIDs, err := normalizeTaskAssetIDs(assetIDs)
if err != nil {
return mutation, err
}
_, sourceSummary, err = normalizeTaskAssetSource("manual", sourceSummary)
if err != nil {
return mutation, err
}
if sourceSummary == "" {
return mutation, fmt.Errorf("%w: source_summary is required", ErrTaskAssetInvalid)
}
mutation.Requested = len(assetIDs)
tx, err := s.db.Begin()
if err != nil {
return mutation, err
}
defer tx.Rollback() //nolint:errcheck
var taskExists bool
if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM tasks WHERE id=$1 AND deleted_at IS NULL)`, taskID).Scan(&taskExists); err != nil {
return mutation, err
}
if !taskExists {
return mutation, ErrTaskAssetTaskNotFound
}
var found, existing int
if err := tx.QueryRow(`
SELECT count(*), count(*) FILTER (WHERE $1=ANY(task_ids))
FROM assets WHERE id=ANY($2::bigint[])`, taskID, assetIDs).Scan(&found, &existing); err != nil {
return mutation, err
}
if found != len(assetIDs) {
return mutation, ErrTaskAssetAssetNotFound
}
if _, err := tx.Exec(`
UPDATE assets
SET task_ids=CASE WHEN $1=ANY(task_ids) THEN task_ids ELSE array_append(task_ids,$1) END
WHERE id=ANY($2::bigint[])`, taskID, assetIDs); err != nil {
return mutation, err
}
if _, err := tx.Exec(`
INSERT INTO task_asset_links(task_id, asset_id, source, source_summary)
SELECT $1, id, 'manual', $3 FROM assets WHERE id=ANY($2::bigint[])
ON CONFLICT (task_id, asset_id) DO UPDATE
SET source='manual', source_summary=EXCLUDED.source_summary, source_node_id=NULL`,
taskID, assetIDs, sourceSummary); err != nil {
return mutation, err
}
mutation.Existing = existing
mutation.Attached = len(assetIDs) - existing
return mutation, tx.Commit()
}
// DetachAssetFromTask removes only the task association. The global asset and
// exploration anchors remain available for historical blackboard auditing.
func (s *AssetStore) DetachAssetFromTask(taskID, assetID int64) (bool, error) {
if taskID <= 0 || assetID <= 0 {
return false, fmt.Errorf("%w: task and asset ids must be positive", ErrTaskAssetInvalid)
}
var detachedID int64
err := s.db.QueryRow(`
UPDATE assets SET task_ids=array_remove(task_ids,$1)
WHERE id=$2 AND $1=ANY(task_ids)
RETURNING id`, taskID, assetID).Scan(&detachedID)
if err == sql.ErrNoRows {
return false, nil
}
if err != nil {
return false, err
}
return detachedID == assetID, nil
}
func (s *AssetStore) hydrateTaskAssetSources(taskID int64, assets []*Asset) error {
if len(assets) == 0 {
return nil
}
ids := make([]int64, 0, len(assets))
byID := make(map[int64]*Asset, len(assets))
for _, asset := range assets {
ids = append(ids, asset.ID)
byID[asset.ID] = asset
}
rows, err := s.db.Query(`
SELECT asset_id, source, source_summary, source_node_id
FROM task_asset_links
WHERE task_id=$1 AND asset_id=ANY($2::bigint[])`, taskID, ids)
if err != nil {
return err
}
defer rows.Close()
for rows.Next() {
var assetID int64
var source, summary string
var sourceNodeID sql.NullInt64
if err := rows.Scan(&assetID, &source, &summary, &sourceNodeID); err != nil {
return err
}
if asset := byID[assetID]; asset != nil {
asset.TaskSource = source
asset.TaskSourceSummary = summary
if sourceNodeID.Valid {
id := sourceNodeID.Int64
asset.TaskSourceNodeID = &id
}
}
}
return rows.Err()
}
// IntentAssets returns all local worker targets plus immutable targets from the
// task's direct sources. Inherited non-terminal intents remain hidden, matching
// the existing source-aware session contract.
func (s *AssetStore) IntentAssets(taskID int64) ([]IntentAsset, error) {
rows, err := s.db.Query(`
WITH context AS (
SELECT task.id AS task_id, task.exploration_id, false AS inherited
FROM tasks task
WHERE task.id=$1 AND task.deleted_at IS NULL
UNION ALL
SELECT source.id, source.exploration_id, true
FROM task_relations relation
JOIN tasks source ON source.id=relation.source_task_id AND source.deleted_at IS NULL
WHERE relation.task_id=$1
)
SELECT intent.id, asset.id, asset.type,
CASE asset.type
WHEN 'root_domain' THEN COALESCE(asset.domain,'')
WHEN 'subdomain' THEN COALESCE(asset.domain,'')
WHEN 'ip' THEN COALESCE(asset.ip,'')
WHEN 'app' THEN COALESCE(asset.app_name,'')
WHEN 'service' THEN COALESCE(NULLIF(asset.url,''), NULLIF(concat_ws(':', COALESCE(NULLIF(asset.domain,''), NULLIF(asset.ip,'')), asset.port::text),''), NULLIF(asset.service_name,''), '#' || asset.id::text)
WHEN 'endpoint' THEN COALESCE(NULLIF(asset.url,''), '#' || asset.id::text)
ELSE '#' || asset.id::text
END,
COALESCE(link.source,'anchor'),
COALESCE(NULLIF(link.source_summary,''), '意图在黑板中锚定该资产'),
link.source_node_id, context.task_id, context.inherited
FROM context
JOIN exploration_nodes intent ON intent.exploration_id=context.exploration_id AND intent.kind='intent'
JOIN exploration_anchors anchor ON anchor.node_id=intent.id
JOIN assets asset ON asset.id=anchor.asset_id
LEFT JOIN task_asset_links link ON link.task_id=context.task_id AND link.asset_id=asset.id
WHERE NOT context.inherited OR intent.state IN ('done','blocked','exhausted','stopped')
ORDER BY context.inherited, intent.id DESC, asset.id`, taskID)
if err != nil {
return nil, err
}
defer rows.Close()
out := []IntentAsset{}
for rows.Next() {
var asset IntentAsset
var sourceNodeID sql.NullInt64
if err := rows.Scan(&asset.IntentID, &asset.AssetID, &asset.Type, &asset.Label,
&asset.Source, &asset.SourceSummary, &sourceNodeID, &asset.SourceTaskID, &asset.Inherited); err != nil {
return nil, err
}
if sourceNodeID.Valid {
id := sourceNodeID.Int64
asset.SourceNodeID = &id
}
out = append(out, asset)
}
return out, rows.Err()
}
db/task_assets_test.go
+188
View File
@@ -0,0 +1,188 @@
package db
import (
"errors"
"fmt"
"testing"
"time"
)
func TestRegisterTaskAssetScopesCreatesAssetsAndPersistsTextScope(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
task, err := d.CreateTask("manual scope registration", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
suffix := time.Now().UnixNano()
domain := fmt.Sprintf("manual-scope-%d.example.test", suffix)
ip := fmt.Sprintf("203.0.%d.%d", (suffix/250)%250, suffix%250+1)
inputs := []ScopeInput{
{Kind: "domain", Value: domain},
{Kind: "ip", Value: ip},
{Kind: "cidr", Value: "198.51.100.0/24"},
{Kind: "icp", Value: " 京 ICP 备 12345678 号-1 "},
{Kind: "keyword", Value: " Acme Security "},
}
first, err := d.Assets().RegisterTaskAssetScopes(task.ID, inputs)
if err != nil {
t.Fatal(err)
}
if first.Requested != 5 || first.AssetsLinked != 2 || first.AssetsExisting != 0 || first.ScopesAdded != 5 || first.ScopesExisting != 0 {
t.Fatalf("unexpected first mutation: %+v", first)
}
assets, err := d.Assets().QueryByTask(task.ID, "", 10, 0)
if err != nil || len(assets) != 2 {
t.Fatalf("task assets=%+v err=%v", assets, err)
}
assetIDs := make([]int64, 0, len(assets))
for _, asset := range assets {
assetIDs = append(assetIDs, asset.ID)
if asset.TaskSource != "manual" || asset.TaskSourceSummary != manualTaskScopeSummary {
t.Fatalf("unexpected task asset provenance: %+v", asset)
}
}
t.Cleanup(func() { _, _ = d.Assets().DeleteByIDs(assetIDs) })
scopes, err := d.Assets().ListTaskScope(task.ID)
if err != nil || len(scopes) != 5 {
t.Fatalf("task scope=%+v err=%v", scopes, err)
}
values := map[string]string{}
for _, scope := range scopes {
values[scope.Kind] = scope.Value
}
if values["icp"] != NormalizeICP(inputs[3].Value) || values["keyword"] != normalizeKeyword(inputs[4].Value) {
t.Fatalf("text scope not normalized: %+v", values)
}
second, err := d.Assets().RegisterTaskAssetScopes(task.ID, inputs)
if err != nil {
t.Fatal(err)
}
if second.AssetsExisting != 2 || second.AssetsLinked != 0 || second.ScopesExisting != 5 || second.ScopesAdded != 0 {
t.Fatalf("unexpected idempotent mutation: %+v", second)
}
rollbackDomain := fmt.Sprintf("rollback-%d.example.test", suffix)
_, err = d.Assets().RegisterTaskAssetScopes(task.ID, []ScopeInput{
{Kind: "domain", Value: rollbackDomain},
{Kind: "cidr", Value: "10.0.0.0/8"},
})
if !errors.Is(err, ErrTaskAssetInvalid) {
t.Fatalf("invalid CIDR error=%v", err)
}
var rollbackAssets int
if err := d.QueryRow(`SELECT count(*) FROM assets WHERE type='root_domain' AND domain=$1`, rollbackDomain).Scan(&rollbackAssets); err != nil {
t.Fatal(err)
}
if rollbackAssets != 0 {
t.Fatalf("invalid request created %d assets", rollbackAssets)
}
}
func TestTaskAssetAttachDetachPreservesGlobalAssetAndAnchors(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
task, err := d.CreateTask("asset editing", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
assetID, err := d.Assets().UpsertRootDomain(UpsertRootDomainReq{Domain: fmt.Sprintf("manual-%d.example.test", time.Now().UnixNano())})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _, _ = d.Assets().DeleteByIDs([]int64{assetID}) })
mutation, err := d.Assets().AttachAssetsToTask(task.ID, []int64{assetID, assetID}, "授权资产清单第 3 项")
if err != nil {
t.Fatal(err)
}
if mutation.Requested != 1 || mutation.Attached != 1 || mutation.Existing != 0 {
t.Fatalf("unexpected first mutation: %+v", mutation)
}
assets, err := d.Assets().QueryByTask(task.ID, "", 10, 0)
if err != nil || len(assets) != 1 {
t.Fatalf("task assets=%+v err=%v", assets, err)
}
if assets[0].TaskSource != "manual" || assets[0].TaskSourceSummary != "授权资产清单第 3 项" {
t.Fatalf("unexpected provenance: %+v", assets[0])
}
intentID, err := d.Exploration(task.ExplorationID).AddIntent(map[string]any{"summary": "test asset"}, 5, []int64{assetID}, "human")
if err != nil {
t.Fatal(err)
}
detached, err := d.Assets().DetachAssetFromTask(task.ID, assetID)
if err != nil || !detached {
t.Fatalf("detach=%v err=%v", detached, err)
}
if assets, err := d.Assets().QueryByTask(task.ID, "", 10, 0); err != nil || len(assets) != 0 {
t.Fatalf("detached task assets=%+v err=%v", assets, err)
}
var global, anchors, links int
_ = d.QueryRow(`SELECT count(*) FROM assets WHERE id=$1`, assetID).Scan(&global)
_ = d.QueryRow(`SELECT count(*) FROM exploration_anchors WHERE node_id=$1 AND asset_id=$2`, intentID, assetID).Scan(&anchors)
_ = d.QueryRow(`SELECT count(*) FROM task_asset_links WHERE task_id=$1 AND asset_id=$2`, task.ID, assetID).Scan(&links)
if global != 1 || anchors != 1 || links != 0 {
t.Fatalf("global=%d anchors=%d links=%d", global, anchors, links)
}
}
func TestIntentAssetsIncludesDirectSourceProvenance(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer d.Close()
source, err := d.CreateTask("source assets", "goal", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
current, err := d.CreateTaskWithOptions("current assets", "goal", TaskCreateOptions{SourceTaskIDs: []int64{source.ID}})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteTask(current.ID); _ = d.DeleteTask(source.ID) })
assetID, err := d.Assets().UpsertHTTPService(UpsertHTTPServiceReq{
URL: fmt.Sprintf("https://intent-%d.example.test", time.Now().UnixNano()), TaskID: source.ID,
})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _, _ = d.Assets().DeleteByIDs([]int64{assetID}) })
intentID, err := d.Exploration(source.ExplorationID).AddIntent(map[string]any{"summary": "source worker"}, 5, []int64{assetID}, "planner")
if err != nil {
t.Fatal(err)
}
if err := d.Exploration(source.ExplorationID).SetNodeState(intentID, "done"); err != nil {
t.Fatal(err)
}
nodeID := intentID
if err := d.Assets().SetTaskAssetSource(source.ID, assetID, "agent", "Worker 通过 insert_assets 登记", &nodeID); err != nil {
t.Fatal(err)
}
assets, err := d.Assets().IntentAssets(current.ID)
if err != nil {
t.Fatal(err)
}
if len(assets) != 1 || assets[0].IntentID != intentID || assets[0].SourceTaskID != source.ID || !assets[0].Inherited {
t.Fatalf("unexpected intent assets: %+v", assets)
}
if assets[0].Source != "agent" || assets[0].SourceSummary == "" || assets[0].SourceNodeID == nil || *assets[0].SourceNodeID != intentID {
t.Fatalf("unexpected intent provenance: %+v", assets[0])
}
}
db/task_scope.go
+73 -25
View File
@@ -1,6 +1,7 @@
package db
import (
"database/sql"
"fmt"
"net"
"strconv"
@@ -10,17 +11,18 @@ import (
// TaskScope is one row of a task's test scope — the coverage denominator and the
// per-task authorization edge. Rows come either from insertAssets (source='auto',
// conservative, one per explicitly-inserted asset) or from the add_task_scope tool
// (source='agent', for company / whole-root-domain / subdomain / ip).
// (source='agent', for company / domain / network / ICP / keyword scope).
type TaskScope struct {
ID int64 `json:"id"`
TaskID int64 `json:"task_id"`
Kind string `json:"kind"` // company|root_domain|subdomain|ip|cidr
Kind string `json:"kind"` // company|root_domain|subdomain|ip|cidr|icp|keyword
CompanyID *int64 `json:"company_id,omitempty"`
// CompanyName is resolved for kind=company so callers can label a scope row
// without a second lookup. Empty when the row is not a company reference.
CompanyName string `json:"company_name,omitempty"`
Domain string `json:"domain,omitempty"`
Net string `json:"net,omitempty"`
Value string `json:"value,omitempty"`
Source string `json:"source"`
Reason string `json:"reason,omitempty"`
}
@@ -53,11 +55,11 @@ func ipToHostCIDR(ip string) string {
// upsertTaskScope inserts one scope row idempotently (uq_task_scope). A duplicate is
// silently ignored. taskID<=0 or empty kind → no-op.
func (s *AssetStore) upsertTaskScope(ts TaskScope) error {
func (s *AssetStore) upsertTaskScopeResult(ts TaskScope) (bool, error) {
if ts.TaskID <= 0 || ts.Kind == "" {
return nil
return false, nil
}
var domainVal, netVal, companyVal any
var domainVal, netVal, companyVal, valueVal any
if ts.Domain != "" {
domainVal = ts.Domain
}
@@ -67,14 +69,35 @@ func (s *AssetStore) upsertTaskScope(ts TaskScope) error {
if ts.CompanyID != nil && *ts.CompanyID > 0 {
companyVal = *ts.CompanyID
}
if ts.Value != "" {
valueVal = ts.Value
}
src := ts.Source
if src == "" {
src = "auto"
}
_, err := s.db.Exec(`
INSERT INTO task_scope(task_id, kind, company_id, domain, net, source, reason)
VALUES ($1,$2,$3,$4,$5::cidr,$6,NULLIF($7,''))
ON CONFLICT DO NOTHING`, ts.TaskID, ts.Kind, companyVal, domainVal, netVal, src, ts.Reason)
query := `
INSERT INTO task_scope(task_id, kind, company_id, domain, net, value, source, reason)
VALUES ($1,$2,$3,$4,$5::cidr,$6,$7,NULLIF($8,''))
ON CONFLICT DO NOTHING`
var (
result sql.Result
err error
)
if s.tx != nil {
result, err = s.tx.Exec(query, ts.TaskID, ts.Kind, companyVal, domainVal, netVal, valueVal, src, ts.Reason)
} else {
result, err = s.db.Exec(query, ts.TaskID, ts.Kind, companyVal, domainVal, netVal, valueVal, src, ts.Reason)
}
if err != nil {
return false, err
}
rows, err := result.RowsAffected()
return rows > 0, err
}
func (s *AssetStore) upsertTaskScope(ts TaskScope) error {
_, err := s.upsertTaskScopeResult(ts)
return err
}
@@ -116,11 +139,15 @@ func (s *AssetStore) AddAutoScope(taskID int64, assetType, domain, rawURL, ip st
return nil
}
// AddAgentScope parses an agent-provided (kind,value) and records it (source='agent').
// AddAgentScope parses a (kind,value) pair and records it as task scope.
// source is "agent" when called from an LLM tool, "manual" from the UI.
// company: value = company name or id (must already exist). root_domain/subdomain:
// value = a domain. ip/cidr: value = an IP or CIDR (bare IP → /32,/128).
func (s *AssetStore) AddAgentScope(taskID int64, kind, value, reason string) (TaskScope, error) {
ts := TaskScope{TaskID: taskID, Kind: kind, Source: "agent", Reason: reason}
func (s *AssetStore) AddAgentScope(taskID int64, kind, value, reason, source string) (TaskScope, error) {
if source == "" {
source = "agent"
}
ts := TaskScope{TaskID: taskID, Kind: kind, Source: source, Reason: reason}
if taskID <= 0 {
return ts, fmt.Errorf("需要 task_id")
}
@@ -178,8 +205,14 @@ func (s *AssetStore) AddAgentScope(taskID int64, kind, value, reason string) (Ta
return ts, fmt.Errorf("无效 ip/cidr: %s", value)
}
ts.Net = v
case "icp", "keyword":
parsed, err := ParseScopeInput(ScopeInput{Kind: kind, Value: value})
if err != nil {
return ts, err
}
ts.Value = parsed.Value
default:
return ts, fmt.Errorf("不支持的 kind: %s（company/root_domain/subdomain/ip/cidr）", kind)
return ts, fmt.Errorf("不支持的 kind: %s（company/root_domain/subdomain/ip/cidr/icp/keyword）", kind)
}
if err := s.upsertTaskScope(ts); err != nil {
return ts, err
@@ -187,11 +220,22 @@ func (s *AssetStore) AddAgentScope(taskID int64, kind, value, reason string) (Ta
return ts, nil
}
// DeleteTaskScope removes a single scope row by id, scoped to the given task.
// Returns whether a row was actually deleted.
func (s *AssetStore) DeleteTaskScope(taskID, scopeID int64) (bool, error) {
res, err := s.db.Exec(`DELETE FROM task_scope WHERE id=$1 AND task_id=$2`, scopeID, taskID)
if err != nil {
return false, err
}
n, _ := res.RowsAffected()
return n > 0, nil
}
// ListTaskScope returns all scope rows for a task.
func (s *AssetStore) ListTaskScope(taskID int64) ([]TaskScope, error) {
rows, err := s.db.Query(`
SELECT ts.id, ts.kind, COALESCE(ts.company_id,0), COALESCE(c.name,''), COALESCE(ts.domain,''),
COALESCE(ts.net::text,''), ts.source, COALESCE(ts.reason,'')
COALESCE(ts.net::text,''), COALESCE(ts.value,''), ts.source, COALESCE(ts.reason,'')
FROM task_scope ts
LEFT JOIN companies c ON c.id=ts.company_id
WHERE ts.task_id=$1 ORDER BY ts.id`, taskID)
@@ -203,7 +247,7 @@ WHERE ts.task_id=$1 ORDER BY ts.id`, taskID)
for rows.Next() {
var t TaskScope
var cid int64
if err := rows.Scan(&t.ID, &t.Kind, &cid, &t.CompanyName, &t.Domain, &t.Net, &t.Source, &t.Reason); err != nil {
if err := rows.Scan(&t.ID, &t.Kind, &cid, &t.CompanyName, &t.Domain, &t.Net, &t.Value, &t.Source, &t.Reason); err != nil {
return nil, err
}
t.TaskID = taskID
@@ -264,9 +308,13 @@ target AS (
JOIN task_scope ts ON ts.task_id = $1 AND (
(ts.kind='company' AND a.company_id = ts.company_id)
OR (ts.kind='root_domain' AND a.root_domain = ts.domain)
OR (ts.kind='subdomain' AND a.domain = ts.domain)
OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet)
)
OR (ts.kind='subdomain' AND a.domain = ts.domain)
OR (ts.kind IN ('ip','cidr') AND a.ip IS NOT NULL AND a.ip !~ '[^0-9a-fA-F:.]' AND ts.net >>= a.ip::inet)
OR (ts.kind='icp' AND (
lower(regexp_replace(COALESCE(a.icp,''), '[[:space:]]+', '', 'g')) = ts.value
OR lower(regexp_replace(COALESCE(a.app_icp,''), '[[:space:]]+', '', 'g')) = ts.value
))
)
),
tested AS (
SELECT DISTINCT ea.asset_id
@@ -368,23 +416,23 @@ func hostPortOf(n *CoverageGraphNode) (string, int) {
return host, port
}
// BuildCoverageGraph assembles the full coverage graph for a task: every in-scope
// asset (all types) plus the connector root domains / companies needed to link
// them, with a tested flag on each in-scope asset. The frontend does the folding
// and highlighting; this only returns nodes + derived containment edges.
func (s *AssetStore) BuildCoverageGraph(taskID, expID int64) (*CoverageGraphData, error) {
// BuildCoverageGraph assembles the full coverage graph for a task and its direct
// read-only sources: every in-scope asset plus connector root domains/companies,
// with current-or-source fact anchors reflected in Tested. The legacy expID
// argument is retained for API compatibility; the task registry is authoritative.
func (s *AssetStore) BuildCoverageGraph(taskID, _ int64) (*CoverageGraphData, error) {
g := &CoverageGraphData{Nodes: []CoverageGraphNode{}, Edges: []CoverageGraphEdge{}}
if taskID <= 0 {
return g, nil
}
rows, err := s.db.Query(`WITH `+covTargetCTE+`
rows, err := s.db.Query(`WITH `+contextCoverageCTE+`
SELECT a.id, a.type, COALESCE(a.company_id,0),
COALESCE(a.domain,''), COALESCE(a.root_domain,''), COALESCE(a.ip,''),
COALESCE(a.url,''), COALESCE(a.port,0), COALESCE(a.service_type,''),
COALESCE(a.app_name,''), COALESCE(a.page_title,''), COALESCE(a.status_code,0),
(a.id IN (SELECT asset_id FROM tested)) AS tested
FROM assets a JOIN target t ON t.id = a.id
ORDER BY a.id`, taskID, expID)
ORDER BY a.id`, taskID)
if err != nil {
return nil, err
}