Merge tag 'v0.3.4'
This commit is contained in:
commit
48e85b33d6
33 files changed
+1456
-120
No files matched your search
@@ -3,6 +3,7 @@ package db
|
||||
import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/url"
|
||||
@@ -300,6 +301,27 @@ RETURNING id`, domain, icpVal, companyID, taskIDs).Scan(&id)
|
||||
return id, err
|
||||
}
|
||||
|
||||
// ErrAssetIPInvalid marks a non-address value in an asset's ip field. Both
|
||||
// insert_assets and the asset API report it per item with the item index, so one
|
||||
// bad entry never costs the rest of the batch.
|
||||
var ErrAssetIPInvalid = errors.New("invalid asset ip")
|
||||
|
||||
// ValidateAssetIP keeps hostnames out of assets.ip. Network attribution casts
|
||||
// that column to inet (see try_inet in schema.sql), so a hostname stored here is
|
||||
// silently invisible to every IP/CIDR scope rule — the asset simply never gets
|
||||
// attributed and nobody can tell why. The message states the fix rather than
|
||||
// just the fault, so an agent can correct the item on its next turn. An empty
|
||||
// value is accepted: the ip field is optional for service and endpoint assets.
|
||||
func ValidateAssetIP(value string) error {
|
||||
if value == "" || net.ParseIP(value) != nil {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf(
|
||||
"%w: ip 必须是 IPv4/IPv6 地址,收到 %q。若这是主机名,请改用 type=subdomain 并填 domain 字段;"+
|
||||
"若确实要登记地址,请先解析出 A/AAAA 记录,再用解析出的地址填 ip",
|
||||
ErrAssetIPInvalid, value)
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// UpsertIP
|
||||
// =====================================================================
|
||||
@@ -317,6 +339,9 @@ func (s *AssetStore) UpsertIP(req UpsertIPReq) (int64, error) {
|
||||
if req.IP == "" {
|
||||
return 0, fmt.Errorf("ip is required")
|
||||
}
|
||||
if err := ValidateAssetIP(req.IP); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if s.tx == nil {
|
||||
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
|
||||
return scoped.UpsertIP(req)
|
||||
@@ -659,6 +684,9 @@ func (s *AssetStore) UpsertHTTPService(req UpsertHTTPServiceReq) (int64, error)
|
||||
if req.URL == "" {
|
||||
return 0, fmt.Errorf("url is required")
|
||||
}
|
||||
if err := ValidateAssetIP(req.IP); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if s.tx == nil {
|
||||
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
|
||||
return scoped.UpsertHTTPService(req)
|
||||
@@ -793,6 +821,9 @@ func (s *AssetStore) UpsertOtherService(req UpsertOtherServiceReq) (int64, error
|
||||
if req.Domain == "" && req.IP == "" {
|
||||
return 0, fmt.Errorf("domain or ip is required")
|
||||
}
|
||||
if err := ValidateAssetIP(req.IP); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if req.Port == 0 {
|
||||
return 0, fmt.Errorf("port is required")
|
||||
}
|
||||
@@ -936,6 +967,9 @@ func (s *AssetStore) UpsertEndpoint(req UpsertEndpointReq) (int64, error) {
|
||||
if req.Method == "" {
|
||||
return 0, fmt.Errorf("method is required")
|
||||
}
|
||||
if err := ValidateAssetIP(req.IP); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if s.tx == nil {
|
||||
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
|
||||
return scoped.UpsertEndpoint(req)
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
@@ -946,3 +948,54 @@ func TestAssetPaginationUsesStableIDTieBreaker(t *testing.T) {
|
||||
t.Fatalf("DSL pagination order=%v want=%v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAssetIPRejectsHostname pins the write-side guard added after a hostname in
|
||||
// assets.ip broke company attribution: every asset type that persists an ip must
|
||||
// refuse a non-address, and the message must tell the caller how to fix it.
|
||||
func TestAssetIPRejectsHostname(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
hostname := "d63cd476.cdn.ucloud.com.cn"
|
||||
cases := []struct {
|
||||
name string
|
||||
call func() (int64, error)
|
||||
}{
|
||||
{"ip", func() (int64, error) {
|
||||
return av2.UpsertIP(UpsertIPReq{IP: hostname})
|
||||
}},
|
||||
{"http service", func() (int64, error) {
|
||||
return av2.UpsertHTTPService(UpsertHTTPServiceReq{URL: "https://example.com/", IP: hostname})
|
||||
}},
|
||||
{"other service", func() (int64, error) {
|
||||
return av2.UpsertOtherService(UpsertOtherServiceReq{IP: hostname, Port: 22, ServiceName: "ssh"})
|
||||
}},
|
||||
{"endpoint", func() (int64, error) {
|
||||
return av2.UpsertEndpoint(UpsertEndpointReq{URL: "https://example.com/a", Method: "GET", IP: hostname})
|
||||
}},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
id, err := tc.call()
|
||||
if !errors.Is(err, ErrAssetIPInvalid) {
|
||||
deleteAsset(d, id)
|
||||
t.Fatalf("%s: err=%v, want %v", tc.name, err, ErrAssetIPInvalid)
|
||||
}
|
||||
// The caller has to learn the remedy, not just that something was wrong.
|
||||
if !strings.Contains(err.Error(), hostname) || !strings.Contains(err.Error(), "type=subdomain") {
|
||||
t.Fatalf("%s: message is not actionable: %v", tc.name, err)
|
||||
}
|
||||
}
|
||||
|
||||
// An empty ip stays legal for the types where it is optional, and a real
|
||||
// address is still accepted.
|
||||
id, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{URL: "https://ip-guard.example.com/"})
|
||||
if err != nil {
|
||||
t.Fatalf("empty ip rejected: %v", err)
|
||||
}
|
||||
deleteAsset(d, id)
|
||||
id, err = av2.UpsertIP(UpsertIPReq{IP: "203.0.113.7"})
|
||||
if err != nil {
|
||||
t.Fatalf("valid ip rejected: %v", err)
|
||||
}
|
||||
deleteAsset(d, id)
|
||||
}
|
||||
+98
-12
@@ -4,6 +4,7 @@ import (
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"net"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
@@ -161,9 +162,11 @@ RETURNING id`, name, nkey, logoVal).Scan(&id); err != nil {
|
||||
return 0, 0, 0, invalid, validationErrors, err
|
||||
}
|
||||
if needsAttribution {
|
||||
if err := recomputeAttributionTx(tx); err != nil {
|
||||
warning, err := recomputeAttributionTx(tx)
|
||||
if err != nil {
|
||||
return 0, 0, 0, invalid, validationErrors, err
|
||||
}
|
||||
logAttributionWarning(warning)
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return 0, 0, 0, invalid, validationErrors, err
|
||||
@@ -244,9 +247,13 @@ func (s *CompanyStore) DeleteCompanyWithAssets(id int64, deleteAssets bool) (ass
|
||||
if deleted == 0 {
|
||||
return 0, ErrCompanyNotFound
|
||||
}
|
||||
if err := recomputeAttributionTx(tx); err != nil {
|
||||
// This path has no per-request warning channel, so the log is the only place
|
||||
// the operator can learn about unparseable ip rows here.
|
||||
warning, err := recomputeAttributionTx(tx)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
logAttributionWarning(warning)
|
||||
if err := tx.Commit(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
@@ -370,9 +377,11 @@ func (s *CompanyStore) AddScopeInputsChecked(companyID int64, inputs []ScopeInpu
|
||||
return 0, 0, invalid, errors, err
|
||||
}
|
||||
if needsAttribution {
|
||||
if err := recomputeAttributionTx(tx); err != nil {
|
||||
warning, err := recomputeAttributionTx(tx)
|
||||
if err != nil {
|
||||
return 0, 0, invalid, errors, fmt.Errorf("重新计算企业归属失败: %w", err)
|
||||
}
|
||||
logAttributionWarning(warning)
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return 0, 0, invalid, errors, err
|
||||
@@ -502,20 +511,26 @@ func (s *CompanyStore) RecomputeAttribution() error {
|
||||
if err := lockCompanyScopeMutation(tx); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := recomputeAttributionTx(tx); err != nil {
|
||||
warning, err := recomputeAttributionTx(tx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
logAttributionWarning(warning)
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
func recomputeAttributionTx(tx *sql.Tx) error {
|
||||
// recomputeAttributionTx rebuilds scope-derived ownership. It returns a warning
|
||||
// for assets whose ip column cannot be parsed: try_inet skips them instead of
|
||||
// aborting the statement, so without this they would silently never receive a
|
||||
// network-based company. Callers surface the warning and it is always logged.
|
||||
func recomputeAttributionTx(tx *sql.Tx) (string, 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
|
||||
SET company_id = NULL, company_source = 'scope'
|
||||
WHERE company_source = 'scope'`); err != nil {
|
||||
return err
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Domain-based attribution (root_domain exact match).
|
||||
@@ -533,7 +548,7 @@ 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 "", err
|
||||
}
|
||||
|
||||
// IP/CIDR attribution for still-unowned assets.
|
||||
@@ -541,7 +556,7 @@ WHERE a.id = matched.asset_id`); 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 IN ('ip','cidr') AND cs.net >>= a.ip::inet
|
||||
JOIN company_scope cs ON cs.kind IN ('ip','cidr') AND cs.net >>= try_inet(a.ip)
|
||||
WHERE a.company_id IS NULL
|
||||
AND a.type IN ('ip','subdomain','service','endpoint')
|
||||
AND a.ip IS NOT NULL
|
||||
@@ -551,7 +566,7 @@ 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 "", err
|
||||
}
|
||||
|
||||
// Exact normalized ICP attribution after domain/network precedence.
|
||||
@@ -572,9 +587,78 @@ 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 "", err
|
||||
}
|
||||
return nil
|
||||
return malformedIPAssetWarning(tx)
|
||||
}
|
||||
|
||||
// malformedIPAssetsSampled bounds how many offending ids one warning names, so a
|
||||
// large batch of bad rows stays readable in a toast and in the log.
|
||||
const malformedIPAssetsSampled = 5
|
||||
|
||||
// logAttributionWarning records a recompute warning in the server log. Every
|
||||
// recompute path calls it, so the warning is reported even for triggers with no
|
||||
// per-request response (company deletion, scopesentry sync, agent asset writes).
|
||||
func logAttributionWarning(warning string) {
|
||||
if warning != "" {
|
||||
log.Printf("[assets] %s", warning)
|
||||
}
|
||||
}
|
||||
|
||||
// malformedIPAssetQueryer is satisfied by both *sql.Tx and *DB so the warning
|
||||
// can be produced inside a recompute transaction or read standalone by the API.
|
||||
type malformedIPAssetQueryer interface {
|
||||
Query(query string, args ...any) (*sql.Rows, error)
|
||||
}
|
||||
|
||||
// MalformedIPAssetWarning reports assets with an unparseable ip outside of any
|
||||
// mutation, letting the API attach the warning to a scope response without
|
||||
// widening the mutation signatures — an unrelated data problem is not one of
|
||||
// this request's validation errors.
|
||||
func (s *CompanyStore) MalformedIPAssetWarning() (string, error) {
|
||||
return malformedIPAssetWarning(s.db)
|
||||
}
|
||||
|
||||
// malformedIPAssetWarning describes assets whose ip column is not a valid
|
||||
// address. They are invisible to network attribution, so the operator has to be
|
||||
// told which rows to fix — silently skipping them would look like scope rules
|
||||
// that simply do not work.
|
||||
func malformedIPAssetWarning(q malformedIPAssetQueryer) (string, error) {
|
||||
rows, err := q.Query(`
|
||||
SELECT id, ip, count(*) OVER () AS total
|
||||
FROM assets
|
||||
WHERE ip IS NOT NULL AND ip <> '' AND try_inet(ip) IS NULL
|
||||
AND type IN ('ip','subdomain','service','endpoint')
|
||||
ORDER BY id
|
||||
LIMIT $1`, malformedIPAssetsSampled)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer rows.Close()
|
||||
var total int
|
||||
samples := make([]string, 0, malformedIPAssetsSampled)
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
var ip string
|
||||
if err := rows.Scan(&id, &ip, &total); err != nil {
|
||||
return "", err
|
||||
}
|
||||
samples = append(samples, fmt.Sprintf("#%d %s", id, ip))
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if total == 0 {
|
||||
return "", nil
|
||||
}
|
||||
warning := fmt.Sprintf(
|
||||
"%d 条资产的 ip 字段不是合法 IP,已跳过 IP/CIDR 范围匹配(这些资产不会被网段规则归属到企业):%s",
|
||||
total, strings.Join(samples, "、"),
|
||||
)
|
||||
if total > len(samples) {
|
||||
warning += fmt.Sprintf(" 等 %d 条", total)
|
||||
}
|
||||
return warning, nil
|
||||
}
|
||||
|
||||
// UpdateScope replaces all scope rules for a company and reattributes.
|
||||
@@ -635,9 +719,11 @@ func (s *CompanyStore) UpdateScopeInputsChecked(companyID int64, inputs []ScopeI
|
||||
}
|
||||
// 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 {
|
||||
warning, err := recomputeAttributionTx(tx)
|
||||
if err != nil {
|
||||
return 0, invalid, errs, fmt.Errorf("重新计算企业归属失败: %w", err)
|
||||
}
|
||||
logAttributionWarning(warning)
|
||||
if err := tx.Commit(); err != nil {
|
||||
return 0, invalid, errs, err
|
||||
}
|
||||
|
||||
@@ -857,6 +857,16 @@ WHERE exploration_id=$1 AND kind='intent' AND state IN ('open','running'))`, s.e
|
||||
return exists, err
|
||||
}
|
||||
|
||||
// HasOpenGoal reports whether this exploration still has any goal in state 'open'.
|
||||
// false ⇒ 所有目标已 met/abandoned(或本任务无目标)⇒ 进入 goalless(人工直投)分支:
|
||||
// planner 停跑,任务是否结束改由 frontier 是否抽干决定。met 与 abandoned 都算"已了结"。
|
||||
func (s *ExplorationStore) HasOpenGoal() (bool, error) {
|
||||
var exists bool
|
||||
err := s.db.QueryRow(`SELECT EXISTS(SELECT 1 FROM exploration_nodes
|
||||
WHERE exploration_id=$1 AND kind='goal' AND state='open')`, s.expID).Scan(&exists)
|
||||
return exists, err
|
||||
}
|
||||
|
||||
// ClaimIntent atomically moves an open intent to running. Returns true if claimed.
|
||||
func (s *ExplorationStore) ClaimIntent(id int64, owner string) (bool, error) {
|
||||
res, err := s.db.Exec(`UPDATE exploration_nodes SET state='running', owner=$1
|
||||
|
||||
@@ -8,6 +8,18 @@ CREATE OR REPLACE FUNCTION set_updated_at() RETURNS trigger AS $$
|
||||
BEGIN NEW.updated_at = now(); RETURN NEW; END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
|
||||
-- 安全的 text→inet 转换:非法值返回 NULL 而不是抛 22P02。assets.ip 是自由文本
|
||||
-- (Agent / 资产 API 可能写进主机名),裸转 a.ip::inet 会让单独一行脏数据把整条
|
||||
-- 企业归属重算语句打挂。调用方用 try_inet(...) IS NULL 找出这些行并告警。
|
||||
-- 不用 pg_input_is_valid 是因为那要 PG16+,这里要兼容更老的存量库。
|
||||
CREATE OR REPLACE FUNCTION try_inet(value text) RETURNS inet AS $$
|
||||
BEGIN
|
||||
RETURN value::inet;
|
||||
EXCEPTION WHEN others THEN
|
||||
RETURN NULL;
|
||||
END;
|
||||
$$ LANGUAGE plpgsql IMMUTABLE STRICT;
|
||||
|
||||
-- =====================================================================
|
||||
-- A. 资产层:companies / assets / company_scope
|
||||
-- =====================================================================
|
||||
|
||||
@@ -288,6 +288,10 @@ FROM assets WHERE id=ANY($2::bigint[])`, taskID, assetIDs).Scan(&found, &existin
|
||||
if found != len(assetIDs) {
|
||||
return mutation, ErrTaskAssetAssetNotFound
|
||||
}
|
||||
// Order matters: this UPDATE fires trg_assets_task_links, which creates the
|
||||
// link rows with the generic source='system'. The INSERT below must stay
|
||||
// after it so the operator-authored 'manual' provenance wins; swapping the
|
||||
// two statements silently degrades every manual attach back to 'system'.
|
||||
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
|
||||
|
||||
+12
-12
@@ -29,12 +29,12 @@ target AS (
|
||||
JOIN task_scope ts ON (
|
||||
(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='icp' AND (
|
||||
lower(regexp_replace(COALESCE(a.icp,''), '[[:space:]]+', '', 'g')) = ts.value
|
||||
OR lower(regexp_replace(COALESCE(a.app_icp,''), '[[:space:]]+', '', 'g')) = ts.value
|
||||
))
|
||||
OR (ts.kind='subdomain' AND a.domain = ts.domain)
|
||||
OR (ts.kind IN ('ip','cidr') AND ts.net >>= try_inet(a.ip))
|
||||
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
|
||||
))
|
||||
)
|
||||
JOIN context_tasks ctx ON ctx.task_id=ts.task_id
|
||||
UNION
|
||||
@@ -179,12 +179,12 @@ context_assets AS (
|
||||
JOIN task_scope ts ON (
|
||||
(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='icp' AND (
|
||||
lower(regexp_replace(COALESCE(a.icp,''), '[[:space:]]+', '', 'g')) = ts.value
|
||||
OR lower(regexp_replace(COALESCE(a.app_icp,''), '[[:space:]]+', '', 'g')) = ts.value
|
||||
))
|
||||
OR (ts.kind='subdomain' AND a.domain=ts.domain)
|
||||
OR (ts.kind IN ('ip','cidr') AND ts.net >>= try_inet(a.ip))
|
||||
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
|
||||
))
|
||||
)
|
||||
JOIN context_tasks ctx ON ctx.task_id=ts.task_id
|
||||
)
|
||||
|
||||
@@ -13,6 +13,10 @@ import (
|
||||
|
||||
const MaxTaskCategoryNameRunes = 80
|
||||
|
||||
// MaxTaskCategoryBatchSize bounds one batch move so a single request cannot lock
|
||||
// an unbounded number of task rows.
|
||||
const MaxTaskCategoryBatchSize = 100
|
||||
|
||||
var (
|
||||
ErrTaskCategoryInvalid = errors.New("invalid task category")
|
||||
ErrTaskCategoryNameConflict = errors.New("task category name already exists")
|
||||
@@ -199,3 +203,80 @@ FROM selected, updated`, taskID, *categoryID))
|
||||
}
|
||||
return &category, nil
|
||||
}
|
||||
|
||||
// SetTasksCategory moves several tasks into one category (nil = uncategorized)
|
||||
// inside a single transaction, so a half-applied batch is never observable.
|
||||
// It returns the ids that were actually updated — ids missing from that slice
|
||||
// were deleted between selection and submit — plus the refreshed category row
|
||||
// whose task_count already reflects this move.
|
||||
func (d *DB) SetTasksCategory(taskIDs []int64, categoryID *int64) ([]int64, *TaskCategory, error) {
|
||||
if len(taskIDs) == 0 {
|
||||
return nil, nil, fmt.Errorf("%w: task ids are required", ErrTaskCategoryInvalid)
|
||||
}
|
||||
if len(taskIDs) > MaxTaskCategoryBatchSize {
|
||||
return nil, nil, fmt.Errorf("%w: at most %d tasks per request", ErrTaskCategoryInvalid, MaxTaskCategoryBatchSize)
|
||||
}
|
||||
for _, id := range taskIDs {
|
||||
if id <= 0 {
|
||||
return nil, nil, fmt.Errorf("%w: task id must be positive", ErrTaskCategoryInvalid)
|
||||
}
|
||||
}
|
||||
if categoryID != nil && *categoryID <= 0 {
|
||||
return nil, nil, fmt.Errorf("%w: category id must be positive", ErrTaskCategoryInvalid)
|
||||
}
|
||||
tx, err := d.Begin()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
defer tx.Rollback() //nolint:errcheck
|
||||
|
||||
// Checking the category inside the transaction keeps a concurrent delete from
|
||||
// turning the UPDATE below into a foreign key violation.
|
||||
if categoryID != nil {
|
||||
var exists bool
|
||||
if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM task_categories WHERE id=$1)`, *categoryID).Scan(&exists); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if !exists {
|
||||
return nil, nil, ErrTaskCategoryNotFound
|
||||
}
|
||||
}
|
||||
rows, err := tx.Query(`
|
||||
UPDATE tasks SET category_id=$2
|
||||
WHERE id=ANY($1::bigint[]) AND deleted_at IS NULL
|
||||
RETURNING id`, taskIDs, categoryID)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
updated := make([]int64, 0, len(taskIDs))
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
rows.Close()
|
||||
return nil, nil, err
|
||||
}
|
||||
updated = append(updated, id)
|
||||
}
|
||||
rows.Close()
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
var category *TaskCategory
|
||||
if categoryID != nil {
|
||||
fetched, err := scanTaskCategory(tx.QueryRow(`
|
||||
SELECT `+taskCategoryCols+`
|
||||
FROM task_categories category
|
||||
LEFT JOIN tasks task ON task.category_id=category.id
|
||||
WHERE category.id=$1
|
||||
GROUP BY category.id`, *categoryID))
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
category = &fetched
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return updated, category, nil
|
||||
}
|
||||
@@ -89,3 +89,95 @@ func TestCreateTaskRejectsMissingCategoryAtomically(t *testing.T) {
|
||||
t.Fatalf("failed create leaked %d exploration rows", count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetTasksCategoryBatch(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
|
||||
stamp := time.Now().UnixNano()
|
||||
category, err := d.CreateTaskCategory(fmt.Sprintf("Batch %d", stamp))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _, _ = d.DeleteTaskCategory(category.ID) })
|
||||
|
||||
ids := make([]int64, 0, 3)
|
||||
for i := range 3 {
|
||||
task, err := d.CreateTaskWithOptions(fmt.Sprintf("batch-move-%d-%d", stamp, i), "goal", TaskCreateOptions{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
|
||||
ids = append(ids, task.ID)
|
||||
}
|
||||
|
||||
// A deleted id must not abort the move for the surviving tasks; it is simply
|
||||
// absent from the returned slice so the caller can report it.
|
||||
missing := int64(1 << 62)
|
||||
updated, moved, err := d.SetTasksCategory(append(append([]int64{}, ids...), missing), &category.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(updated) != len(ids) {
|
||||
t.Fatalf("updated=%v, want the %d live ids only", updated, len(ids))
|
||||
}
|
||||
if moved == nil || moved.TaskCount != len(ids) {
|
||||
t.Fatalf("category=%+v, want task_count=%d", moved, len(ids))
|
||||
}
|
||||
for _, id := range ids {
|
||||
loaded, err := d.GetTask(id)
|
||||
if err != nil || loaded.CategoryID == nil || *loaded.CategoryID != category.ID {
|
||||
t.Fatalf("task %d not moved: %+v err=%v", id, loaded, err)
|
||||
}
|
||||
}
|
||||
|
||||
// A nil category clears the assignment for the whole batch.
|
||||
updated, moved, err = d.SetTasksCategory(ids, nil)
|
||||
if err != nil || len(updated) != len(ids) || moved != nil {
|
||||
t.Fatalf("clear updated=%v category=%+v err=%v", updated, moved, err)
|
||||
}
|
||||
for _, id := range ids {
|
||||
loaded, _ := d.GetTask(id)
|
||||
if loaded.CategoryID != nil || loaded.CategoryName != "" {
|
||||
t.Fatalf("task %d still categorized: %+v", id, loaded)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetTasksCategoryRejectsBadInputAtomically(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
|
||||
stamp := time.Now().UnixNano()
|
||||
task, err := d.CreateTaskWithOptions(fmt.Sprintf("batch-reject-%d", stamp), "goal", TaskCreateOptions{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
|
||||
|
||||
missingCategory := int64(1 << 62)
|
||||
if _, _, err := d.SetTasksCategory([]int64{task.ID}, &missingCategory); !errors.Is(err, ErrTaskCategoryNotFound) {
|
||||
t.Fatalf("missing category error=%v, want %v", err, ErrTaskCategoryNotFound)
|
||||
}
|
||||
if _, _, err := d.SetTasksCategory(nil, nil); !errors.Is(err, ErrTaskCategoryInvalid) {
|
||||
t.Fatalf("empty ids error=%v, want %v", err, ErrTaskCategoryInvalid)
|
||||
}
|
||||
oversized := make([]int64, MaxTaskCategoryBatchSize+1)
|
||||
for i := range oversized {
|
||||
oversized[i] = task.ID
|
||||
}
|
||||
if _, _, err := d.SetTasksCategory(oversized, nil); !errors.Is(err, ErrTaskCategoryInvalid) {
|
||||
t.Fatalf("oversized error=%v, want %v", err, ErrTaskCategoryInvalid)
|
||||
}
|
||||
// The rejected calls must leave the task untouched.
|
||||
loaded, _ := d.GetTask(task.ID)
|
||||
if loaded.CategoryID != nil {
|
||||
t.Fatalf("rejected batch mutated task: %+v", loaded)
|
||||
}
|
||||
}
|
||||
+7
-7
@@ -308,13 +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='icp' AND (
|
||||
lower(regexp_replace(COALESCE(a.icp,''), '[[:space:]]+', '', 'g')) = ts.value
|
||||
OR lower(regexp_replace(COALESCE(a.app_icp,''), '[[:space:]]+', '', 'g')) = ts.value
|
||||
))
|
||||
)
|
||||
OR (ts.kind='subdomain' AND a.domain = ts.domain)
|
||||
OR (ts.kind IN ('ip','cidr') AND ts.net >>= try_inet(a.ip))
|
||||
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
|
||||
|
||||
+3
-1
@@ -322,7 +322,9 @@ func (d *DB) SetParentRef(id int64, parentRef string) error {
|
||||
|
||||
// ListTasks returns alive tasks, newest first.
|
||||
func (d *DB) ListTasks() ([]*Task, error) {
|
||||
rows, err := d.Query(`SELECT ` + taskCols + ` FROM tasks WHERE deleted_at IS NULL ORDER BY created_at DESC`)
|
||||
// 按 id 降序:id 是 BIGSERIAL,单调递增即插入顺序,最新的在最前。不用 created_at
|
||||
// 是因为同一时刻创建的多个任务时间戳相同,排序不稳定,列表会频繁换位(见 List)。
|
||||
rows, err := d.Query(`SELECT ` + taskCols + ` FROM tasks WHERE deleted_at IS NULL ORDER BY id DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -563,3 +563,45 @@ func TestCreateTaskWithCompanyScopes(t *testing.T) {
|
||||
t.Fatalf("failed company association leaked %d exploration rows", leaked)
|
||||
}
|
||||
}
|
||||
|
||||
// TestListTasksOrderByIDDesc pins list ordering to id-descending. created_at is
|
||||
// deliberately not the sort key: tasks created in the same instant share a
|
||||
// timestamp and would reorder between polls; id is unique and monotonic.
|
||||
func TestListTasksOrderByIDDesc(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
|
||||
stamp := time.Now().UnixNano()
|
||||
var ids []int64
|
||||
for i := range 3 {
|
||||
tk, err := d.CreateTask(fmt.Sprintf("order-%d-%d", stamp, i), "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(tk.ID) })
|
||||
ids = append(ids, tk.ID)
|
||||
}
|
||||
|
||||
list, err := d.ListTasks()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Reduce to just the ids created here; other rows may exist in the shared DB.
|
||||
mine := map[int64]bool{ids[0]: true, ids[1]: true, ids[2]: true}
|
||||
var seen []int64
|
||||
for _, task := range list {
|
||||
if mine[task.ID] {
|
||||
seen = append(seen, task.ID)
|
||||
}
|
||||
}
|
||||
if len(seen) != 3 {
|
||||
t.Fatalf("found %d of the created tasks in the list, want 3", len(seen))
|
||||
}
|
||||
// Newest (largest id) first.
|
||||
if seen[0] != ids[2] || seen[1] != ids[1] || seen[2] != ids[0] {
|
||||
t.Fatalf("order=%v, want descending %v", seen, []int64{ids[2], ids[1], ids[0]})
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user