diff --git a/db/asset_dsl.go b/db/asset_dsl.go index 24c7f25..dbe8683 100644 --- a/db/asset_dsl.go +++ b/db/asset_dsl.go @@ -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 diff --git a/db/assets.go b/db/assets.go index fee5d94..ec3a7cb 100644 --- a/db/assets.go +++ b/db/assets.go @@ -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 diff --git a/db/assets_test.go b/db/assets_test.go index 4c5b5b2..c2ae2e5 100644 --- a/db/assets_test.go +++ b/db/assets_test.go @@ -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) + } +} diff --git a/db/companies.go b/db/companies.go index 84be395..e360002 100644 --- a/db/companies.go +++ b/db/companies.go @@ -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 } diff --git a/db/companies_test.go b/db/companies_test.go index 091e153..f1592c8 100644 --- a/db/companies_test.go +++ b/db/companies_test.go @@ -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 { diff --git a/db/company_lock_test.go b/db/company_lock_test.go new file mode 100644 index 0000000..5f4ea24 --- /dev/null +++ b/db/company_lock_test.go @@ -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) + } + } +} diff --git a/db/company_scope.go b/db/company_scope.go index 42207c9..8fb8c08 100644 --- a/db/company_scope.go +++ b/db/company_scope.go @@ -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 -} diff --git a/db/company_scope_consistency_test.go b/db/company_scope_consistency_test.go new file mode 100644 index 0000000..19e83fd --- /dev/null +++ b/db/company_scope_consistency_test.go @@ -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) + } +} diff --git a/db/company_scope_test.go b/db/company_scope_test.go index 09e0974..713938e 100644 --- a/db/company_scope_test.go +++ b/db/company_scope_test.go @@ -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) { diff --git a/db/task_assets.go b/db/task_assets.go new file mode 100644 index 0000000..db87fef --- /dev/null +++ b/db/task_assets.go @@ -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() +} diff --git a/db/task_assets_test.go b/db/task_assets_test.go new file mode 100644 index 0000000..4c19747 --- /dev/null +++ b/db/task_assets_test.go @@ -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]) + } +} diff --git a/db/task_scope.go b/db/task_scope.go index 52d35b6..f584b5d 100644 --- a/db/task_scope.go +++ b/db/task_scope.go @@ -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 } diff --git a/server/assets.go b/server/assets.go index 6712a8a..dc854b2 100644 --- a/server/assets.go +++ b/server/assets.go @@ -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}) diff --git a/server/assets_scope_test.go b/server/assets_scope_test.go new file mode 100644 index 0000000..c3e6b7c --- /dev/null +++ b/server/assets_scope_test.go @@ -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()) + } +} diff --git a/server/task_assets.go b/server/task_assets.go new file mode 100644 index 0000000..cda4287 --- /dev/null +++ b/server/task_assets.go @@ -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}) +} diff --git a/web/src/app/(main)/function/assets/page.tsx b/web/src/app/(main)/function/assets/page.tsx index b1892d4..b9b1b76 100644 --- a/web/src/app/(main)/function/assets/page.tsx +++ b/web/src/app/(main)/function/assets/page.tsx @@ -1,56 +1,25 @@ "use client"; import * as React from "react"; + import { BuildingIcon, - GlobeIcon, - NetworkIcon, - LayoutTemplateIcon, - LinkIcon, - SmartphoneIcon, - SearchIcon, - KeyRoundIcon, ChevronLeftIcon, ChevronRightIcon, - RefreshCwIcon, - Trash2Icon, + GlobeIcon, + KeyRoundIcon, + LayoutTemplateIcon, + LinkIcon, type LucideIcon, + NetworkIcon, + RefreshCwIcon, + SearchIcon, + SmartphoneIcon, + Trash2Icon, } from "lucide-react"; - import { toast } from "sonner"; -import { Card } from "@/components/ui/card"; -import { Input } from "@/components/ui/input"; -import { Badge } from "@/components/ui/badge"; -import { Button } from "@/components/ui/button"; -import { Label } from "@/components/ui/label"; -import { Textarea } from "@/components/ui/textarea"; -import { Avatar, AvatarFallback, AvatarImage } from "@/components/ui/avatar"; -import { - Dialog, - DialogContent, - DialogDescription, - DialogFooter, - DialogHeader, - DialogTitle, - DialogTrigger, -} from "@/components/ui/dialog"; -import { - Select, - SelectContent, - SelectItem, - SelectTrigger, - SelectValue, -} from "@/components/ui/select"; -import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; -import { - Table, - TableBody, - TableCell, - TableHead, - TableHeader, - TableRow, -} from "@/components/ui/table"; +import { ScopeTextEditor } from "@/components/scope-text-editor"; import { AlertDialog, AlertDialogAction, @@ -61,38 +30,205 @@ import { AlertDialogHeader, AlertDialogTitle, } from "@/components/ui/alert-dialog"; +import { Avatar, AvatarFallback, AvatarImage } from "@/components/ui/avatar"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Card } from "@/components/ui/card"; import { Checkbox } from "@/components/ui/checkbox"; -import { cn } from "@/lib/utils"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, + DialogTrigger, +} from "@/components/ui/dialog"; +import { Field, FieldGroup, FieldLabel } from "@/components/ui/field"; +import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectGroup, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { + Sheet, + SheetContent, + SheetDescription, + SheetFooter, + SheetHeader, + SheetTitle, + SheetTrigger, +} from "@/components/ui/sheet"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { api } from "@/lib/api"; -import type { Asset, Company } from "@/lib/types"; +import { MAX_COMPANY_SCOPE_RULES, parseCompanyScopeText } from "@/lib/company-scope"; +import type { Asset, Company, CompanyScopeRule } from "@/lib/types"; +import { cn } from "@/lib/utils"; // ── DSL autocomplete ────────────────────────────────────────────────────────── const DSL_FIELDS: { name: string; desc: string; ops: { op: string; desc: string }[] }[] = [ - { name: "domain", desc: "域名(根域名/子域名/服务域名)", ops: [{ op: "=", desc: "模糊匹配" }, { op: "==", desc: "精确匹配" }, { op: "!=", desc: "排除" }] }, - { name: "ip", desc: "IPv4/IPv6 地址", ops: [{ op: "=", desc: "模糊匹配" }, { op: "==", desc: "精确匹配" }, { op: "!=", desc: "排除" }] }, - { name: "url", desc: "完整 URL(服务/接口)", ops: [{ op: "=", desc: "模糊匹配" }, { op: "==", desc: "精确匹配" }, { op: "!=", desc: "排除" }] }, - { name: "root_domain", desc: "根域名", ops: [{ op: "=", desc: "模糊匹配" }, { op: "==", desc: "精确匹配" }, { op: "!=", desc: "排除" }] }, - { name: "page_title", desc: "页面标题(HTTP 服务)", ops: [{ op: "=", desc: "模糊匹配" }, { op: "==", desc: "精确匹配" }, { op: "!=", desc: "排除" }] }, - { name: "icp", desc: "ICP 备案号", ops: [{ op: "=", desc: "模糊匹配" }, { op: "==", desc: "精确匹配" }, { op: "!=", desc: "排除" }] }, - { name: "service_name", desc: "服务名称(非 HTTP 服务)", ops: [{ op: "=", desc: "模糊匹配" }, { op: "==", desc: "精确匹配" }, { op: "!=", desc: "排除" }] }, - { name: "app_name", desc: "应用名称", ops: [{ op: "=", desc: "模糊匹配" }, { op: "==", desc: "精确匹配" }, { op: "!=", desc: "排除" }] }, - { name: "bundle_id", desc: "应用 Bundle ID", ops: [{ op: "=", desc: "模糊匹配" }, { op: "==", desc: "精确匹配" }, { op: "!=", desc: "排除" }] }, - { name: "category", desc: "应用分类", ops: [{ op: "=", desc: "模糊匹配" }, { op: "==", desc: "精确匹配" }, { op: "!=", desc: "排除" }] }, - { name: "app_icp", desc: "应用 ICP 备案", ops: [{ op: "=", desc: "模糊匹配" }, { op: "==", desc: "精确匹配" }, { op: "!=", desc: "排除" }] }, - { name: "method", desc: "HTTP 方法 GET/POST/PUT/…", ops: [{ op: "==", desc: "精确匹配" }, { op: "!=", desc: "排除" }] }, - { name: "service_type", desc: "服务类型:http | other", ops: [{ op: "==", desc: "精确匹配" }, { op: "!=", desc: "排除" }] }, - { name: "record_type", desc: "DNS 解析类型 A/CNAME/MX/…", ops: [{ op: "==", desc: "精确匹配" }, { op: "!=", desc: "排除" }] }, - { name: "technology", desc: "技术指纹(数组字段)", ops: [{ op: "=", desc: "模糊匹配" }, { op: "==", desc: "精确匹配" }, { op: "!=", desc: "排除" }] }, - { name: "port", desc: "端口号(整数)", ops: [{ op: "==", desc: "等于" }, { op: "!=", desc: "不等于" }, { op: ">", desc: "大于" }, { op: ">=", desc: "大于等于" }, { op: "<", desc: "小于" }, { op: "<=", desc: "小于等于" }] }, - { name: "status_code", desc: "HTTP 状态码(整数)", ops: [{ op: "==", desc: "等于" }, { op: "!=", desc: "不等于" }, { op: ">", desc: "大于" }, { op: ">=", desc: "大于等于" }, { op: "<", desc: "小于" }, { op: "<=", desc: "小于等于" }] }, - { name: "company_id", desc: "归属企业 ID(整数)", ops: [{ op: "==", desc: "等于" }] }, - { name: "task_id", desc: "来源任务 ID(整数)", ops: [{ op: "==", desc: "等于" }] }, + { + name: "domain", + desc: "域名(根域名/子域名/服务域名)", + ops: [ + { op: "=", desc: "模糊匹配" }, + { op: "==", desc: "精确匹配" }, + { op: "!=", desc: "排除" }, + ], + }, + { + name: "ip", + desc: "IPv4/IPv6 地址", + ops: [ + { op: "=", desc: "模糊匹配" }, + { op: "==", desc: "精确匹配" }, + { op: "!=", desc: "排除" }, + ], + }, + { + name: "url", + desc: "完整 URL(服务/接口)", + ops: [ + { op: "=", desc: "模糊匹配" }, + { op: "==", desc: "精确匹配" }, + { op: "!=", desc: "排除" }, + ], + }, + { + name: "root_domain", + desc: "根域名", + ops: [ + { op: "=", desc: "模糊匹配" }, + { op: "==", desc: "精确匹配" }, + { op: "!=", desc: "排除" }, + ], + }, + { + name: "page_title", + desc: "页面标题(HTTP 服务)", + ops: [ + { op: "=", desc: "模糊匹配" }, + { op: "==", desc: "精确匹配" }, + { op: "!=", desc: "排除" }, + ], + }, + { + name: "icp", + desc: "ICP 备案号", + ops: [ + { op: "=", desc: "模糊匹配" }, + { op: "==", desc: "精确匹配" }, + { op: "!=", desc: "排除" }, + ], + }, + { + name: "service_name", + desc: "服务名称(非 HTTP 服务)", + ops: [ + { op: "=", desc: "模糊匹配" }, + { op: "==", desc: "精确匹配" }, + { op: "!=", desc: "排除" }, + ], + }, + { + name: "app_name", + desc: "应用名称", + ops: [ + { op: "=", desc: "模糊匹配" }, + { op: "==", desc: "精确匹配" }, + { op: "!=", desc: "排除" }, + ], + }, + { + name: "bundle_id", + desc: "应用 Bundle ID", + ops: [ + { op: "=", desc: "模糊匹配" }, + { op: "==", desc: "精确匹配" }, + { op: "!=", desc: "排除" }, + ], + }, + { + name: "category", + desc: "应用分类", + ops: [ + { op: "=", desc: "模糊匹配" }, + { op: "==", desc: "精确匹配" }, + { op: "!=", desc: "排除" }, + ], + }, + { + name: "app_icp", + desc: "应用 ICP 备案", + ops: [ + { op: "=", desc: "模糊匹配" }, + { op: "==", desc: "精确匹配" }, + { op: "!=", desc: "排除" }, + ], + }, + { + name: "method", + desc: "HTTP 方法 GET/POST/PUT/…", + ops: [ + { op: "==", desc: "精确匹配" }, + { op: "!=", desc: "排除" }, + ], + }, + { + name: "service_type", + desc: "服务类型:http | other", + ops: [ + { op: "==", desc: "精确匹配" }, + { op: "!=", desc: "排除" }, + ], + }, + { + name: "record_type", + desc: "DNS 解析类型 A/CNAME/MX/…", + ops: [ + { op: "==", desc: "精确匹配" }, + { op: "!=", desc: "排除" }, + ], + }, + { + name: "technology", + desc: "技术指纹(数组字段)", + ops: [ + { op: "=", desc: "模糊匹配" }, + { op: "==", desc: "精确匹配" }, + { op: "!=", desc: "排除" }, + ], + }, + { + name: "port", + desc: "端口号(整数)", + ops: [ + { op: "==", desc: "等于" }, + { op: "!=", desc: "不等于" }, + { op: ">", desc: "大于" }, + { op: ">=", desc: "大于等于" }, + { op: "<", desc: "小于" }, + { op: "<=", desc: "小于等于" }, + ], + }, + { + name: "status_code", + desc: "HTTP 状态码(整数)", + ops: [ + { op: "==", desc: "等于" }, + { op: "!=", desc: "不等于" }, + { op: ">", desc: "大于" }, + { op: ">=", desc: "大于等于" }, + { op: "<", desc: "小于" }, + { op: "<=", desc: "小于等于" }, + ], + }, + { name: "company_id", desc: "归属企业 ID(整数)", ops: [{ op: "==", desc: "等于" }] }, + { name: "task_id", desc: "来源任务 ID(整数)", ops: [{ op: "==", desc: "等于" }] }, ]; const LOGIC_OPS = [ { label: "AND", desc: "且(两个条件都满足)" }, - { label: "OR", desc: "或(满足其中之一)" }, + { label: "OR", desc: "或(满足其中之一)" }, ]; interface DslSuggestion { @@ -129,23 +265,18 @@ function getDslSuggestions(text: string, cursor: number): DslSuggestion[] { // Everything before the current token (trimmed) const beforeToken = before.slice(0, tokenStart).trimEnd(); - const afterExpression = - beforeToken.length > 0 && - !/\b(AND|OR)\s*$/i.test(beforeToken) && - !beforeToken.endsWith("("); + const afterExpression = beforeToken.length > 0 && !/\b(AND|OR)\s*$/i.test(beforeToken) && !beforeToken.endsWith("("); // Current token is a prefix of AND/OR and follows a complete expression if (/^(a|an|and|o|or)$/i.test(currentToken) && afterExpression) { - return LOGIC_OPS.filter((l) => l.label.startsWith(currentToken.toUpperCase())).map( - ({ label, desc }) => ({ - kind: "logic", - label, - desc, - replaceStart: tokenStart, - replaceEnd: cursor, - insertText: `${label} `, - }), - ); + return LOGIC_OPS.filter((l) => l.label.startsWith(currentToken.toUpperCase())).map(({ label, desc }) => ({ + kind: "logic", + label, + desc, + replaceStart: tokenStart, + replaceEnd: cursor, + insertText: `${label} `, + })); } // No current token, after a complete expression → suggest AND/OR @@ -172,18 +303,15 @@ function getDslSuggestions(text: string, cursor: number): DslSuggestion[] { })); } -function applyDslSuggestion( - text: string, - s: DslSuggestion, -): { text: string; cursor: number } { +function applyDslSuggestion(text: string, s: DslSuggestion): { text: string; cursor: number } { const newText = text.slice(0, s.replaceStart) + s.insertText + text.slice(s.replaceEnd); return { text: newText, cursor: s.replaceStart + s.insertText.length }; } const KIND_STYLE: Record = { - field: "text-blue-500 dark:text-blue-400", + field: "text-blue-500 dark:text-blue-400", operator: "text-amber-500 dark:text-amber-400", - logic: "text-emerald-500 dark:text-emerald-400", + logic: "text-emerald-500 dark:text-emerald-400", }; const METHOD_COLOR: Record = { @@ -199,7 +327,12 @@ const METHOD_COLOR: Record = { function MethodBadge({ method }: { method: string }) { const m = method.toUpperCase(); return ( - + {m || "—"} ); @@ -216,13 +349,13 @@ function statusTone(code: number) { const PAGE_SIZES = [25, 50, 100, 200]; const TABS: { key: string; label: string; icon: LucideIcon }[] = [ - { key: "company", label: "企业", icon: BuildingIcon }, + { key: "company", label: "企业", icon: BuildingIcon }, { 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 }, + { 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 }, ]; export default function AssetsPage() { @@ -238,6 +371,8 @@ export default function AssetsPage() { const [size, setSize] = React.useState(50); const [dslError, setDslError] = React.useState(""); const [refreshKey, setRefreshKey] = React.useState(0); + const assetsRequestRef = React.useRef(0); + const [rowsTab, setRowsTab] = React.useState(""); // asset selection & delete const [selected, setSelected] = React.useState>(new Set()); @@ -252,8 +387,18 @@ export default function AssetsPage() { // Companies + per-type counts (tab badges) — loaded on demand, no background polling. const loadMeta = React.useCallback(() => { - api.companies().then(setCompanies).catch(() => {}); - api.assetCounts().then(setCounts).catch(() => {}); + api + .companies() + .then(setCompanies) + .catch(() => { + /* Keep the last successful company snapshot on a transient failure. */ + }); + api + .assetCounts() + .then(setCounts) + .catch(() => { + /* Keep the last successful counters on a transient failure. */ + }); }, []); // Manual refresh: reload counts/companies and re-fetch the current page. @@ -266,7 +411,8 @@ export default function AssetsPage() { const toggleSelect = (id: number) => { setSelected((prev) => { const next = new Set(prev); - if (next.has(id)) next.delete(id); else next.add(id); + if (next.has(id)) next.delete(id); + else next.add(id); return next; }); }; @@ -275,8 +421,15 @@ export default function AssetsPage() { setSelected((prev) => { const allSelected = ids.every((id) => prev.has(id)); const next = new Set(prev); - if (allSelected) ids.forEach((id) => next.delete(id)); - else ids.forEach((id) => next.add(id)); + if (allSelected) { + ids.forEach((id) => { + next.delete(id); + }); + } else { + ids.forEach((id) => { + next.add(id); + }); + } return next; }); }; @@ -306,9 +459,10 @@ export default function AssetsPage() { setCompanyDeleting(true); try { const res = await api.deleteCompany(companyDeleteTarget.id, companyDeleteAssets); - const msg = companyDeleteAssets && res.assets_deleted > 0 - ? `已删除企业,同时删除 ${res.assets_deleted} 条资产` - : "已删除企业"; + const msg = + companyDeleteAssets && res.assets_deleted > 0 + ? `已删除企业,同时删除 ${res.assets_deleted} 条资产` + : "已删除企业"; toast.success(msg); refresh(); } catch (e) { @@ -325,43 +479,64 @@ export default function AssetsPage() { }, [loadMeta]); // Reset query + page + selection when switching tabs + // biome-ignore lint/correctness/useExhaustiveDependencies: tab changes intentionally reset tab-local controls. React.useEffect(() => { setQuery(""); setDslError(""); setSelected(new Set()); + setRows([]); + setRowsTab(""); + setTotal(0); + setLoaded(false); }, [tab]); + // biome-ignore lint/correctness/useExhaustiveDependencies: these controls intentionally reset server pagination. React.useEffect(() => setPage(0), [tab, size, query]); const dslMode = query.trim() !== ""; // Server-side paginated page loader for the active data tab (company tab excluded). + // biome-ignore lint/correctness/useExhaustiveDependencies: refreshKey is an explicit manual-reload trigger. React.useEffect(() => { if (tab === "company") return; + const request = ++assetsRequestRef.current; const dsl = query.trim(); setLoading(true); const offset = page * size; const run = async () => { try { - const r = dsl - ? await api.searchAssets(dsl, tab, size, offset) - : await api.assets(tab, size, offset); + const r = dsl ? await api.searchAssets(dsl, tab, size, offset) : await api.assets(tab, size, offset); + if (assetsRequestRef.current !== request) return; setRows(r.assets); + setRowsTab(tab); setTotal(r.total); setDslError(""); } catch (e) { + if (assetsRequestRef.current !== request) return; setDslError(String((e as Error)?.message ?? e)); setRows([]); + setRowsTab(tab); setTotal(0); } finally { - setLoading(false); - setLoaded(true); + if (assetsRequestRef.current === request) { + setLoading(false); + setLoaded(true); + } } }; const tid = setTimeout(run, dsl ? 400 : 0); - return () => clearTimeout(tid); + return () => { + clearTimeout(tid); + if (assetsRequestRef.current === request) assetsRequestRef.current++; + }; }, [tab, page, size, query, refreshKey]); + React.useEffect(() => { + if (tab === "company") return; + const lastPage = Math.max(0, Math.ceil(total / size) - 1); + if (page > lastPage) setPage(lastPage); + }, [page, size, tab, total]); + const companyById = React.useMemo(() => { const m = new Map(); for (const c of companies) m.set(c.id, c.name); @@ -374,7 +549,8 @@ export default function AssetsPage() { const totalAssets = Object.values(counts).reduce((a, b) => a + b, 0); // rows are already the current server-side page. - const tabData = (_type: string): Asset[] => rows; + const tabData = (type: string): Asset[] => (rowsTab === type ? rows : []); + const currentRows = tabData(tab); const slice = (list: T[]) => list; const searchBox = ( @@ -398,11 +574,7 @@ export default function AssetsPage() { 共 {totalAssets} 项资产 {selected.size > 0 && ( - )} @@ -453,7 +625,9 @@ export default function AssetsPage() { {c.scope?.length ? (
{c.scope.map((s, i) => ( - {s.raw} + + {s.raw} + ))}
) : ( @@ -468,7 +642,11 @@ export default function AssetsPage() { variant="ghost" size="icon" className="size-7 text-muted-foreground hover:text-destructive" - onClick={() => { setCompanyDeleteTarget(c); setCompanyDeleteAssets(false); }} + onClick={() => { + setCompanyDeleteTarget(c); + setCompanyDeleteAssets(false); + }} + aria-label={`删除企业 ${c.name}`} > @@ -500,7 +678,7 @@ export default function AssetsPage() { size={size} onSize={setSize} onPage={setPage} - rows={rows} + rows={currentRows} selected={selected} onToggleAll={(ids) => toggleSelectAll(ids)} > @@ -520,6 +698,7 @@ export default function AssetsPage() { size="icon" className="size-7 text-muted-foreground hover:text-destructive" onClick={() => openDelete([a.id])} + aria-label={`删除资产 ${a.domain || a.id}`} > @@ -540,7 +719,7 @@ export default function AssetsPage() { size={size} onSize={setSize} onPage={setPage} - rows={rows} + rows={currentRows} selected={selected} onToggleAll={(ids) => toggleSelectAll(ids)} > @@ -551,12 +730,12 @@ export default function AssetsPage() { {a.ip} {a.c_segment || "—"} - + + + - p.service ? `${p.port}/${p.service}` : String(p.port), - )} + items={(a.open_ports ?? []).map((p) => (p.service ? `${p.port}/${p.service}` : String(p.port)))} mono /> @@ -566,6 +745,7 @@ export default function AssetsPage() { size="icon" className="size-7 text-muted-foreground hover:text-destructive" onClick={() => openDelete([a.id])} + aria-label={`删除资产 ${a.ip || a.id}`} > @@ -586,7 +766,7 @@ export default function AssetsPage() { size={size} onSize={setSize} onPage={setPage} - rows={rows} + rows={currentRows} selected={selected} onToggleAll={(ids) => toggleSelectAll(ids)} > @@ -598,13 +778,16 @@ export default function AssetsPage() { {a.domain} {a.root_domain || "—"} {a.record_type || "—"} - {(Array.isArray(a.record_value) ? a.record_value.join(", ") : a.record_value) || "—"} + + {(Array.isArray(a.record_value) ? a.record_value.join(", ") : a.record_value) || "—"} + @@ -625,7 +808,7 @@ export default function AssetsPage() { size={size} onSize={setSize} onPage={setPage} - rows={rows} + rows={currentRows} selected={selected} onToggleAll={(ids) => toggleSelectAll(ids)} > @@ -644,6 +827,7 @@ export default function AssetsPage() { size="icon" className="size-7 text-muted-foreground hover:text-destructive" onClick={() => openDelete([a.id])} + aria-label={`删除资产 ${a.app_name || a.id}`} > @@ -664,7 +848,7 @@ export default function AssetsPage() { size={size} onSize={setSize} onPage={setPage} - rows={rows} + rows={currentRows} selected={selected} onToggleAll={(ids) => toggleSelectAll(ids)} > @@ -706,10 +890,16 @@ export default function AssetsPage() { {a.status_code} - ) : "—"} + ) : ( + "—" + )} + + + {a.page_title || "—"} + + + - {a.page_title || "—"} - {(a.auth ?? []).length === 0 ? ( — @@ -732,6 +922,7 @@ export default function AssetsPage() { size="icon" className="size-7 text-muted-foreground hover:text-destructive" onClick={() => openDelete([a.id])} + aria-label={`删除资产 ${a.url || a.id}`} > @@ -753,7 +944,7 @@ export default function AssetsPage() { size={size} onSize={setSize} onPage={setPage} - rows={rows} + rows={currentRows} selected={selected} onToggleAll={(ids) => toggleSelectAll(ids)} > @@ -762,8 +953,12 @@ export default function AssetsPage() { toggleSelect(a.id)} /> - - {a.url || "—"} + + + + + {a.url || "—"} + { @@ -779,6 +974,7 @@ export default function AssetsPage() { size="icon" className="size-7 text-muted-foreground hover:text-destructive" onClick={() => openDelete([a.id])} + aria-label={`删除资产 ${a.url || a.id}`} > @@ -794,13 +990,17 @@ export default function AssetsPage() { 确认删除 - 将永久删除 {deleteIds.length} 条资产记录,此操作不可撤销。 + 将永久删除 {deleteIds.length}{" "} + 条资产记录,此操作不可撤销。 取消 { e.preventDefault(); confirmDelete(); }} + onClick={(e) => { + e.preventDefault(); + void confirmDelete(); + }} disabled={deleting} className="bg-destructive text-destructive-foreground hover:bg-destructive/90" > @@ -810,14 +1010,25 @@ export default function AssetsPage() { - { if (!o) { setCompanyDeleteTarget(null); setCompanyDeleteAssets(false); } }}> + { + if (!o) { + setCompanyDeleteTarget(null); + setCompanyDeleteAssets(false); + } + }} + > 删除企业 · {companyDeleteTarget?.name}

此操作将永久删除该企业及其资产范围配置,不可撤销。

-
@@ -836,7 +1045,10 @@ export default function AssetsPage() { 取消 { e.preventDefault(); confirmDeleteCompany(); }} + onClick={(e) => { + e.preventDefault(); + void confirmDeleteCompany(); + }} disabled={companyDeleting} className="bg-destructive text-destructive-foreground hover:bg-destructive/90" > @@ -910,14 +1122,16 @@ function TabSearchBox({ setSelIdx((i) => Math.max(i - 1, 0)); } else if (e.key === "Tab" || e.key === "Enter") { const s = suggestions[selIdx]; - if (s) { e.preventDefault(); apply(s); } + if (s) { + e.preventDefault(); + apply(s); + } } else if (e.key === "Escape") { setOpen(false); } }; - const cursorPos = () => - inputRef.current?.selectionStart ?? query.length; + const cursorPos = () => inputRef.current?.selectionStart ?? query.length; return (
@@ -937,31 +1151,29 @@ function TabSearchBox({ {open && suggestions.length > 0 && (
{suggestions.map((s, i) => ( -
setSelIdx(i)} - onMouseDown={(e) => { e.preventDefault(); apply(s); }} + onMouseDown={(e) => { + e.preventDefault(); + apply(s); + }} > - - {s.label} - + {s.label} {s.desc} -
+ ))}
)}
{query.trim() && !open && (

- {loading - ? "搜索中…" - : error - ? {error} - : `找到 ${count ?? 0} 条`} + {loading ? "搜索中…" : error ? {error} : `找到 ${count ?? 0} 条`}

)} @@ -1018,12 +1230,14 @@ function AssetCard({ ) : ( {c} - ) + ), )} - {childRows.length > 0 ? childRows : ( + {childRows.length > 0 ? ( + childRows + ) : ( {loaded ? "暂无数据。" : "加载中…"} @@ -1036,19 +1250,43 @@ function AssetCard({ {total > 0 && (
- {start}–{end} / {total} + + {start}–{end} / {total} + {pageCount > 1 && (
- - {page + 1} / {pageCount} -
@@ -1065,7 +1303,9 @@ function Chips({ items, mono }: { items: string[]; mono?: boolean }) { return (
{clean.map((s, i) => ( - {s} + + {s} + ))}
); @@ -1081,26 +1321,44 @@ function CompanyAvatar({ name, logo }: { name: string; logo?: string }) { ); } -// 新增企业弹窗 +function savedScopeRules(company: Company): CompanyScopeRule[] { + return (company.scope ?? []) + .map((scope) => ({ kind: scope.kind, value: scope.raw.trim() })) + .filter((scope) => scope.value); +} + +function savedScopeText(company: Company): string { + return savedScopeRules(company) + .map((scope) => scope.value) + .join("\n"); +} + +// 新增企业使用与任务、LLM 编辑一致的右侧抽屉。 function CompanyDialog({ onSaved }: { onSaved: () => void }) { const [open, setOpen] = React.useState(false); const [name, setName] = React.useState(""); - const [logo, setLogo] = React.useState(""); - const [scope, setScope] = React.useState(""); + const [scopeText, setScopeText] = React.useState(""); const [busy, setBusy] = React.useState(false); + const parsedScope = React.useMemo(() => parseCompanyScopeText(scopeText), [scopeText]); React.useEffect(() => { if (!open) return; setName(""); - setLogo(""); - setScope(""); + setScopeText(""); }, [open]); const submit = async () => { - if (!name.trim()) { toast.error("请填写企业名称"); return; } + if (!name.trim()) { + toast.error("请填写企业名称"); + return; + } + if (parsedScope.errors.length > 0) { + toast.error("请修正无效的资产范围"); + return; + } setBusy(true); try { - const res = await api.createCompany(name.trim(), logo.trim(), scope); + const res = await api.createCompany(name.trim(), parsedScope.rules); const added = res.scope_added ?? 0; const invalid = res.scope_invalid ?? 0; if (invalid > 0) toast.warning(`已创建企业,添加 ${added} 条范围;${invalid} 行无效`); @@ -1110,86 +1368,84 @@ function CompanyDialog({ onSaved }: { onSaved: () => void }) { } catch (e) { const msg = String((e as Error)?.message ?? e); if (/:\s*409$/.test(msg)) toast.error("企业已存在,请换个名称"); - else toast.error("保存失败:" + msg); + else toast.error(`保存失败:${msg}`); } finally { setBusy(false); } }; return ( - - - - - - - 新增企业 - - 资产范围是归属的唯一来源:命中范围的资产(现有 + 未来)会被自动归到该企业。一行一条:根域名(含全部子域)、IP、或 CIDR 网段。 - - -
-
- - setName(e.target.value)} /> -
-
- -
+ + + + + + + 新增企业 + 配置企业及其资产范围。关键词只作为 Agent 提示,不会自动归属资产。 + +
+ + + 企业名称 setLogo(e.target.value)} + id="cn-name" + placeholder="如 Acme Corp(名称唯一)" + value={name} + onChange={(e) => setName(e.target.value)} /> - -
-
-
- -