admin/ARTEX-chixuweihu
1
0 Fork 0
mirror of https://github.com/RuoJi6/ARTEX.git synced 2026-10-27 17:03:49 +08:00
Code Issues 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
}
server/assets.go
+167 -53
View File
@@ -2,12 +2,58 @@ package server
import (
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"strconv"
"strings"
"github.com/Autumn-27/artex/db"
)
// companyScopeInputs accepts both the new [{kind,value}] contract and the
// historical ["example.com","203.0.113.10"] contract.
type companyScopeInputs []db.ScopeInput
const maxCompanyMutationBodyBytes = 2 << 20
func decodeCompanyMutationRequest(w http.ResponseWriter, r *http.Request, value any) bool {
r.Body = http.MaxBytesReader(w, r.Body, maxCompanyMutationBodyBytes)
if err := json.NewDecoder(r.Body).Decode(value); err != nil {
var tooLarge *http.MaxBytesError
if errors.As(err, &tooLarge) {
writeErr(w, http.StatusRequestEntityTooLarge, "请求正文过大")
} else {
writeErr(w, http.StatusBadRequest, "invalid JSON: "+err.Error())
}
return false
}
return true
}
func (items *companyScopeInputs) UnmarshalJSON(data []byte) error {
var rawItems []json.RawMessage
if err := json.Unmarshal(data, &rawItems); err != nil {
return err
}
out := make([]db.ScopeInput, 0, len(rawItems))
for i, raw := range rawItems {
var legacy string
if err := json.Unmarshal(raw, &legacy); err == nil {
out = append(out, db.ScopeInput{Value: legacy})
continue
}
var structured db.ScopeInput
if err := json.Unmarshal(raw, &structured); err != nil {
return fmt.Errorf("scope[%d] must be a string or {kind,value}", i)
}
out = append(out, structured)
}
*items = out
return nil
}
// assetStore returns the asset store.
func (s *Server) assetStore() *db.AssetStore {
if s.m.pg == nil {
@@ -53,32 +99,45 @@ func (s *Server) createCompany(w http.ResponseWriter, r *http.Request) {
return
}
var req struct {
Name string `json:"name"`
Logo string `json:"logo"`
Scope []string `json:"scope"`
Name string `json:"name"`
Logo string `json:"logo"`
Scope companyScopeInputs `json:"scope"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
writeErr(w, 400, "invalid JSON: "+err.Error())
if !decodeCompanyMutationRequest(w, r, &req) {
return
}
req.Name = strings.TrimSpace(req.Name)
if req.Name == "" {
writeErr(w, 400, "name required")
return
}
id, created, err := cs.UpsertCompany(req.Name, req.Logo)
if err != nil {
writeErr(w, 500, err.Error())
if err := db.ValidateCompanyScopeInputBounds(req.Scope); err != nil {
writeErr(w, http.StatusBadRequest, err.Error())
return
}
out := map[string]any{"id": id, "created": created}
if len(req.Scope) > 0 {
added, skipped, invalid, errs := cs.AddScope(id, req.Scope, "api")
out["scope_added"] = added
out["scope_skipped"] = skipped
out["scope_invalid"] = invalid
if len(errs) > 0 {
out["scope_errors"] = errs
id, added, skipped, invalid, scopeErrs, err := cs.CreateCompanyWithScope(req.Name, req.Logo, req.Scope, "api")
if err != nil {
if errors.Is(err, db.ErrCompanyNameConflict) {
writeErr(w, http.StatusConflict, "企业名称已存在")
return
}
var validationErr *db.CompanyScopeValidationError
if errors.As(err, &validationErr) {
writeErr(w, http.StatusBadRequest, validationErr.Error())
return
}
writeErr(w, http.StatusInternalServerError, err.Error())
return
}
out := map[string]any{
"id": id,
"created": true,
"scope_added": added,
"scope_skipped": skipped,
"scope_invalid": invalid,
}
if len(scopeErrs) > 0 {
out["scope_errors"] = scopeErrs
}
writeJSON(w, 201, out)
}
@@ -131,20 +190,37 @@ func (s *Server) addCompanyScope(w http.ResponseWriter, r *http.Request) {
return
}
var req struct {
Scope []string `json:"scope"`
Reason string `json:"reason"`
Reset bool `json:"reset"` // if true, replace existing scope
Scope companyScopeInputs `json:"scope"`
Reason string `json:"reason"`
Reset bool `json:"reset"` // if true, replace existing scope
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
writeErr(w, 400, "invalid JSON: "+err.Error())
if !decodeCompanyMutationRequest(w, r, &req) {
return
}
if err := db.ValidateCompanyScopeInputBounds(req.Scope); err != nil {
writeErr(w, http.StatusBadRequest, err.Error())
return
}
var added, skipped, invalid int
var errs []string
var mutationErr error
if req.Reset {
added, invalid, errs = cs.UpdateScope(id, req.Scope, req.Reason)
added, invalid, errs, mutationErr = cs.UpdateScopeInputsChecked(id, req.Scope, req.Reason)
} else {
added, skipped, invalid, errs = cs.AddScope(id, req.Scope, req.Reason)
added, skipped, invalid, errs, mutationErr = cs.AddScopeInputsChecked(id, req.Scope, req.Reason)
}
if mutationErr != nil {
if errors.Is(mutationErr, db.ErrCompanyNotFound) {
writeErr(w, http.StatusNotFound, "company not found")
return
}
var validationErr *db.CompanyScopeValidationError
if errors.As(mutationErr, &validationErr) {
writeErr(w, http.StatusBadRequest, validationErr.Error())
return
}
writeErr(w, http.StatusInternalServerError, mutationErr.Error())
return
}
out := map[string]any{
"added": added,
@@ -175,21 +251,32 @@ func (s *Server) deleteCompany(w http.ResponseWriter, r *http.Request) {
var req struct {
DeleteAssets bool `json:"delete_assets"`
}
// Body is optional; ignore decode errors (e.g. empty body).
_ = json.NewDecoder(r.Body).Decode(&req)
assetsDeleted := int64(0)
if req.DeleteAssets {
as := s.assetStore()
if as != nil {
assetsDeleted, err = as.DeleteByCompanyID(id)
if err != nil {
writeErr(w, 500, err.Error())
return
// The body is optional. Only an actual empty body is ignored; malformed or
// trailing JSON is a client error.
decoder := json.NewDecoder(r.Body)
if err := decoder.Decode(&req); err != nil {
if !errors.Is(err, io.EOF) {
writeErr(w, http.StatusBadRequest, "invalid JSON: "+err.Error())
return
}
} else {
var trailing any
if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) {
if err == nil {
writeErr(w, http.StatusBadRequest, "invalid JSON: multiple values")
} else {
writeErr(w, http.StatusBadRequest, "invalid JSON: "+err.Error())
}
return
}
}
if err := cs.DeleteCompany(id); err != nil {
assetsDeleted, err := s.m.DeleteCompanyWithAssets(id, req.DeleteAssets)
if err != nil {
if errors.Is(err, db.ErrCompanyNotFound) {
writeErr(w, http.StatusNotFound, "company not found")
return
}
writeErr(w, 500, err.Error())
return
}
@@ -217,6 +304,11 @@ func (s *Server) reattribute(w http.ResponseWriter, r *http.Request) {
// GET /api/assets
// =====================================================================
const (
defaultAssetPageSize = 50
maxAssetPageSize = 200
)
func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) {
as := s.assetStore()
if as == nil {
@@ -228,43 +320,56 @@ func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) {
limit, _ := strconv.Atoi(q.Get("limit"))
offset, _ := strconv.Atoi(q.Get("offset"))
if limit <= 0 {
limit = 50
limit = defaultAssetPageSize
} else if limit > maxAssetPageSize {
limit = maxAssetPageSize
}
if offset < 0 {
offset = 0
}
var assets []*db.Asset
assets := []*db.Asset{}
var err error
// total is the full match count ignoring limit/offset — lets the UI paginate on
// the server. For company/task scopes it falls back to the returned length.
// total is the full match count ignoring limit/offset. Count first so an offset
// beyond the last row can return an empty page without an expensive scan.
total := 0
if dsl := q.Get("dsl"); dsl != "" {
assets, err = as.QueryDSL(dsl, typ, limit, offset)
if err == nil {
total, _ = as.CountDSL(dsl, typ)
if err := db.ValidateDSL(dsl); err != nil {
writeErr(w, http.StatusBadRequest, err.Error())
return
}
total, err = as.CountDSL(dsl, typ)
if err == nil && offset < total {
assets, err = as.QueryDSL(dsl, typ, limit, offset)
}
} else {
companyID, _ := strconv.ParseInt(q.Get("company_id"), 10, 64)
taskID, _ := strconv.ParseInt(q.Get("task_id"), 10, 64)
switch {
case companyID > 0:
assets, err = as.QueryByCompany(companyID, typ, limit)
total = len(assets)
total, err = as.CountByCompany(companyID, typ)
if err == nil && offset < total {
assets, err = as.QueryByCompany(companyID, typ, limit, offset)
}
case taskID > 0:
assets, err = as.QueryByTask(taskID, typ, limit)
total = len(assets)
total, err = as.CountByTask(taskID, typ)
if err == nil && offset < total {
assets, err = as.QueryByTask(taskID, typ, limit, offset)
}
default:
if typ == "" {
typ = "root_domain"
}
assets, err = as.QueryByType(typ, limit, offset)
if err == nil {
total, _ = as.CountByType(typ)
total, err = as.CountByType(typ)
if err == nil && offset < total {
assets, err = as.QueryByType(typ, limit, offset)
}
}
}
if err != nil {
writeErr(w, 400, err.Error())
writeErr(w, http.StatusInternalServerError, err.Error())
return
}
writeJSON(w, 200, map[string]any{
@@ -284,7 +389,13 @@ func (s *Server) assetCounts(w http.ResponseWriter, r *http.Request) {
writeErr(w, 503, "database unavailable")
return
}
counts, err := as.CountsByType()
var counts map[string]int
var err error
if taskID, _ := strconv.ParseInt(r.URL.Query().Get("task_id"), 10, 64); taskID > 0 {
counts, err = as.CountsByTypeForTask(taskID)
} else {
counts, err = as.CountsByType()
}
if err != nil {
writeErr(w, 500, err.Error())
return
@@ -336,8 +447,8 @@ func (s *Server) insertAssets(w http.ResponseWriter, r *http.Request) {
Assets []struct {
Type string `json:"type"`
// root_domain / subdomain
Domain string `json:"domain"`
ICP string `json:"icp"`
Domain string `json:"domain"`
ICP string `json:"icp"`
RecordType string `json:"record_type"`
RecordValue []string `json:"record_value"`
// ip
@@ -420,6 +531,9 @@ func (s *Server) insertAssets(w http.ResponseWriter, r *http.Request) {
errs = append(errs, errEntry{Index: i, Error: err.Error()})
continue
}
if req.TaskID > 0 {
_ = as.SetTaskAssetSource(req.TaskID, id, "api", "通过资产 API 登记", nil)
}
results = append(results, result{Index: i, ID: id, Type: a.Type})
}
writeJSON(w, 200, map[string]any{"results": results, "errors": errs})
server/assets_scope_test.go
+265
View File
@@ -0,0 +1,265 @@
package server
import (
"bytes"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Autumn-27/artex/db"
)
func TestCompanyScopeInputsAcceptStructuredAndLegacyRules(t *testing.T) {
var inputs companyScopeInputs
if err := json.Unmarshal([]byte(`[
"example.com",
{"kind":"icp","value":"京 ICP备 123号"},
{"kind":"keyword","value":"Acme Security"}
]`), &inputs); err != nil {
t.Fatal(err)
}
if len(inputs) != 3 {
t.Fatalf("input count=%d want 3", len(inputs))
}
if inputs[0].Kind != "" || inputs[0].Value != "example.com" {
t.Fatalf("legacy rule changed: %+v", inputs[0])
}
if inputs[1].Kind != "icp" || inputs[2].Kind != "keyword" {
t.Fatalf("structured rules changed: %+v", inputs)
}
}
func TestCreateCompanyRejectsNormalizedDuplicateWithoutChangingScope(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
stamp := time.Now().UnixNano()
name := fmt.Sprintf("HTTP Strict Company %d", stamp)
companyID, _, _, _, _, err := m.pg.Companies().CreateCompanyWithScope(name, "", []db.ScopeInput{
{Kind: "domain", Value: fmt.Sprintf("existing-%d.example", stamp)},
}, "test")
if err != nil {
t.Fatal(err)
}
defer func() { _ = m.pg.Companies().DeleteCompany(companyID) }()
s := &Server{m: m}
body, err := json.Marshal(map[string]any{
"name": " " + name + " ",
"scope": []map[string]string{{"kind": "domain", "value": fmt.Sprintf("new-%d.example", stamp)}},
})
if err != nil {
t.Fatal(err)
}
req := httptest.NewRequest(http.MethodPost, "/api/companies", bytes.NewReader(body))
rec := httptest.NewRecorder()
s.createCompany(rec, req)
if rec.Code != http.StatusConflict {
t.Fatalf("duplicate create status=%d body=%s", rec.Code, rec.Body.String())
}
scope, err := m.pg.Companies().GetScope(companyID)
if err != nil {
t.Fatal(err)
}
if len(scope) != 1 || scope[0].Domain != fmt.Sprintf("existing-%d.example", stamp) {
t.Fatalf("duplicate HTTP create changed existing scope: %+v", scope)
}
}
func TestDeleteCompanyRefreshesLiveTaskCompanyIDs(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
companyID, _, err := m.pg.Companies().UpsertCompany(
fmt.Sprintf("Delete Company DTO %d", time.Now().UnixNano()), "",
)
if err != nil {
t.Fatal(err)
}
task, err := m.CreateTaskWithOptions("company deletion dto", "goal", db.TaskCreateOptions{
CompanyIDs: []int64{companyID},
})
if err != nil {
t.Fatal(err)
}
defer func() { _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{}) }()
s := &Server{m: m}
req := httptest.NewRequest(http.MethodDelete, "/api/companies/1", bytes.NewBufferString(`{}`))
req.SetPathValue("id", i64s(companyID))
rec := httptest.NewRecorder()
s.deleteCompany(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("delete company status=%d body=%s", rec.Code, rec.Body.String())
}
if dto := taskDTO(task, "created"); len(dto.CompanyIDs) != 0 {
t.Fatalf("live task DTO retained deleted company: %+v", dto.CompanyIDs)
}
persisted, err := m.pg.GetTask(mustTaskID(t, task.ID))
if err != nil {
t.Fatal(err)
}
if persisted == nil || len(persisted.CompanyIDs) != 0 {
t.Fatalf("persisted task retained deleted company: %+v", persisted)
}
}
func TestCompanyScopeHTTPErrorClassificationAndBounds(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
s := &Server{m: m}
t.Run("missing company is 404", func(t *testing.T) {
body := bytes.NewBufferString(`{"scope":[]}`)
req := httptest.NewRequest(http.MethodPost, "/api/companies/999999999/scope", body)
req.SetPathValue("id", "999999999")
rec := httptest.NewRecorder()
s.addCompanyScope(rec, req)
if rec.Code != http.StatusNotFound {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
})
t.Run("oversized Unicode scope is 400", func(t *testing.T) {
body, err := json.Marshal(map[string]any{
"name": "Oversized Scope",
"scope": []map[string]string{{
"kind": "keyword", "value": strings.Repeat("界", db.MaxCompanyScopeRawRunes+1),
}},
})
if err != nil {
t.Fatal(err)
}
req := httptest.NewRequest(http.MethodPost, "/api/companies", bytes.NewReader(body))
rec := httptest.NewRecorder()
s.createCompany(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
})
t.Run("too many scope rules is 400", func(t *testing.T) {
scope := make([]map[string]string, db.MaxCompanyScopeRules+1)
for i := range scope {
scope[i] = map[string]string{"kind": "keyword", "value": fmt.Sprintf("keyword-%d", i)}
}
body, err := json.Marshal(map[string]any{"name": "Too Many Scope Rules", "scope": scope})
if err != nil {
t.Fatal(err)
}
req := httptest.NewRequest(http.MethodPost, "/api/companies", bytes.NewReader(body))
rec := httptest.NewRecorder()
s.createCompany(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
})
t.Run("oversized request body is 413", func(t *testing.T) {
body := `{"name":"` + strings.Repeat("x", maxCompanyMutationBodyBytes+1) + `"}`
req := httptest.NewRequest(http.MethodPost, "/api/companies", strings.NewReader(body))
rec := httptest.NewRecorder()
s.createCompany(rec, req)
if rec.Code != http.StatusRequestEntityTooLarge {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
})
}
func TestListAssetsClassifiesValidationAndDatabaseErrors(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
s := &Server{m: m}
req := httptest.NewRequest(http.MethodGet, "/api/assets?dsl=(", nil)
rec := httptest.NewRecorder()
s.listAssets(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("invalid DSL status=%d body=%s", rec.Code, rec.Body.String())
}
if err := m.pg.Close(); err != nil {
t.Fatal(err)
}
req = httptest.NewRequest(http.MethodGet, "/api/assets?type=root_domain", nil)
rec = httptest.NewRecorder()
s.listAssets(rec, req)
if rec.Code != http.StatusInternalServerError {
t.Fatalf("database failure status=%d body=%s", rec.Code, rec.Body.String())
}
}
func TestDeleteCompanyRejectsBadJSONAndReportsMissing(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
s := &Server{m: m}
companyID, _, err := m.pg.Companies().UpsertCompany(
fmt.Sprintf("Delete JSON Company %d", time.Now().UnixNano()), "",
)
if err != nil {
t.Fatal(err)
}
defer func() { _ = m.pg.Companies().DeleteCompany(companyID) }()
for _, body := range []string{`{"delete_assets":`, `{}` + `{}`} {
req := httptest.NewRequest(http.MethodDelete, "/api/companies/1", strings.NewReader(body))
req.SetPathValue("id", i64s(companyID))
rec := httptest.NewRecorder()
s.deleteCompany(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("body=%q status=%d response=%s", body, rec.Code, rec.Body.String())
}
company, err := m.pg.Companies().GetCompany(companyID)
if err != nil || company == nil {
t.Fatalf("bad JSON deleted company: company=%v err=%v", company, err)
}
}
req := httptest.NewRequest(http.MethodDelete, "/api/companies/999999999", nil)
req.SetPathValue("id", "999999999")
rec := httptest.NewRecorder()
s.deleteCompany(rec, req)
if rec.Code != http.StatusNotFound {
t.Fatalf("missing delete status=%d body=%s", rec.Code, rec.Body.String())
}
}
func TestCompanyScopeSystemFailureIsHTTP500(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer m.Close()
s := &Server{m: m}
if err := m.pg.Close(); err != nil {
t.Fatal(err)
}
req := httptest.NewRequest(http.MethodPost, "/api/companies/1/scope", bytes.NewBufferString(`{"scope":[]}`))
req.SetPathValue("id", "1")
rec := httptest.NewRecorder()
s.addCompanyScope(rec, req)
if rec.Code != http.StatusInternalServerError {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
}
server/task_assets.go
+108
View File
@@ -0,0 +1,108 @@
package server
import (
"errors"
"net/http"
"strings"
"github.com/Autumn-27/artex/db"
)
const maxTaskAssetRequestBytes = 512 << 10
func writeTaskAssetError(w http.ResponseWriter, err error) {
switch {
case errors.Is(err, db.ErrTaskAssetInvalid):
writeErr(w, http.StatusBadRequest, err.Error())
case errors.Is(err, db.ErrTaskAssetTaskNotFound), errors.Is(err, db.ErrTaskAssetAssetNotFound):
writeErr(w, http.StatusNotFound, err.Error())
default:
writeErr(w, http.StatusInternalServerError, err.Error())
}
}
func (s *Server) attachTaskAssets(w http.ResponseWriter, r *http.Request) {
task, ok := s.m.Task(r.PathValue("id"))
if !ok {
writeErr(w, http.StatusNotFound, "task not found")
return
}
r.Body = http.MaxBytesReader(w, r.Body, maxTaskAssetRequestBytes)
var request struct {
AssetIDs []int64 `json:"asset_ids"`
SourceSummary string `json:"source_summary"`
Scope companyScopeInputs `json:"scope"`
}
if err := decode(r, &request); err != nil {
var tooLarge *http.MaxBytesError
if errors.As(err, &tooLarge) {
writeErr(w, http.StatusRequestEntityTooLarge, "请求正文过大")
} else {
writeErr(w, http.StatusBadRequest, err.Error())
}
return
}
request.SourceSummary = strings.TrimSpace(request.SourceSummary)
taskID, _ := parseTaskID(task.ID)
if request.Scope != nil && len(request.AssetIDs) > 0 {
writeErr(w, http.StatusBadRequest, "scope 与 asset_ids 不能同时提交")
return
}
if request.Scope != nil {
mutation, err := s.m.Assets().RegisterTaskAssetScopes(taskID, request.Scope)
if err != nil {
writeTaskAssetError(w, err)
return
}
task.Notify()
writeJSON(w, http.StatusOK, mutation)
return
}
mutation, err := s.m.Assets().AttachAssetsToTask(taskID, request.AssetIDs, request.SourceSummary)
if err != nil {
writeTaskAssetError(w, err)
return
}
task.Notify()
writeJSON(w, http.StatusOK, mutation)
}
func (s *Server) detachTaskAsset(w http.ResponseWriter, r *http.Request) {
task, ok := s.m.Task(r.PathValue("id"))
if !ok {
writeErr(w, http.StatusNotFound, "task not found")
return
}
assetID, ok := pathInt(r, "assetID")
if !ok {
writeErr(w, http.StatusBadRequest, "bad asset id")
return
}
taskID, _ := parseTaskID(task.ID)
detached, err := s.m.Assets().DetachAssetFromTask(taskID, assetID)
if err != nil {
writeTaskAssetError(w, err)
return
}
if !detached {
writeErr(w, http.StatusNotFound, "asset is not associated with this task")
return
}
task.Notify()
writeJSON(w, http.StatusOK, map[string]any{"detached": assetID})
}
func (s *Server) taskIntentAssets(w http.ResponseWriter, r *http.Request) {
task, ok := s.m.Task(r.PathValue("id"))
if !ok {
writeErr(w, http.StatusNotFound, "task not found")
return
}
taskID, _ := parseTaskID(task.ID)
assets, err := s.m.Assets().IntentAssets(taskID)
if err != nil {
writeTaskAssetError(w, err)
return
}
writeJSON(w, http.StatusOK, map[string]any{"assets": assets})
}
web/src/app/(main)/function/assets/page.tsx
+541 -263
View File
File diff suppressed because it is too large. Load diff
web/src/app/(main)/function/tasks/detail/_tabs/assets-tab.tsx
+510 -214
View File
@@ -1,56 +1,102 @@
"use client";
import * as React from "react";
import {
GlobeIcon,
NetworkIcon,
LayoutTemplateIcon,
LinkIcon,
KeyRoundIcon,
ChevronLeftIcon,
ChevronRightIcon,
GlobeIcon,
KeyRoundIcon,
LayoutTemplateIcon,
LinkIcon,
type LucideIcon,
NetworkIcon,
PlusIcon,
SmartphoneIcon,
Trash2Icon,
} from "lucide-react";
import { cn } from "@/lib/utils";
import { toast } from "sonner";
import { ScopeTextEditor } from "@/components/scope-text-editor";
import {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Card } from "@/components/ui/card";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import { Tabs, TabsContent } from "@/components/ui/tabs";
import {
Table,
TableBody,
TableCell,
TableHead,
TableHeader,
TableRow,
} from "@/components/ui/table";
import { Card, CardContent, CardFooter } from "@/components/ui/card";
import { Select, SelectContent, SelectGroup, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Sheet, SheetContent, SheetDescription, SheetFooter, SheetHeader, SheetTitle } from "@/components/ui/sheet";
import { Spinner } from "@/components/ui/spinner";
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
import { api } from "@/lib/api";
import type { Asset } from "@/lib/types";
import { parseCompanyScopeText } from "@/lib/company-scope";
import { taskAssetSourceLabel } from "@/lib/task-assets";
import type { Asset, NewAssetType } from "@/lib/types";
import { cn } from "@/lib/utils";
const PAGE_SIZES = [25, 50, 100, 200];
const METHOD_COLOR: Record<string, string> = {
GET: "bg-emerald-100 text-emerald-700",
POST: "bg-blue-100 text-blue-700",
PUT: "bg-amber-100 text-amber-700",
PATCH: "bg-orange-100 text-orange-700",
DELETE: "bg-red-100 text-red-700",
GET: "bg-emerald-100 text-emerald-700",
HEAD: "bg-purple-100 text-purple-700",
OPTIONS: "bg-slate-100 text-slate-600",
PATCH: "bg-orange-100 text-orange-700",
POST: "bg-blue-100 text-blue-700",
PUT: "bg-amber-100 text-amber-700",
};
const TABS: { key: NewAssetType; label: string; icon: LucideIcon }[] = [
{ key: "root_domain", label: "根域名", icon: GlobeIcon },
{ key: "ip", label: "IP", icon: NetworkIcon },
{ key: "subdomain", label: "子域名", icon: GlobeIcon },
{ key: "app", label: "应用", icon: SmartphoneIcon },
{ key: "service", label: "服务", icon: LayoutTemplateIcon },
{ key: "endpoint", label: "接口", icon: LinkIcon },
];
function firstText(values: Array<string | undefined>, fallback: string): string {
return values.find((value) => value?.trim()) ?? fallback;
}
function assetLabel(asset: Asset): string {
switch (asset.type) {
case "root_domain":
case "subdomain":
return firstText([asset.domain], `#${asset.id}`);
case "ip":
return firstText([asset.ip], `#${asset.id}`);
case "app":
return firstText([asset.app_name, asset.url], `#${asset.id}`);
case "service": {
const host = firstText([asset.ip, asset.domain], "");
const address = [host, asset.port].filter((value) => value !== undefined && value !== "").join(":");
return firstText([asset.url, address, asset.service_name], `#${asset.id}`);
}
case "endpoint":
return firstText([[asset.method, asset.url].filter(Boolean).join(" ")], `#${asset.id}`);
}
}
function MethodBadge({ method }: { method: string }) {
const m = method.toUpperCase();
const normalized = method.toUpperCase();
return (
<span className={cn("inline-block rounded px-1.5 py-0.5 font-mono text-[10px] font-semibold leading-none", METHOD_COLOR[m] ?? "bg-muted text-muted-foreground")}>
{m || "—"}
<span
className={cn(
"inline-block rounded px-1.5 py-0.5 font-mono font-semibold text-[10px] leading-none",
METHOD_COLOR[normalized] ?? "bg-muted text-muted-foreground",
)}
>
{normalized || "—"}
</span>
);
}
@@ -63,259 +109,470 @@ function statusTone(code: number) {
return "text-muted-foreground";
}
function fmtBytes(n?: number | null) {
if (!n || !Number.isFinite(n) || n <= 0) return "—";
const u = ["B", "KB", "MB"];
const i = Math.min(Math.floor(Math.log(n) / Math.log(1024)), u.length - 1);
return `${i === 0 ? n : (n / 1024 ** i).toFixed(1)} ${u[i]}`;
function fmtBytes(value?: number | null) {
if (!value || !Number.isFinite(value) || value <= 0) return "—";
const units = ["B", "KB", "MB"];
const index = Math.min(Math.floor(Math.log(value) / Math.log(1024)), units.length - 1);
return `${index === 0 ? value : (value / 1024 ** index).toFixed(1)} ${units[index]}`;
}
function Chips({ items, mono }: { items: string[]; mono?: boolean }) {
const clean = items.filter(Boolean);
if (clean.length === 0) return <span className="text-xs text-muted-foreground">—</span>;
if (clean.length === 0) return <span className="text-muted-foreground text-xs">—</span>;
return (
<div className="flex flex-wrap gap-1">
{clean.map((s, i) => (
<Badge key={i} variant="outline" className={cn("text-[10px]", mono && "font-mono")}>{s}</Badge>
{clean.map((item) => (
<Badge key={item} variant="outline" className={cn(mono && "font-mono")}>
{item}
</Badge>
))}
</div>
);
}
function SourceCell({ asset }: { asset: Asset }) {
const source = firstText([asset.task_source], "legacy");
const summary = firstText([asset.task_source_summary], "由历史任务资产关联迁移，暂无更详细来源说明");
return (
<Tooltip>
<TooltipTrigger asChild>
<Badge variant="outline" className="max-w-28 shrink-0 font-normal">
<span className="truncate">{taskAssetSourceLabel(source)}</span>
</Badge>
</TooltipTrigger>
<TooltipContent side="left" align="start" className="max-w-sm">
<div className="flex min-w-0 flex-col gap-1">
<span className="font-medium">{taskAssetSourceLabel(source)}</span>
<span className="[overflow-wrap:anywhere]">{summary}</span>
{asset.task_source_node_id ? (
<span className="font-mono opacity-80">来源节点 #{asset.task_source_node_id}</span>
) : null}
</div>
</TooltipContent>
</Tooltip>
);
}
function AssetCard({
cols, total, page, size, onSize, onPage, children,
children,
cols,
loaded,
onPage,
onSize,
page,
size,
total,
}: {
children: React.ReactNode;
cols: string[];
total: number;
loaded: boolean;
onPage: (page: number) => void;
onSize: (size: number) => void;
page: number;
size: number;
onSize: (v: number) => void;
onPage: (fn: (p: number) => number) => void;
children: React.ReactNode;
total: number;
}) {
const rows = React.Children.toArray(children);
const pageCount = Math.max(1, Math.ceil(total / size));
const start = total === 0 ? 0 : page * size + 1;
const end = page * size + rows.length;
const end = Math.min(total, page * size + rows.length);
let tableRows: React.ReactNode;
if (!loaded) {
tableRows = (
<TableRow>
<TableCell colSpan={cols.length} className="py-14 text-center">
<Spinner className="mx-auto" />
</TableCell>
</TableRow>
);
} else if (rows.length > 0) {
tableRows = rows;
} else {
tableRows = (
<TableRow>
<TableCell colSpan={cols.length} className="py-10 text-center text-muted-foreground text-sm">
当前分类暂无测试资产
</TableCell>
</TableRow>
);
}
return (
<Card className="flex min-h-0 flex-1 flex-col overflow-hidden py-0">
<div className="min-h-0 flex-1 overflow-auto scrollbar-thin scrollbar-track-transparent">
<Card className="flex min-h-0 flex-1 flex-col py-0">
<CardContent className="scrollbar-thin scrollbar-track-transparent min-h-0 flex-1 overflow-auto p-0">
<Table>
<TableHeader className="sticky top-0 z-10 bg-card">
<TableRow>
{cols.map((c) => <TableHead key={c}>{c}</TableHead>)}
{cols.map((column) => (
<TableHead key={column}>{column}</TableHead>
))}
</TableRow>
</TableHeader>
<TableBody>
{rows.length > 0 ? rows : (
<TableRow>
<TableCell colSpan={cols.length} className="py-12 text-center text-sm text-muted-foreground">
暂无数据
</TableCell>
</TableRow>
)}
</TableBody>
<TableBody>{tableRows}</TableBody>
</Table>
</div>
{total > 0 && (
<div className="flex shrink-0 items-center gap-2 border-t px-3 py-1.5 text-xs text-muted-foreground">
<Select value={String(size)} onValueChange={(v) => onSize(Number(v))}>
<SelectTrigger size="sm" className="h-7 w-24"><SelectValue /></SelectTrigger>
</CardContent>
{total > 0 ? (
<CardFooter className="gap-2 px-3 py-1.5 text-muted-foreground text-xs">
<Select
value={String(size)}
onValueChange={(value) => {
onSize(Number(value));
onPage(0);
}}
>
<SelectTrigger size="sm" className="w-24">
<SelectValue />
</SelectTrigger>
<SelectContent>
{PAGE_SIZES.map((n) => <SelectItem key={n} value={String(n)}>{n} / 页</SelectItem>)}
<SelectGroup>
{PAGE_SIZES.map((value) => (
<SelectItem key={value} value={String(value)}>
{value} / 页
</SelectItem>
))}
</SelectGroup>
</SelectContent>
</Select>
<span className="tabular-nums">{start}–{end} / {total}</span>
{pageCount > 1 && (
<div className="flex items-center gap-2">
<Button variant="outline" size="icon" className="size-7" disabled={page <= 0} onClick={() => onPage((p) => Math.max(0, p - 1))}>
<span className="tabular-nums">
{start}–{end} / {total}
</span>
{pageCount > 1 ? (
<div className="ml-auto flex items-center gap-2">
<Button
variant="outline"
size="icon-sm"
disabled={page <= 0}
onClick={() => onPage(Math.max(0, page - 1))}
aria-label="上一页"
>
<ChevronLeftIcon />
</Button>
<span className="tabular-nums">{page + 1} / {pageCount}</span>
<Button variant="outline" size="icon" className="size-7" disabled={page + 1 >= pageCount} onClick={() => onPage((p) => Math.min(pageCount - 1, p + 1))}>
<span className="tabular-nums">
{page + 1} / {pageCount}
</span>
<Button
variant="outline"
size="icon-sm"
disabled={page + 1 >= pageCount}
onClick={() => onPage(Math.min(pageCount - 1, page + 1))}
aria-label="下一页"
>
<ChevronRightIcon />
</Button>
</div>
)}
</div>
)}
) : null}
</CardFooter>
) : null}
</Card>
);
}
const TABS: { key: string; label: string; icon: LucideIcon }[] = [
{ key: "root_domain", label: "根域名", icon: GlobeIcon },
{ key: "ip", label: "IP", icon: NetworkIcon },
{ key: "subdomain", label: "子域名", icon: GlobeIcon },
{ key: "service", label: "服务", icon: LayoutTemplateIcon },
{ key: "endpoint", label: "接口", icon: LinkIcon },
];
export function AssetsTab({ taskId }: { taskId: string }) {
const [assets, setAssets] = React.useState<Asset[]>([]);
const [loaded, setLoaded] = React.useState(false);
const [tab, setTab] = React.useState("root_domain");
const [page, setPage] = React.useState(0);
const [size, setSize] = React.useState(50);
React.useEffect(() => setPage(0), [tab, size]);
function AddTaskAssetsSheet({
onAttached,
onOpenChange,
open,
taskId,
}: {
onAttached: () => void;
onOpenChange: (open: boolean) => void;
open: boolean;
taskId: string;
}) {
const [scopeText, setScopeText] = React.useState("");
const [saving, setSaving] = React.useState(false);
const parsedScope = React.useMemo(() => parseCompanyScopeText(scopeText), [scopeText]);
React.useEffect(() => {
let alive = true;
const load = async () => {
try {
const results = await Promise.all(
TABS.map((t) => api.taskAssets(taskId, t.key)),
);
if (!alive) return;
setAssets(results.flat());
setLoaded(true);
} catch {
if (alive) setLoaded(true);
}
};
load();
const timer = setInterval(load, 10_000);
return () => { alive = false; clearInterval(timer); };
}, [taskId]);
if (!open) return;
setScopeText("");
}, [open]);
const byType = React.useMemo(() => {
const m: Record<string, Asset[]> = { root_domain: [], ip: [], subdomain: [], service: [], endpoint: [] };
for (const a of assets) {
if (a.type in m) m[a.type].push(a);
const attach = async () => {
if (parsedScope.rules.length === 0 || parsedScope.errors.length > 0) return;
setSaving(true);
try {
const result = await api.registerTaskAssetScopes(taskId, parsedScope.rules);
const assetSummary = result.assets_linked + result.assets_existing;
toast.success(`已登记 ${result.requested} 条范围，关联 ${assetSummary} 项域名/IP 资产`);
onAttached();
onOpenChange(false);
} catch (reason) {
toast.error(`新增失败：${String((reason as Error)?.message ?? reason)}`);
} finally {
setSaving(false);
}
return m;
}, [assets]);
const counts: Record<string, number> = {
root_domain: byType.root_domain.length,
ip: byType.ip.length,
subdomain: byType.subdomain.length,
service: byType.service.length,
endpoint: byType.endpoint.length,
};
const totalAll = assets.length;
return (
<Sheet open={open} onOpenChange={onOpenChange}>
<SheetContent className="w-full sm:max-w-xl">
<SheetHeader>
<SheetTitle>新增测试资产</SheetTitle>
<SheetDescription>
直接填写测试范围。域名和 IP 会创建或复用全局资产；CIDR、ICP 和关键词作为 Agent 范围上下文。
</SheetDescription>
</SheetHeader>
<div className="flex min-h-0 flex-1 flex-col overflow-y-auto px-4">
<ScopeTextEditor
id="task-asset-scope"
value={scopeText}
onValueChange={setScopeText}
parsed={parsedScope}
label="测试资产与范围"
description="每行一条，自动识别域名、IP、CIDR、ICP 备案和关键词。"
/>
</div>
<SheetFooter>
<Button variant="outline" onClick={() => onOpenChange(false)} disabled={saving}>
取消
</Button>
<Button
onClick={() => void attach()}
disabled={saving || parsedScope.rules.length === 0 || parsedScope.errors.length > 0}
>
{saving ? <Spinner data-icon="inline-start" /> : <PlusIcon data-icon="inline-start" />}
登记 {parsedScope.rules.length > 0 ? parsedScope.rules.length : ""} 条
</Button>
</SheetFooter>
</SheetContent>
</Sheet>
);
}
if (!loaded) {
return <div className="py-20 text-center text-sm text-muted-foreground">加载中…</div>;
}
if (totalAll === 0) {
return (
<div className="py-20 text-center text-sm text-muted-foreground">
该任务暂无涉及资产
<p className="mt-1 text-xs">资产在任务执行过程中由 agent 发现并记录</p>
</div>
);
}
export function AssetsTab({ taskId }: { taskId: string }) {
const [rows, setRows] = React.useState<Asset[]>([]);
const [total, setTotal] = React.useState(0);
const [counts, setCounts] = React.useState<Record<string, number>>({});
const [loaded, setLoaded] = React.useState(false);
const [tab, setTab] = React.useState<NewAssetType>("root_domain");
const [page, setPage] = React.useState(0);
const [size, setSize] = React.useState(50);
const [refreshKey, setRefreshKey] = React.useState(0);
const [addOpen, setAddOpen] = React.useState(false);
const [removeTarget, setRemoveTarget] = React.useState<Asset | null>(null);
const [removing, setRemoving] = React.useState(false);
const cur = byType[tab] ?? [];
const slice = cur.slice(page * size, page * size + size);
React.useEffect(() => {
setPage(0);
setRows([]);
setLoaded(false);
}, []);
React.useEffect(() => {
let active = true;
const load = async () => {
try {
const [current, nextCounts] = await Promise.all([
api.taskAssets(taskId, tab, size, page * size),
api.assetCounts(taskId),
]);
if (!active) return;
setRows(current.assets);
setTotal(current.total);
setCounts(nextCounts ?? {});
} catch (reason) {
if (active) toast.error(`加载任务资产失败：${String((reason as Error)?.message ?? reason)}`);
} finally {
if (active) setLoaded(true);
}
};
void load();
const timer = setInterval(load, 10_000);
return () => {
active = false;
clearInterval(timer);
};
}, [page, refreshKey, size, tab, taskId]);
React.useEffect(() => {
const maxPage = Math.max(0, Math.ceil(total / size) - 1);
if (page > maxPage) setPage(maxPage);
}, [page, size, total]);
const refresh = React.useCallback(() => {
setLoaded(false);
setRefreshKey((current) => current + 1);
}, []);
const remove = async () => {
if (!removeTarget) return;
setRemoving(true);
try {
await api.detachTaskAsset(taskId, removeTarget.id);
toast.success(`已将 ${assetLabel(removeTarget)} 移出当前任务`);
setRemoveTarget(null);
refresh();
} catch (reason) {
toast.error(`移出失败：${String((reason as Error)?.message ?? reason)}`);
} finally {
setRemoving(false);
}
};
const removeButton = (asset: Asset) => (
<Button
variant="ghost"
size="icon-sm"
onClick={() => setRemoveTarget(asset)}
aria-label={`将资产 ${assetLabel(asset)} 移出任务`}
title="移出任务"
>
<Trash2Icon />
</Button>
);
const commonCardProps = { loaded, onPage: setPage, onSize: setSize, page, size, total };
const totalAll = TABS.reduce((sum, item) => sum + (counts[item.key] ?? 0), 0);
return (
<Tabs orientation="vertical" value={tab} onValueChange={setTab} className="items-start gap-4">
{/* left nav */}
<div className="w-36 shrink-0 rounded-lg border bg-card p-1 shadow-sm">
{TABS.map((t) => (
<button
key={t.key}
onClick={() => setTab(t.key)}
className={cn(
"flex w-full items-center gap-2 rounded-md px-2.5 py-1.5 text-sm transition-colors",
tab === t.key
? "bg-accent text-accent-foreground font-medium"
: "text-muted-foreground hover:bg-accent/60 hover:text-foreground",
)}
>
<t.icon className="size-3.5 shrink-0" />
<span className="flex-1 text-left">{t.label}</span>
<span className="tabular-nums text-xs opacity-70">{counts[t.key] ?? 0}</span>
</button>
))}
<div className="flex min-h-0 flex-1 flex-col gap-4">
<div className="flex flex-wrap items-center justify-between gap-3">
<div className="min-w-0">
<h2 className="font-medium text-sm">测试资产</h2>
<p className="text-muted-foreground text-xs">当前任务共关联 {totalAll} 项资产</p>
</div>
<Button size="sm" onClick={() => setAddOpen(true)}>
<PlusIcon data-icon="inline-start" />
新增测试资产
</Button>
</div>
{/* right content */}
<div className="min-w-0 flex-1 flex flex-col gap-4">
<Tabs value={tab} onValueChange={(value) => setTab(value as NewAssetType)} className="min-h-0 flex-1">
<div className="overflow-x-auto overflow-y-hidden">
<TabsList className="w-max">
{TABS.map((item) => (
<TabsTrigger key={item.key} value={item.key}>
<item.icon data-icon="inline-start" />
{item.label}
<span className="text-muted-foreground tabular-nums">{counts[item.key] ?? 0}</span>
</TabsTrigger>
))}
</TabsList>
</div>
<TabsContent value="root_domain" className="mt-0 flex min-h-0 flex-1 flex-col">
<AssetCard cols={["域名", "ICP 备案"]} total={cur.length} page={page} size={size} onSize={setSize} onPage={setPage}>
{slice.map((a) => (
<TableRow key={a.id}>
<TableCell className="font-mono text-xs font-medium">{a.domain}</TableCell>
<TableCell className="text-xs">{a.icp || "—"}</TableCell>
<TabsContent value="root_domain" className="mt-2 flex min-h-0 flex-1 flex-col">
<AssetCard cols={["域名", "ICP 备案", "来源", "操作"]} {...commonCardProps}>
{rows.map((asset) => (
<TableRow key={asset.id}>
<TableCell className="font-medium font-mono text-xs">{asset.domain}</TableCell>
<TableCell className="text-xs">{asset.icp || "—"}</TableCell>
<TableCell>
<SourceCell asset={asset} />
</TableCell>
<TableCell className="w-16">{removeButton(asset)}</TableCell>
</TableRow>
))}
</AssetCard>
</TabsContent>
<TabsContent value="ip" className="mt-0 flex min-h-0 flex-1 flex-col">
<AssetCard cols={["IP", "C段", "绑定域名", "开放端口"]} total={cur.length} page={page} size={size} onSize={setSize} onPage={setPage}>
{slice.map((a) => (
<TableRow key={a.id}>
<TableCell className="font-mono text-xs font-medium">{a.ip}</TableCell>
<TableCell className="font-mono text-xs">{a.c_segment || "—"}</TableCell>
<TableCell><Chips items={a.bound_domains ?? []} mono /></TableCell>
<TabsContent value="ip" className="mt-2 flex min-h-0 flex-1 flex-col">
<AssetCard cols={["IP", "C段", "绑定域名", "开放端口", "来源", "操作"]} {...commonCardProps}>
{rows.map((asset) => (
<TableRow key={asset.id}>
<TableCell className="font-medium font-mono text-xs">{asset.ip}</TableCell>
<TableCell className="font-mono text-xs">{asset.c_segment || "—"}</TableCell>
<TableCell>
<Chips items={asset.bound_domains ?? []} mono />
</TableCell>
<TableCell>
<Chips
items={(a.open_ports ?? []).map((p) =>
p.service ? `${p.port}/${p.service}` : String(p.port),
items={(asset.open_ports ?? []).map((item) =>
item.service ? `${item.port}/${item.service}` : String(item.port),
)}
mono
/>
</TableCell>
<TableCell>
<SourceCell asset={asset} />
</TableCell>
<TableCell className="w-16">{removeButton(asset)}</TableCell>
</TableRow>
))}
</AssetCard>
</TabsContent>
<TabsContent value="subdomain" className="mt-0 flex min-h-0 flex-1 flex-col">
<AssetCard cols={["域名", "根域名", "解析类型", "解析值"]} total={cur.length} page={page} size={size} onSize={setSize} onPage={setPage}>
{slice.map((a) => (
<TableRow key={a.id}>
<TableCell className="font-mono text-xs font-medium">{a.domain}</TableCell>
<TableCell className="font-mono text-xs">{a.root_domain || "—"}</TableCell>
<TableCell className="text-xs">{a.record_type || "—"}</TableCell>
<TableCell className="max-w-xs truncate font-mono text-xs">{(Array.isArray(a.record_value) ? a.record_value.join(", ") : a.record_value) || "—"}</TableCell>
<TabsContent value="subdomain" className="mt-2 flex min-h-0 flex-1 flex-col">
<AssetCard cols={["域名", "根域名", "解析类型", "解析值", "来源", "操作"]} {...commonCardProps}>
{rows.map((asset) => (
<TableRow key={asset.id}>
<TableCell className="font-medium font-mono text-xs">{asset.domain}</TableCell>
<TableCell className="font-mono text-xs">{asset.root_domain || "—"}</TableCell>
<TableCell className="text-xs">{asset.record_type || "—"}</TableCell>
<TableCell className="max-w-xs truncate font-mono text-xs">
{(Array.isArray(asset.record_value) ? asset.record_value.join(", ") : asset.record_value) || "—"}
</TableCell>
<TableCell>
<SourceCell asset={asset} />
</TableCell>
<TableCell className="w-16">{removeButton(asset)}</TableCell>
</TableRow>
))}
</AssetCard>
</TabsContent>
<TabsContent value="service" className="mt-0 flex min-h-0 flex-1 flex-col">
<TabsContent value="app" className="mt-2 flex min-h-0 flex-1 flex-col">
<AssetCard cols={["应用", "地址", "分类", "标题", "指纹", "来源", "操作"]} {...commonCardProps}>
{rows.map((asset) => (
<TableRow key={asset.id}>
<TableCell className="max-w-48 truncate font-medium text-xs">{asset.app_name || "—"}</TableCell>
<TableCell className="max-w-xs truncate font-mono text-xs" title={asset.url}>
{asset.url || "—"}
</TableCell>
<TableCell className="text-xs">{asset.category || "—"}</TableCell>
<TableCell className="max-w-48 truncate text-xs">{asset.page_title || "—"}</TableCell>
<TableCell>
<Chips items={asset.technologies ?? []} />
</TableCell>
<TableCell>
<SourceCell asset={asset} />
</TableCell>
<TableCell className="w-16">{removeButton(asset)}</TableCell>
</TableRow>
))}
</AssetCard>
</TabsContent>
<TabsContent value="service" className="mt-2 flex min-h-0 flex-1 flex-col">
<AssetCard
cols={["地址 / 服务", "状态码", "标题", "响应长度", "指纹", "认证"]}
total={cur.length}
page={page}
size={size}
onSize={setSize}
onPage={setPage}
cols={["地址 / 服务", "状态码", "标题", "响应长度", "指纹", "认证", "来源", "操作"]}
{...commonCardProps}
>
{slice.map((a) => {
const isHttp = a.service_type === "http";
const addr = isHttp
? (a.url || "")
: (a.service_name || [a.ip || a.domain, a.port].filter(Boolean).join(":"));
{rows.map((asset) => {
const isHttp = asset.service_type === "http";
const address = isHttp
? asset.url || ""
: asset.service_name || [asset.ip || asset.domain, asset.port].filter(Boolean).join(":");
return (
<TableRow key={a.id}>
<TableCell className="max-w-xs truncate font-mono text-xs" title={addr}>{addr || "—"}</TableCell>
<TableCell>
{a.status_code != null ? (
<span className={cn("font-mono text-xs font-semibold tabular-nums", statusTone(a.status_code))}>
{a.status_code}
</span>
) : "—"}
<TableRow key={asset.id}>
<TableCell className="max-w-xs truncate font-mono text-xs" title={address}>
{address || "—"}
</TableCell>
<TableCell className="max-w-[12rem] truncate text-xs">{a.page_title || "—"}</TableCell>
<TableCell className="text-xs tabular-nums">{fmtBytes(a.content_length)}</TableCell>
<TableCell><Chips items={a.technologies ?? []} /></TableCell>
<TableCell>
{(a.auth ?? []).length === 0 ? (
<span className="text-xs text-muted-foreground">—</span>
{asset.status_code != null ? (
<span
className={cn("font-mono font-semibold text-xs tabular-nums", statusTone(asset.status_code))}
>
{asset.status_code}
</span>
) : (
(a.auth ?? []).map((authItem, i) => {
"—"
)}
</TableCell>
<TableCell className="max-w-48 truncate text-xs">{asset.page_title || "—"}</TableCell>
<TableCell className="text-xs tabular-nums">{fmtBytes(asset.content_length)}</TableCell>
<TableCell>
<Chips items={asset.technologies ?? []} />
</TableCell>
<TableCell>
{(asset.auth ?? []).length === 0 ? (
<span className="text-muted-foreground text-xs">—</span>
) : (
(asset.auth ?? []).map((authItem) => {
const item = authItem as Record<string, string>;
return (
<span key={i} className="inline-flex items-center gap-1 text-[11px]">
<span
key={`${asset.id}-${item.type ?? ""}-${item.username ?? ""}-${JSON.stringify(item)}`}
className="inline-flex items-center gap-1 text-[11px]"
>
<KeyRoundIcon className="size-3 text-muted-foreground" />
<span className="font-mono">{item.type || item.username || "认证"}</span>
</span>
@@ -323,33 +580,72 @@ export function AssetsTab({ taskId }: { taskId: string }) {
})
)}
</TableCell>
<TableCell>
<SourceCell asset={asset} />
</TableCell>
<TableCell className="w-16">{removeButton(asset)}</TableCell>
</TableRow>
);
})}
</AssetCard>
</TabsContent>
<TabsContent value="endpoint" className="mt-0 flex min-h-0 flex-1 flex-col">
<AssetCard cols={["方法", "完整地址", "参数"]} total={cur.length} page={page} size={size} onSize={setSize} onPage={setPage}>
{slice.map((a) => (
<TableRow key={a.id}>
<TableCell className="w-16"><MethodBadge method={a.method || ""} /></TableCell>
<TableCell className="max-w-sm truncate font-mono text-xs" title={a.url}>{a.url || "—"}</TableCell>
<TabsContent value="endpoint" className="mt-2 flex min-h-0 flex-1 flex-col">
<AssetCard cols={["方法", "完整地址", "参数", "来源", "操作"]} {...commonCardProps}>
{rows.map((asset) => (
<TableRow key={asset.id}>
<TableCell className="w-16">
<MethodBadge method={asset.method || ""} />
</TableCell>
<TableCell className="max-w-sm truncate font-mono text-xs" title={asset.url}>
{asset.url || "—"}
</TableCell>
<TableCell>
<Chips
items={(a.params ?? []).map((p) => {
const item = p as Record<string, string>;
return item.name ? `${item.name}(${item.location || "?"})` : "?";
items={(asset.params ?? []).map((param) => {
const item = param as Record<string, string>;
return item.name ? `${item.name}(${item.location || item.in || "?"})` : "?";
})}
mono
/>
</TableCell>
<TableCell>
<SourceCell asset={asset} />
</TableCell>
<TableCell className="w-16">{removeButton(asset)}</TableCell>
</TableRow>
))}
</AssetCard>
</TabsContent>
</Tabs>
</div>
</Tabs>
<AddTaskAssetsSheet onAttached={refresh} onOpenChange={setAddOpen} open={addOpen} taskId={taskId} />
<AlertDialog open={Boolean(removeTarget)} onOpenChange={(open) => !open && !removing && setRemoveTarget(null)}>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>移出当前任务？</AlertDialogTitle>
<AlertDialogDescription className="[overflow-wrap:anywhere]">
{removeTarget ? `将“${assetLabel(removeTarget)}”从当前任务的测试资产中移出。` : ""}
全局资产、关联流量和历史黑板锚点会继续保留。
</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel disabled={removing}>取消</AlertDialogCancel>
<AlertDialogAction
variant="destructive"
disabled={removing}
onClick={(event) => {
event.preventDefault();
void remove();
}}
>
{removing ? <Spinner data-icon="inline-start" /> : <Trash2Icon data-icon="inline-start" />}
{removing ? "移出中" : "确认移出"}
</AlertDialogAction>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
</div>
);
}
web/src/components/scope-text-editor.tsx
+78
View File
@@ -0,0 +1,78 @@
"use client";
import * as React from "react";
import { Badge } from "@/components/ui/badge";
import { Field, FieldDescription, FieldError, FieldLabel } from "@/components/ui/field";
import { Textarea } from "@/components/ui/textarea";
import type { ParsedCompanyScopeText } from "@/lib/company-scope";
import type { CompanyScopeKind } from "@/lib/types";
const SCOPE_KIND_LABELS: Record<CompanyScopeKind, string> = {
domain: "域名",
ip: "IP",
cidr: "CIDR",
icp: "ICP",
keyword: "关键词",
};
export function ScopeTextEditor({
id,
value,
onValueChange,
parsed,
label = "资产范围",
description = "每行一条，自动识别域名、IP、CIDR、ICP 备案和企业关键词。",
}: {
id: string;
value: string;
onValueChange: (value: string) => void;
parsed: ParsedCompanyScopeText;
label?: string;
description?: string;
}) {
const counts = React.useMemo(() => {
const result = new Map<CompanyScopeKind, number>();
for (const rule of parsed.rules) result.set(rule.kind, (result.get(rule.kind) ?? 0) + 1);
return result;
}, [parsed.rules]);
return (
<Field data-invalid={parsed.errors.length > 0}>
<FieldLabel htmlFor={id}>{label}</FieldLabel>
<FieldDescription>{description}</FieldDescription>
<Textarea
id={id}
rows={8}
value={value}
aria-invalid={parsed.errors.length > 0}
placeholder={"example.com\n203.0.113.10\n198.51.100.0/24\n京ICP备12345678号-1\n企业名称关键词"}
className="min-h-36 resize-y font-mono text-sm"
onChange={(event) => onValueChange(event.target.value)}
/>
{parsed.rules.length > 0 && (
<div className="flex flex-wrap items-center gap-1.5 text-muted-foreground text-xs">
<span>已识别 {parsed.rules.length} 条</span>
{Object.entries(SCOPE_KIND_LABELS).map(([kind, kindLabel]) => {
const count = counts.get(kind as CompanyScopeKind) ?? 0;
return count > 0 ? (
<Badge key={kind} variant="secondary" className="font-mono tabular-nums">
{kindLabel} {count}
</Badge>
) : null;
})}
</div>
)}
{parsed.errors.length > 0 && (
<FieldError>
{parsed.errors.slice(0, 5).map((item) => (
<span key={`${item.line}-${item.error}`} className="block">
第 {item.line} 行：{item.error}
</span>
))}
{parsed.errors.length > 5 && <span className="block">另有 {parsed.errors.length - 5} 行错误</span>}
</FieldError>
)}
</Field>
);
}
web/src/lib/company-scope.ts
+212
View File
@@ -0,0 +1,212 @@
import type { CompanyScopeKind, CompanyScopeRule } from "@/lib/types";
export const MAX_COMPANY_SCOPE_RULES = 256;
export const MAX_COMPANY_SCOPE_VALUE_LENGTH = 1024;
export type CompanyScopeTextIssue = {
line: number;
rule?: CompanyScopeRule;
error?: string;
};
export type ParsedCompanyScopeText = {
rules: CompanyScopeRule[];
errors: Array<{ line: number; error: string }>;
};
export const COMPANY_SCOPE_KINDS = new Set<CompanyScopeKind>(["domain", "ip", "cidr", "icp", "keyword"]);
export function isCompanyScopeKind(value: unknown): value is CompanyScopeKind {
return typeof value === "string" && COMPANY_SCOPE_KINDS.has(value as CompanyScopeKind);
}
function isIPv4(value: string): boolean {
const parts = value.split(".");
return (
parts.length === 4 && parts.every((part) => (part === "0" || /^[1-9]\d{0,2}$/.test(part)) && Number(part) <= 255)
);
}
function isIPv6(value: string): boolean {
if (!value.includes(":") || /[^0-9a-f:.]/i.test(value) || value.includes(":::")) return false;
const halves = value.split("::");
if (halves.length > 2) return false;
const countSegments = (half: string): number | null => {
if (!half) return 0;
const segments = half.split(":");
let count = 0;
for (const [index, segment] of segments.entries()) {
if (segment.includes(".")) {
if (index !== segments.length - 1 || !isIPv4(segment)) return null;
count += 2;
} else {
if (!/^[0-9a-f]{1,4}$/i.test(segment)) return null;
count++;
}
}
return count;
};
const left = countSegments(halves[0]);
const right = countSegments(halves[1] ?? "");
if (left === null || right === null) return false;
return halves.length === 2 ? left + right < 8 : left === 8;
}
function ipVersion(value: string): 4 | 6 | null {
if (isIPv4(value)) return 4;
if (isIPv6(value)) return 6;
return null;
}
function looksLikeIPAddress(value: string): boolean {
const trimmed = value.trim();
if (trimmed.includes("://") || /\s/.test(trimmed)) return false;
if ((trimmed.match(/:/g) ?? []).length >= 2) {
const parts = trimmed.split(":");
let validSegments = 0;
for (const part of parts) {
if (!part) {
if (validSegments > 0 || trimmed.startsWith("::")) return true;
continue;
}
if (part.length > 4 || !/^[0-9a-f]+$/i.test(part)) return false;
validSegments++;
if (validSegments >= 2) return true;
}
return false;
}
return trimmed.includes(".") && /^[0-9.]+$/.test(trimmed);
}
function domainHostname(value: string): string | null {
try {
let candidate = value;
if (value.startsWith("//")) candidate = `http:${value}`;
else if (!value.includes("://")) candidate = `http://${value}`;
const url = new URL(candidate);
return (
url.hostname
.replace(/^\[|\]$/g, "")
.replace(/\.$/, "")
.toLowerCase() || null
);
} catch {
return null;
}
}
export function companyScopeRuleError(rule: CompanyScopeRule): string {
const value = rule.value.trim();
if (!value) return "请填写范围值";
if (Array.from(value).length > MAX_COMPANY_SCOPE_VALUE_LENGTH) {
return `最多 ${MAX_COMPANY_SCOPE_VALUE_LENGTH} 个字符`;
}
if (rule.kind === "domain") {
if (/\s/.test(value)) return "请输入不含空格的有效域名或 URL";
const hostname = domainHostname(value);
if (!hostname || ipVersion(hostname) !== null) return "请输入有效域名或 URL";
const labels = hostname.split(".");
if (
hostname.length > 253 ||
labels.length < 2 ||
labels.some(
(label) =>
!label || label.length > 63 || label.startsWith("-") || label.endsWith("-") || !/^[a-z0-9-]+$/i.test(label),
)
) {
return "请输入有效域名或 URL";
}
}
if (rule.kind === "ip" && ipVersion(value) === null) return "请输入有效 IP";
if (rule.kind === "cidr") {
const separator = value.lastIndexOf("/");
if (separator <= 0) return "请输入 CIDR 网段";
const address = value.slice(0, separator);
const prefixText = value.slice(separator + 1);
const version = ipVersion(address);
if (version === null || !/^\d+$/.test(prefixText)) return "请输入 CIDR 网段";
const prefix = Number(prefixText);
if (version === 4 && (prefix < 16 || prefix > 32)) return "IPv4 网段前缀需为 /16 至 /32";
if (version === 6 && (prefix < 32 || prefix > 128)) return "IPv6 网段前缀需为 /32 至 /128";
}
return "";
}
export function classifyCompanyScopeLine(
raw: string,
line: number,
preservedRule?: CompanyScopeRule,
): CompanyScopeTextIssue {
const value = raw.trim();
if (Array.from(value).length > MAX_COMPANY_SCOPE_VALUE_LENGTH) {
return { line, error: `最多 ${MAX_COMPANY_SCOPE_VALUE_LENGTH} 个字符` };
}
if (preservedRule) {
const rule = { kind: preservedRule.kind, value };
const error = companyScopeRuleError(rule);
return error ? { line, error } : { line, rule };
}
const separator = value.lastIndexOf("/");
if (
separator > 0 &&
(ipVersion(value.slice(0, separator)) !== null || looksLikeIPAddress(value.slice(0, separator)))
) {
const rule: CompanyScopeRule = { kind: "cidr", value };
const error = companyScopeRuleError(rule);
return error ? { line, error } : { line, rule };
}
if (ipVersion(value) !== null) return { line, rule: { kind: "ip", value } };
if (looksLikeIPAddress(value)) return { line, error: "请输入有效 IP" };
const looksLikeDomain = value.includes("://") || (!/\s/.test(value) && value.includes("."));
if (looksLikeDomain) {
const hostname = domainHostname(value);
if (hostname && ipVersion(hostname) !== null) return { line, rule: { kind: "ip", value: hostname } };
const rule: CompanyScopeRule = { kind: "domain", value };
const error = companyScopeRuleError(rule);
return error ? { line, error } : { line, rule };
}
if (/icp|备案/i.test(value)) return { line, rule: { kind: "icp", value } };
return { line, rule: { kind: "keyword", value } };
}
export function parseCompanyScopeText(
value: string,
options: { preservedRules?: CompanyScopeRule[]; maxRules?: number } = {},
): ParsedCompanyScopeText {
const nonEmpty = value
.split(/\r?\n/)
.map((raw, index) => ({ raw, line: index + 1 }))
.filter(({ raw }) => raw.trim());
const preserved = new Map<string, CompanyScopeRule[]>();
for (const rule of options.preservedRules ?? []) {
const key = rule.value.trim();
preserved.set(key, [...(preserved.get(key) ?? []), rule]);
}
const issues = nonEmpty.map(({ raw, line }) => {
const key = raw.trim();
const queue = preserved.get(key);
const preservedRule = queue?.shift();
return classifyCompanyScopeLine(raw, line, preservedRule);
});
const maxRules = options.maxRules ?? MAX_COMPANY_SCOPE_RULES;
if (nonEmpty.length > maxRules) {
issues.push({
line: nonEmpty[Math.max(0, maxRules)]?.line ?? 1,
error: `最多允许 ${maxRules} 条范围，当前 ${nonEmpty.length} 条`,
});
}
return {
rules: issues.flatMap((item) => (item.rule ? [item.rule] : [])),
errors: issues.flatMap((item) => (item.error ? [{ line: item.line, error: item.error }] : [])),
};
}
export function normalizeCompanyScopeValue(rule: CompanyScopeRule): string {
const value = rule.value.trim();
if (rule.kind === "domain") return domainHostname(value) ?? value.toLowerCase();
if (rule.kind === "icp") return value.toLowerCase().replace(/\s/gu, "");
if (rule.kind === "keyword") return value.toLowerCase().split(/\s+/u).filter(Boolean).join(" ");
return value.toLowerCase();
}
web/src/lib/task-assets.ts
+28
View File
@@ -0,0 +1,28 @@
import type { NewAssetType } from "@/lib/types";
const ASSET_TYPE_LABELS: Record<NewAssetType, string> = {
app: "应用",
endpoint: "接口",
ip: "IP",
root_domain: "根域名",
service: "服务",
subdomain: "子域名",
};
const TASK_ASSET_SOURCE_LABELS: Record<string, string> = {
agent: "Agent 发现",
anchor: "黑板锚点",
api: "资产 API",
legacy: "历史关联",
manual: "人工加入",
system: "系统关联",
task: "任务初始化",
};
export function taskAssetTypeLabel(type: NewAssetType): string {
return ASSET_TYPE_LABELS[type];
}
export function taskAssetSourceLabel(source: string): string {
return TASK_ASSET_SOURCE_LABELS[source] ?? source;
}