feat(notify): 漏洞发现支持推送到钉钉/飞书/企业微信等六种渠道
新增「系统 → 通知推送」页，把扫描发现的漏洞推送到钉钉、飞书、企业微信、 通用 Webhook、Telegram、邮件。同一类型可配任意多个机器人实例，每个实例独立 设启停、限流与过滤规则（最低级别 / 任务资产范围 / 漏洞类型关键词 / 是否接收 处置状态变更）。推送时机分实时与汇总两种模式，「高危实时、其余汇总」建两个 渠道即可表达，策略不写死在代码里。带投递历史与失败重发。 三个关键设计： 1. 写漏洞的事务只做一次盲 INSERT。notification_events 由 RecordFindingTx 在同一事务内写入，保证「漏洞落库」与「推送任务存在」原子一致；这条 INSERT 刻意不读渠道表、不跑用户的过滤规则，否则一条配错的过滤条件就能污染甚至 中止事务，让高危漏洞存不进库。为把失败隔离在这一条语句上（PostgreSQL 中 事务内任一语句报错会让整个事务作废、连 COMMIT 都失败），它被 SAVEPOINT 包住，失败时只记日志。 2. 投递用租约领取而非长事务。FOR UPDATE SKIP LOCKED 领取后把行置为 sending、 把 next_attempt_at 推到未来作为租约，提交后再做网络投递，因此投递期间不 持有数据库锁；崩溃留下的 sending 行会被下一轮重新领取，自愈且不会无限重试。 3. 限流不消耗重试预算。先按渠道令牌桶算出本轮还能发几条，再按这个数量领取； 顺序反过来会让被限流挡下的投递白计一次尝试，三次预算被等待耗光后落入失败。 审计并修复的四个问题： - 推送凭据经错误信息外泄：http.Client.Do 返回的 *url.Error 会把完整 URL 打进 错误文本，而钉钉 access_token、企业微信 key、飞书 hook id、Telegram /bot<token>/ 都就在 URL 里。该串流入 last_error（明文落库）、投递历史接口 （绕过渠道掩码）、服务端日志与前端错误提示。现统一脱敏为 scheme://host + 底层原因，并拒绝跨主机重定向。 - 最低级别门槛打错字会让过滤器静默失效：未知级别序数为 0，判定退化成 rank >= 0 恒真，用户以为限制「仅高危」实际全推。现于写入路径校验取值。 - 汇总批次无上限：一次领取把全部待发行读进内存、消息被渠道长度上限截掉大半 且静默丢失。现单批上限 500 条，超出部分留库成为下一批。 - 邮件渠道把 SMTP 4xx（灰名单等临时失败）判成永久失败，导致这类最该重试的 场景从不重试。现按应答码首位区分 4xx/5xx。 测试：notify 包 82% 覆盖（含用自建最小 SMTP 服务器驱动的完整邮件协议会话、 两种签名算法对 OpenSSL 基准值的比对、UTF-8 边界截断、以及模板能力边界的 反射断言）；db 与 server 侧共 29 个用例跑真实 PostgreSQL，其中端到端用例 真发 HTTP 到假接收端并断言消息内容与投递状态。 Co-Authored-By: Claude Code <[email protected]>
This commit is contained in:
luckoneandClaude Code committed
1 parent
18589b270e
commit
f3fb37f70d
41 files changed
+8740
-5
No files matched your search
+19
-2
@@ -10,6 +10,8 @@ import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/notify"
|
||||
)
|
||||
|
||||
var (
|
||||
@@ -335,7 +337,9 @@ type RecordedFinding struct {
|
||||
Traffic *FindingTraffic `json:"traffic"`
|
||||
}
|
||||
|
||||
func RecordFindingTx(tx *sql.Tx, in RecordFindingInput, prepared []PreparedTrafficEvidence) (*RecordedFinding, error) {
|
||||
// ctx 由调用方传入本次事务所用的上下文（而非在内部取 context.Background）：
|
||||
// 事务内新加的推送事件写入同样应受调用方的取消与超时约束。
|
||||
func RecordFindingTx(ctx context.Context, tx *sql.Tx, in RecordFindingInput, prepared []PreparedTrafficEvidence) (*RecordedFinding, error) {
|
||||
if err := LockTaskEvidenceTx(tx, in.TaskID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -382,6 +386,19 @@ VALUES($1,'finding',$2,9,'confirmed',$3) RETURNING id`, in.ExplorationID, string
|
||||
VALUES(NULLIF($1,0),$2,$3,$4,$5,$6,$7,$8,$9) RETURNING id`, in.TaskID, out.NodeID, in.VulnClass, in.Name, in.Severity, in.Summary, in.Evidence, in.Worker, string(raw)).Scan(&out.FindingID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// 在**同一事务**里登记一条推送事件：提交即保证「漏洞落库」与「推送任务存在」
|
||||
// 原子一致，不存在提交成功却没入队、消息永久丢失的窗口。
|
||||
// 这里的失败被隔离在保存点上、不影响漏洞写入（见函数注释），因此忽略返回值。
|
||||
RecordNotificationEventTx(ctx, tx, notify.EventFindingCreated, out.FindingID, notify.Snapshot{
|
||||
Kind: notify.EventFindingCreated,
|
||||
FindingID: out.FindingID,
|
||||
TaskID: in.TaskID,
|
||||
VulnClass: in.VulnClass,
|
||||
Name: in.Name,
|
||||
Severity: in.Severity,
|
||||
Summary: in.Summary,
|
||||
AssetIDs: assets,
|
||||
})
|
||||
if err := AddFindingTrafficTx(tx, out.FindingID, prepared); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -393,7 +410,7 @@ VALUES(NULLIF($1,0),$2,$3,$4,$5,$6,$7,$8,$9) RETURNING id`, in.TaskID, out.NodeI
|
||||
// RecordFinding is the atomic legacy/no-recorder path used by tools and tests.
|
||||
func (s *ExplorationStore) RecordFinding(ctx context.Context, in RecordFindingInput) (out *RecordedFinding, err error) {
|
||||
in.ExplorationID = s.expID
|
||||
err = s.db.WithEvidenceTx(ctx, func(tx *sql.Tx) error { var e error; out, e = RecordFindingTx(tx, in, nil); return e })
|
||||
err = s.db.WithEvidenceTx(ctx, func(tx *sql.Tx) error { var e error; out, e = RecordFindingTx(ctx, tx, in, nil); return e })
|
||||
return
|
||||
}
|
||||
|
||||
|
||||
@@ -690,6 +690,10 @@ func (d *DB) DeleteFindingsByTask(taskID int64) (int64, error) {
|
||||
}
|
||||
|
||||
// SetFindingStatus updates one finding's triage state. Returns rows affected.
|
||||
//
|
||||
// 底层 setter：只改状态、不登记推送事件。生产代码改状态请走
|
||||
// SetFindingStatusWithNotify —— 直接调本函数会让「状态变更推送」静默失效。
|
||||
// 保留它是为了让不关心通知的用例（参数校验、复测流程）能单独驱动状态。
|
||||
func (d *DB) SetFindingStatus(id int64, status string) (int64, error) {
|
||||
res, err := d.Exec(`UPDATE findings SET status=$1 WHERE id=$2`, status, id)
|
||||
if err != nil {
|
||||
|
||||
@@ -0,0 +1,528 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/notify"
|
||||
)
|
||||
|
||||
// 本文件是 IM 推送的渠道配置与事件层。投递任务的领取与状态流转见
|
||||
// db/notification_delivery.go。
|
||||
//
|
||||
// 两条不变量，改这个文件时务必保持：
|
||||
//
|
||||
// 1. 写漏洞的事务(RecordFindingTx)只调用 InsertNotificationEventTx 做一次盲插，
|
||||
// 不读任何通知相关的表、不做过滤匹配。任何在这里引入的读操作都可能因为
|
||||
// 用户配错的过滤条件而污染甚至中止漏洞写入事务。
|
||||
// 2. 过滤匹配永不报错：配置畸形一律按「命中」处理(见 notify.Match)。宁可多推，
|
||||
// 不可漏推。
|
||||
|
||||
// ErrNotificationChannelNotFound 渠道不存在。
|
||||
var ErrNotificationChannelNotFound = errors.New("通知渠道不存在")
|
||||
|
||||
// 投递状态。
|
||||
const (
|
||||
NotifyStatePending = "pending" // 待发
|
||||
NotifyStateSending = "sending" // 已被某个 dispatcher 领取，租约未到期
|
||||
NotifyStateSent = "sent" // 已送达
|
||||
NotifyStateFailed = "failed" // 重试耗尽或永久失败，可手动重发
|
||||
NotifyStateSkipped = "skipped" // 渠道已停用，不再发送
|
||||
)
|
||||
|
||||
// 推送模式。
|
||||
const (
|
||||
NotifyModeRealtime = "realtime"
|
||||
NotifyModeDigest = "digest"
|
||||
)
|
||||
|
||||
// ValidNotifyMode 白名单校验推送模式（与 findings.status 同理：不用 DB CHECK，
|
||||
// 便于后续扩展）。
|
||||
func ValidNotifyMode(m string) bool {
|
||||
return m == NotifyModeRealtime || m == NotifyModeDigest
|
||||
}
|
||||
|
||||
// NotificationChannel 是一个渠道实例配置。Config 与 Filter 保持原始 JSON，
|
||||
// 解析交给 notify 包——db 层不理解它们的字段含义。
|
||||
type NotificationChannel struct {
|
||||
ID int64 `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Kind string `json:"kind"`
|
||||
Mode string `json:"mode"`
|
||||
Config json.RawMessage `json:"config"`
|
||||
Filter json.RawMessage `json:"filter"`
|
||||
// Enabled 用指针是为了区分「没传这个字段」与「显式传 false」——
|
||||
// 前端开关控件只提交被改动的字段。
|
||||
Enabled *bool `json:"enabled,omitempty"`
|
||||
RatePerMin int `json:"rate_per_min"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// IsEnabled 返回渠道是否启用；Enabled 为 nil（未加载）时按启用处理。
|
||||
func (c *NotificationChannel) IsEnabled() bool { return c.Enabled == nil || *c.Enabled }
|
||||
|
||||
// NotificationEvent 是一条事件事实。
|
||||
type NotificationEvent struct {
|
||||
ID int64 `json:"id"`
|
||||
Kind string `json:"kind"`
|
||||
FindingID int64 `json:"finding_id"`
|
||||
Snapshot json.RawMessage `json:"snapshot"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
const notificationChannelCols = `id, name, kind, enabled, config, mode, filter, rate_per_min, created_at, updated_at`
|
||||
|
||||
func scanNotificationChannel(sc interface{ Scan(...any) error }) (*NotificationChannel, error) {
|
||||
var c NotificationChannel
|
||||
var enabled bool
|
||||
if err := sc.Scan(&c.ID, &c.Name, &c.Kind, &enabled, &c.Config, &c.Mode, &c.Filter, &c.RatePerMin, &c.CreatedAt, &c.UpdatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.Enabled = &enabled
|
||||
return &c, nil
|
||||
}
|
||||
|
||||
// ListNotificationChannels 返回全部渠道实例，启用的排在前面、同级按 id。
|
||||
// 排序放在 SQL 里是为了让 UI 与 dispatcher 看到同一个稳定顺序。
|
||||
func (d *DB) ListNotificationChannels(ctx context.Context) ([]*NotificationChannel, error) {
|
||||
rows, err := d.QueryContext(ctx, `SELECT `+notificationChannelCols+` FROM notification_channels
|
||||
ORDER BY enabled DESC, id`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []*NotificationChannel{}
|
||||
for rows.Next() {
|
||||
c, err := scanNotificationChannel(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, c)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// NotificationChannelByID 取单个渠道。
|
||||
func (d *DB) NotificationChannelByID(ctx context.Context, id int64) (*NotificationChannel, error) {
|
||||
row := d.QueryRowContext(ctx, `SELECT `+notificationChannelCols+` FROM notification_channels WHERE id=$1`, id)
|
||||
c, err := scanNotificationChannel(row)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, ErrNotificationChannelNotFound
|
||||
}
|
||||
return c, err
|
||||
}
|
||||
|
||||
// SaveNotificationChannel 新建或更新一个渠道。
|
||||
//
|
||||
// 更新时只覆盖调用方显式给出的字段（非 nil / 非空），这样前端可以提交局部
|
||||
// 修改的抽屉表单，而不必回传 config 里那些它没展示的字段——回传反而会造成
|
||||
// 「掩码值把真实密钥覆盖掉」的事故。
|
||||
func (d *DB) SaveNotificationChannel(ctx context.Context, c *NotificationChannel) (int64, error) {
|
||||
if c.Mode == "" {
|
||||
c.Mode = NotifyModeRealtime
|
||||
}
|
||||
if c.RatePerMin <= 0 {
|
||||
// 0 表示不限流，是合法配置；但未指定时给一个安全默认值，
|
||||
// 避免新渠道默认就把平台打限流。
|
||||
if ch, ok := notify.Get(c.Kind); ok {
|
||||
c.RatePerMin = ch.DefaultRatePerMin()
|
||||
}
|
||||
}
|
||||
if c.Config == nil {
|
||||
c.Config = json.RawMessage(`{}`)
|
||||
}
|
||||
if c.Filter == nil {
|
||||
c.Filter = json.RawMessage(`{}`)
|
||||
}
|
||||
enabled := c.IsEnabled()
|
||||
|
||||
if c.ID == 0 {
|
||||
var id int64
|
||||
err := d.QueryRowContext(ctx, `INSERT INTO notification_channels(name,kind,enabled,config,mode,filter,rate_per_min)
|
||||
VALUES($1,$2,$3,$4,$5,$6,$7) RETURNING id`,
|
||||
c.Name, c.Kind, enabled, string(c.Config), c.Mode, string(c.Filter), c.RatePerMin).Scan(&id)
|
||||
return id, err
|
||||
}
|
||||
res, err := d.ExecContext(ctx, `UPDATE notification_channels
|
||||
SET name=$2, kind=$3, enabled=$4, config=$5, mode=$6, filter=$7, rate_per_min=$8
|
||||
WHERE id=$1`,
|
||||
c.ID, c.Name, c.Kind, enabled, string(c.Config), c.Mode, string(c.Filter), c.RatePerMin)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if n, _ := res.RowsAffected(); n == 0 {
|
||||
return 0, ErrNotificationChannelNotFound
|
||||
}
|
||||
return c.ID, nil
|
||||
}
|
||||
|
||||
// SetNotificationChannelEnabled 切换启停。
|
||||
//
|
||||
// 停用一个渠道时，把它尚未发出的投递一并标记为 skipped：否则重新启用后
|
||||
// 会突然收到一批「停用期间积压」的旧漏洞，时效已失且容易误判为新增。
|
||||
func (d *DB) SetNotificationChannelEnabled(ctx context.Context, id int64, enabled bool) error {
|
||||
return d.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
||||
res, err := tx.ExecContext(ctx, `UPDATE notification_channels SET enabled=$2 WHERE id=$1`, id, enabled)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := res.RowsAffected(); n == 0 {
|
||||
return ErrNotificationChannelNotFound
|
||||
}
|
||||
if !enabled {
|
||||
if _, err := tx.ExecContext(ctx, `UPDATE notification_deliveries SET state=$2, last_error=$3
|
||||
WHERE channel_id=$1 AND state IN ($4,$5)`,
|
||||
id, NotifyStateSkipped, "渠道已停用", NotifyStatePending, NotifyStateSending); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
// DeleteNotificationChannel 删除渠道。其投递历史随外键级联删除
|
||||
// （渠道配置都没了，历史无从解读）。
|
||||
func (d *DB) DeleteNotificationChannel(ctx context.Context, id int64) error {
|
||||
res, err := d.ExecContext(ctx, `DELETE FROM notification_channels WHERE id=$1`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := res.RowsAffected(); n == 0 {
|
||||
return ErrNotificationChannelNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RecordNotificationEventTx 在调用方的事务里**尽力**写入一条推送事件。
|
||||
//
|
||||
// 这是漏洞写入路径上唯一的通知相关改动：一次 INSERT，不读任何表、不认识渠道、
|
||||
// 不跑过滤。事务提交即保证「漏洞落库」与「推送任务存在」原子一致，
|
||||
// 不存在提交成功却没入队、消息永久丢失的窗口。
|
||||
//
|
||||
// 两个关键设计，都不是随手写的：
|
||||
//
|
||||
// 1. **为什么用 SAVEPOINT**：PostgreSQL 里事务内任一语句报错会让整个事务进入
|
||||
// aborted 状态，此后所有语句（含 COMMIT）一律失败。所以「忽略这条 INSERT
|
||||
// 的错误、让调用方继续提交」在 PG 里是做不到的——除非用保存点把错误隔离在
|
||||
// 这一条语句上。没有保存点，就只剩「整笔回滚」这一个选项。
|
||||
//
|
||||
// 2. **为什么整笔回滚是错的**：推送是便利功能，漏洞记录才是产品本身。一个通知
|
||||
// 表的问题（旧库未迁移、磁盘瞬时故障）不该让高危漏洞存不进库。所以这里隔离
|
||||
// 错误、记日志、返回 false，让漏洞写入照常提交——代价是丢掉这一条推送。
|
||||
// 返回 bool 而非 error 是刻意的：调用方不该把它当作会影响写入成败的错误。
|
||||
func RecordNotificationEventTx(ctx context.Context, tx *sql.Tx, kind string, findingID int64, snap notify.Snapshot) bool {
|
||||
raw, err := json.Marshal(snap)
|
||||
if err != nil {
|
||||
log.Printf("[notify] 序列化推送事件失败 finding=%d: %v", findingID, err)
|
||||
return false
|
||||
}
|
||||
if _, err := tx.ExecContext(ctx, `SAVEPOINT notify_event`); err != nil {
|
||||
log.Printf("[notify] 建立保存点失败 finding=%d: %v", findingID, err)
|
||||
return false
|
||||
}
|
||||
if _, err := tx.ExecContext(ctx, `INSERT INTO notification_events(kind,finding_id,snapshot) VALUES($1,$2,$3)`,
|
||||
kind, findingID, string(raw)); err != nil {
|
||||
log.Printf("[notify] 写入推送事件失败 finding=%d（漏洞记录不受影响）: %v", findingID, err)
|
||||
// 回滚到保存点，把事务从 aborted 状态里救回来。
|
||||
if _, rbErr := tx.ExecContext(ctx, `ROLLBACK TO SAVEPOINT notify_event`); rbErr != nil {
|
||||
log.Printf("[notify] 回滚到保存点失败 finding=%d: %v", findingID, rbErr)
|
||||
}
|
||||
return false
|
||||
}
|
||||
// 释放保存点，避免长事务里积攒无用的保存点。
|
||||
_, _ = tx.ExecContext(ctx, `RELEASE SAVEPOINT notify_event`)
|
||||
return true
|
||||
}
|
||||
|
||||
// AddNotificationEvent 是 InsertNotificationEventTx 的独立事务版本，供不在
|
||||
// 既有事务里的调用点使用（如渠道的「发送测试消息」，它没有真实 finding）。
|
||||
func (d *DB) AddNotificationEvent(ctx context.Context, kind string, findingID int64, snap notify.Snapshot) (int64, error) {
|
||||
raw, err := json.Marshal(snap)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("序列化通知事件快照失败: %w", err)
|
||||
}
|
||||
var id int64
|
||||
err = d.QueryRowContext(ctx, `INSERT INTO notification_events(kind,finding_id,snapshot) VALUES($1,$2,$3) RETURNING id`,
|
||||
kind, findingID, string(raw)).Scan(&id)
|
||||
return id, err
|
||||
}
|
||||
|
||||
// FanOutPendingEvents 把尚未分派的漏洞事件按当前启用的渠道展开成投递任务，
|
||||
// 返回本轮处理的事件数与新建的投递数。
|
||||
//
|
||||
// 整轮操作在一个事务里：事件用 FOR UPDATE SKIP LOCKED 领取，多个进程同时跑
|
||||
// 也各自领到不同的行（项目里归档队列的领取用的是同一套手法，见
|
||||
// db/task_archives.go 的 completeNextArchiveJob）。
|
||||
//
|
||||
// 过滤匹配刻意放在 Go 侧而非 SQL：渠道的过滤条件是一组可选字段的 JSONB，
|
||||
// 用 SQL 表达六种组合的匹配会让查询难以维护，而渠道数量是「人手配的几条」，
|
||||
// 全量加载后在内存里逐条比对更快也更好测。
|
||||
//
|
||||
// 未命中任何渠道的事件同样会被标记 fanned_out ——否则它会永远留在待分派集合里，
|
||||
// 每个 tick 被重扫一遍。
|
||||
func (d *DB) FanOutPendingEvents(ctx context.Context, limit int) (eventCount, deliveryCount int, err error) {
|
||||
if limit <= 0 {
|
||||
limit = 200
|
||||
}
|
||||
tx, err := d.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
defer tx.Rollback() //nolint:errcheck // 提交成功后是 no-op
|
||||
|
||||
channels, err := listEnabledNotificationChannelsTx(ctx, tx)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
rows, err := tx.QueryContext(ctx, `SELECT id, kind, finding_id, snapshot FROM notification_events
|
||||
WHERE NOT fanned_out ORDER BY id FOR UPDATE SKIP LOCKED LIMIT $1`, limit)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
var (
|
||||
events []NotificationEvent
|
||||
parsedSnaps []notify.Snapshot
|
||||
)
|
||||
for rows.Next() {
|
||||
var ev NotificationEvent
|
||||
if err := rows.Scan(&ev.ID, &ev.Kind, &ev.FindingID, &ev.Snapshot); err != nil {
|
||||
rows.Close()
|
||||
return 0, 0, err
|
||||
}
|
||||
var snap notify.Snapshot
|
||||
// 快照是我们自己写的，理论上必定可解析；解析失败不阻断投递流程，
|
||||
// 但这条事件会因字段全空而被所有带过滤条件的渠道跳过——宁可少推一条
|
||||
// 也不让一个坏行卡死整个队列。
|
||||
_ = json.Unmarshal(ev.Snapshot, &snap)
|
||||
// kind 以行内值为准：快照里那份是渲染用的副本，可能被旧版本写过。
|
||||
snap.Kind = ev.Kind
|
||||
events = append(events, ev)
|
||||
parsedSnaps = append(parsedSnaps, snap)
|
||||
}
|
||||
rows.Close()
|
||||
if err := rows.Err(); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
if len(events) == 0 {
|
||||
return 0, 0, tx.Commit()
|
||||
}
|
||||
|
||||
type pending struct {
|
||||
eventID int64
|
||||
channelID int64
|
||||
}
|
||||
var toInsert []pending
|
||||
for i, snap := range parsedSnaps {
|
||||
for _, ch := range channels {
|
||||
if !notify.Match(notify.ParseFilter(ch.Filter), snap) {
|
||||
continue
|
||||
}
|
||||
toInsert = append(toInsert, pending{eventID: events[i].ID, channelID: ch.ID})
|
||||
}
|
||||
}
|
||||
if len(toInsert) > 0 {
|
||||
var (
|
||||
vals []string
|
||||
args []any
|
||||
)
|
||||
for _, p := range toInsert {
|
||||
vals = append(vals, fmt.Sprintf("($%d,$%d)", len(args)+1, len(args)+2))
|
||||
args = append(args, p.eventID, p.channelID)
|
||||
}
|
||||
if _, err := tx.ExecContext(ctx, `INSERT INTO notification_deliveries(event_id,channel_id) VALUES `+strings.Join(vals, ","), args...); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
}
|
||||
|
||||
// 标记本轮事件已分派。未命中任何渠道的事件也一起标记（见函数注释）。
|
||||
ids := make([]string, 0, len(events))
|
||||
markArgs := make([]any, 0, len(events))
|
||||
for _, ev := range events {
|
||||
markArgs = append(markArgs, ev.ID)
|
||||
ids = append(ids, fmt.Sprintf("$%d", len(markArgs)))
|
||||
}
|
||||
if _, err := tx.ExecContext(ctx, `UPDATE notification_events SET fanned_out=true WHERE id IN (`+strings.Join(ids, ",")+`)`, markArgs...); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
return len(events), len(toInsert), tx.Commit()
|
||||
}
|
||||
|
||||
// listEnabledNotificationChannelsTx 在事务里取启用中的渠道。数量很少，
|
||||
// 不做分页也不加缓存——缓存会引入「改了配置何时生效」这个额外的时序问题。
|
||||
func listEnabledNotificationChannelsTx(ctx context.Context, tx *sql.Tx) ([]*NotificationChannel, error) {
|
||||
rows, err := tx.QueryContext(ctx, `SELECT id, name, kind, config, mode, filter, rate_per_min
|
||||
FROM notification_channels WHERE enabled ORDER BY id`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []*NotificationChannel{}
|
||||
for rows.Next() {
|
||||
var c NotificationChannel
|
||||
if err := rows.Scan(&c.ID, &c.Name, &c.Kind, &c.Config, &c.Mode, &c.Filter, &c.RatePerMin); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, &c)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// NotificationAssetNames 把资产 id 解析成简短展示名，供推送消息使用。
|
||||
//
|
||||
// 返回顺序与入参一致、长度可能小于入参（不存在的 id 被跳过）。保持入参顺序是
|
||||
// 为了让同一条漏洞的消息在多次投递里资产顺序稳定——否则重试后收到的消息里
|
||||
// 资产次序变了，会被误读成「资产变了」。
|
||||
func (d *DB) NotificationAssetNames(ctx context.Context, ids []int64) ([]string, error) {
|
||||
if len(ids) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
ph, args := placeholders(1, ids)
|
||||
rows, err := d.QueryContext(ctx, `SELECT id, type, domain, ip, url, app_name, bundle_id FROM assets WHERE id IN (`+ph+`)`, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
labels := map[int64]string{}
|
||||
for rows.Next() {
|
||||
var (
|
||||
id int64
|
||||
typ string
|
||||
domain, ip, url sql.NullString
|
||||
appName, bundleID sql.NullString
|
||||
)
|
||||
if err := rows.Scan(&id, &typ, &domain, &ip, &url, &appName, &bundleID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
labels[id] = assetDisplayName(typ, domain.String, ip.String, url.String, appName.String, bundleID.String)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]string, 0, len(ids))
|
||||
seen := map[int64]bool{}
|
||||
for _, id := range ids {
|
||||
if seen[id] {
|
||||
continue
|
||||
}
|
||||
seen[id] = true
|
||||
if label, ok := labels[id]; ok && label != "" {
|
||||
out = append(out, label)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// assetDisplayName 按资产类型挑选最具辨识度的标识。
|
||||
// 兜底返回空串，由调用方决定怎么呈现「名字取不到的资产」——本函数不臆造占位符，
|
||||
// 否则「资产#42」这种噪音会混进推送消息里，读者还以为是真实域名。
|
||||
func assetDisplayName(typ, domain, ip, url, appName, bundleID string) string {
|
||||
pick := func(vals ...string) string {
|
||||
for _, v := range vals {
|
||||
if strings.TrimSpace(v) != "" {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
switch typ {
|
||||
case "root_domain", "subdomain":
|
||||
return domain
|
||||
case "ip":
|
||||
return ip
|
||||
case "app":
|
||||
return pick(appName, bundleID)
|
||||
case "service", "endpoint":
|
||||
return pick(url, domain, ip)
|
||||
default:
|
||||
return pick(domain, ip, url, appName)
|
||||
}
|
||||
}
|
||||
|
||||
// SetFindingStatusWithNotify 更新漏洞处置状态，并在同一事务里登记一条状态变更
|
||||
// 推送事件。
|
||||
//
|
||||
// 返回 from=变更前的状态；found=漏洞是否存在；notified=事件是否登记成功。
|
||||
//
|
||||
// 三条刻意的行为：
|
||||
// - 状态未实际变化时不登记事件。前端抽屉重复提交同一个值、或自动化脚本
|
||||
// 幂等重放，都不该产出推送噪音。
|
||||
// - 漏洞不存在时返回 found=false 且不做任何写入，由调用方翻译成 404。
|
||||
// - 事件登记失败不影响状态更新（见 RecordNotificationEventTx 的保存点说明），
|
||||
// 所以 notified=false 时状态已经改成功了，调用方不应因此报错。
|
||||
func (d *DB) SetFindingStatusWithNotify(ctx context.Context, id int64, status string) (from string, found bool, notified bool, err error) {
|
||||
err = d.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
||||
var (
|
||||
vulnclass, name, severity, summary string
|
||||
taskID sql.NullInt64
|
||||
assetIDs []byte
|
||||
)
|
||||
scanErr := tx.QueryRowContext(ctx, `SELECT vulnclass, name, severity, summary, task_id, asset_ids, status
|
||||
FROM findings WHERE id=$1 FOR UPDATE`, id).
|
||||
Scan(&vulnclass, &name, &severity, &summary, &taskID, &assetIDs, &from)
|
||||
if scanErr == sql.ErrNoRows {
|
||||
return nil
|
||||
}
|
||||
if scanErr != nil {
|
||||
return scanErr
|
||||
}
|
||||
found = true
|
||||
if from == status {
|
||||
return nil
|
||||
}
|
||||
if _, err := tx.ExecContext(ctx, `UPDATE findings SET status=$2 WHERE id=$1`, id, status); err != nil {
|
||||
return err
|
||||
}
|
||||
var assets []int64
|
||||
_ = json.Unmarshal(assetIDs, &assets)
|
||||
notified = RecordNotificationEventTx(ctx, tx, notify.EventFindingStatusChanged, id, notify.Snapshot{
|
||||
Kind: notify.EventFindingStatusChanged,
|
||||
FindingID: id,
|
||||
TaskID: taskID.Int64,
|
||||
VulnClass: vulnclass,
|
||||
Name: name,
|
||||
Severity: severity,
|
||||
Summary: summary,
|
||||
AssetIDs: assets,
|
||||
FromStatus: from,
|
||||
ToStatus: status,
|
||||
})
|
||||
return nil
|
||||
})
|
||||
return from, found, notified, err
|
||||
}
|
||||
|
||||
// NotificationStats 是通知页顶部的概览计数。
|
||||
type NotificationStats struct {
|
||||
Channels int `json:"channels"`
|
||||
ChannelsOn int `json:"channels_on"`
|
||||
Pending int `json:"pending"`
|
||||
Failed int `json:"failed"`
|
||||
SentToday int `json:"sent_today"`
|
||||
BacklogAgeMS int64 `json:"backlog_age_ms"` // 最老的待发投递距今毫秒数
|
||||
}
|
||||
|
||||
// NotificationStatsSnapshot 汇总通知系统的健康度。
|
||||
// BacklogAgeMS 是「推送是不是卡住了」最直接的指标——比 pending 计数有用得多，
|
||||
// 因为积压 3 条和积压 3 条的差别可以是从 3 秒到 3 小时。
|
||||
func (d *DB) NotificationStatsSnapshot(ctx context.Context) (*NotificationStats, error) {
|
||||
var s NotificationStats
|
||||
if err := d.QueryRowContext(ctx, `SELECT
|
||||
(SELECT count(*) FROM notification_channels),
|
||||
(SELECT count(*) FROM notification_channels WHERE enabled),
|
||||
(SELECT count(*) FROM notification_deliveries WHERE state IN ($1,$2)),
|
||||
(SELECT count(*) FROM notification_deliveries WHERE state=$3),
|
||||
(SELECT count(*) FROM notification_deliveries WHERE state=$4 AND sent_at >= date_trunc('day', now())),
|
||||
COALESCE((SELECT EXTRACT(EPOCH FROM (now() - min(created_at))) * 1000 FROM notification_deliveries WHERE state=$1), 0)::bigint`,
|
||||
NotifyStatePending, NotifyStateSending, NotifyStateFailed, NotifyStateSent).
|
||||
Scan(&s.Channels, &s.ChannelsOn, &s.Pending, &s.Failed, &s.SentToday, &s.BacklogAgeMS); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &s, nil
|
||||
}
|
||||
@@ -0,0 +1,410 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// 本文件是投递任务的领取与状态流转。
|
||||
//
|
||||
// 领取用「租约」而非长事务：把行置为 sending 并把 next_attempt_at 推到未来作为
|
||||
// 租约到期时间，提交事务后再去做网络投递。这样投递期间不持有数据库锁——
|
||||
// 网络请求可能耗时数秒（客户端超时 15 秒），占着行锁不放会拖垮同库的其它写操作。
|
||||
//
|
||||
// 代价是进程若在投递途中崩溃，行会停在 sending。这是**可自愈**的：租约到期后
|
||||
// next_attempt_at 落入过去，下一轮领取会把同一行重新捞起来（见领取条件里的
|
||||
// state IN ('pending','sending')）。重试计数在领取时就已 +1，所以崩溃不会造成
|
||||
// 无限重试——MaxNotifyAttempts 次机会用完后落入 failed 等人工处理。
|
||||
|
||||
// MaxNotifyAttempts 是一条投递的最大尝试次数（含首次）。
|
||||
// 定义在这里而非投递引擎里：它是状态机自身的策略，引擎只是执行者。
|
||||
const MaxNotifyAttempts = 3
|
||||
|
||||
// MaxDigestBatchSize 是单个汇总批次一次最多合并多少条投递。
|
||||
//
|
||||
// 存在的理由是资源：一个汇总周期内如果扫出几万个漏洞（完全可能——一次全量扫描
|
||||
// 就能做到），不设上界的话领取会把全部行读进内存、渲染成一条超长消息，
|
||||
// 然后被渠道的长度上限截掉大半——既浪费内存，又**静默丢失**被截掉的那些漏洞。
|
||||
// 设上界后，超出的部分留在库里成为下一个批次，下个周期自然发出去，不会丢。
|
||||
//
|
||||
// 取 500 的依据：它是渲染成消息后在企微 4096 字节上限内还"有内容可读"的量级；
|
||||
// 再大也只是让截断发生在更靠后的位置而已。
|
||||
const MaxDigestBatchSize = 500
|
||||
|
||||
// NotificationDelivery 是一条投递任务，含渲染所需的渠道配置与事件快照。
|
||||
type NotificationDelivery struct {
|
||||
ID int64 `json:"id"`
|
||||
EventID int64 `json:"event_id"`
|
||||
ChannelID int64 `json:"channel_id"`
|
||||
State string `json:"state"`
|
||||
Attempts int `json:"attempts"`
|
||||
NextAttemptAt time.Time `json:"next_attempt_at"`
|
||||
LastError string `json:"last_error"`
|
||||
BatchID *int64 `json:"batch_id,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
SentAt *time.Time `json:"sent_at,omitempty"`
|
||||
Snapshot json.RawMessage `json:"snapshot,omitempty"`
|
||||
// 联合加载的渲染上下文，不进 JSON（由 server 层组装 DTO）。
|
||||
Channel *NotificationChannel `json:"-"`
|
||||
// FindingID/EventKind 从事件带出，供历史列表直接跳转漏洞详情。
|
||||
FindingID int64 `json:"finding_id,string"`
|
||||
EventKind string `json:"event_kind"`
|
||||
// ChannelName/ChannelKind 是列表展示用的冗余字段，省掉前端二次查询。
|
||||
ChannelName string `json:"channel_name"`
|
||||
ChannelKind string `json:"channel_kind"`
|
||||
}
|
||||
|
||||
const notificationDeliveryCols = `d.id, d.event_id, d.channel_id, d.state, d.attempts, d.next_attempt_at,
|
||||
d.last_error, d.batch_id, d.created_at, d.sent_at`
|
||||
|
||||
// joinedDeliveryQuery 是投递行的统一读取形状：投递 + 事件快照 + 渠道配置。
|
||||
// 渲染一条消息三者缺一不可，分开查会写出三次往返。
|
||||
const joinedDeliveryQuery = `SELECT ` + notificationDeliveryCols + `,
|
||||
e.snapshot, e.kind, e.finding_id,
|
||||
c.id, c.name, c.kind, c.enabled, c.config, c.mode, c.filter, c.rate_per_min
|
||||
FROM notification_deliveries d
|
||||
JOIN notification_events e ON e.id = d.event_id
|
||||
JOIN notification_channels c ON c.id = d.channel_id`
|
||||
|
||||
func scanNotificationDelivery(sc interface{ Scan(...any) error }) (*NotificationDelivery, error) {
|
||||
var (
|
||||
dl NotificationDelivery
|
||||
lastErr sql.NullString
|
||||
batchID sql.NullInt64
|
||||
sentAt sql.NullTime
|
||||
snapshot []byte
|
||||
eventKind string
|
||||
channel NotificationChannel
|
||||
chEnabled bool
|
||||
)
|
||||
if err := sc.Scan(&dl.ID, &dl.EventID, &dl.ChannelID, &dl.State, &dl.Attempts, &dl.NextAttemptAt,
|
||||
&lastErr, &batchID, &dl.CreatedAt, &sentAt,
|
||||
&snapshot, &eventKind, &dl.FindingID,
|
||||
&channel.ID, &channel.Name, &channel.Kind, &chEnabled, &channel.Config, &channel.Mode, &channel.Filter, &channel.RatePerMin); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
dl.LastError = lastErr.String
|
||||
if batchID.Valid {
|
||||
dl.BatchID = &batchID.Int64
|
||||
}
|
||||
if sentAt.Valid {
|
||||
dl.SentAt = &sentAt.Time
|
||||
}
|
||||
dl.Snapshot = json.RawMessage(snapshot)
|
||||
dl.EventKind = eventKind
|
||||
dl.ChannelName = channel.Name
|
||||
dl.ChannelKind = channel.Kind
|
||||
channel.Enabled = &chEnabled
|
||||
dl.Channel = &channel
|
||||
return &dl, nil
|
||||
}
|
||||
|
||||
// claimQuery 描述一次领取：先按 sel 选出候选并加锁，再把它们置为 sending 并
|
||||
// 延长租约。sel 里的 lease 位置由调用方用 $n 占位并自行传参。
|
||||
type claimQuery struct {
|
||||
sql string
|
||||
args []any
|
||||
}
|
||||
|
||||
// ClaimRealtimeDeliveries 领取某渠道一批到期的实时投递，最多 limit 条。
|
||||
//
|
||||
// 刻意按**单个渠道**领取而不是「全局领一批再挑着发」：限流闸在投递引擎里按渠道
|
||||
// 维护，只有先知道这个渠道这一轮还能发几条、再去领同样多的行，限流才不会消耗
|
||||
// 重试次数。若反过来先领后弃，被限流挡下的行已经被计过一次 attempts，
|
||||
// 3 次预算会被纯粹的等待耗光，最后落进 failed。
|
||||
//
|
||||
// 条件含「租约已过期的 sending」——那是崩溃自愈的落点。lease 必须显著大于单次
|
||||
// 投递的最坏耗时（渠道 HTTP 客户端超时 15 秒），否则同一行会被两个 dispatcher
|
||||
// 同时投递。同时挡掉已停用渠道：停用操作已把存量投递标记为 skipped，
|
||||
// 这里再拦一道，避免停用与领取并发时的漏网。
|
||||
func (d *DB) ClaimRealtimeDeliveries(ctx context.Context, channelID int64, limit int, lease time.Duration) ([]*NotificationDelivery, error) {
|
||||
if limit <= 0 {
|
||||
return nil, nil
|
||||
}
|
||||
return d.claimDeliveries(ctx, lease, claimQuery{
|
||||
sql: `SELECT dd.id FROM notification_deliveries dd
|
||||
JOIN notification_channels c ON c.id = dd.channel_id
|
||||
WHERE dd.channel_id = $1 AND dd.state IN ($2,$3) AND dd.next_attempt_at <= now()
|
||||
AND c.enabled AND c.mode = $4
|
||||
ORDER BY dd.next_attempt_at, dd.id
|
||||
FOR UPDATE OF dd SKIP LOCKED
|
||||
LIMIT $5`,
|
||||
args: []any{channelID, NotifyStatePending, NotifyStateSending, NotifyModeRealtime, limit},
|
||||
}, nil)
|
||||
}
|
||||
|
||||
// DigestBatchDue 报告该渠道是否已攒够一个到期批次：存在待发投递，且**最老的那条**
|
||||
// 年龄已达到汇总周期。
|
||||
//
|
||||
// 判定依据是最老投递的年龄而非墙上时钟：这样刚建好的渠道不会因为对齐到整点而
|
||||
// 立刻吐出一条只有一条的「汇总」，积压很久的批次也不会再白等一轮。
|
||||
//
|
||||
// 与 ClaimDigestBatch 分开是因为语义不同：本函数只回答「该不该发」，
|
||||
// 而领取要拿走该渠道**全部**待发行（包括尚未满年龄的那些）——否则一个周期
|
||||
// 会被拆成多条消息，汇总就失去意义了。
|
||||
func (d *DB) DigestBatchDue(ctx context.Context, channelID int64, minAge time.Duration) (bool, error) {
|
||||
var due bool
|
||||
err := d.QueryRowContext(ctx, `SELECT EXISTS (
|
||||
SELECT 1 FROM notification_deliveries d
|
||||
JOIN notification_channels c ON c.id = d.channel_id
|
||||
WHERE d.channel_id = $1 AND d.state IN ($2,$3) AND c.enabled
|
||||
GROUP BY d.channel_id
|
||||
HAVING min(d.created_at) <= now() - make_interval(secs => $4)
|
||||
)`, channelID, NotifyStatePending, NotifyStateSending, int64(minAge.Seconds())).Scan(&due)
|
||||
return due, err
|
||||
}
|
||||
|
||||
// ClaimDigestBatch 领取某渠道当前到期的待发投递，作为一个汇总批次，
|
||||
// 单批最多 MaxDigestBatchSize 条。
|
||||
//
|
||||
// 同批次的所有投递共享 batch_id，用集合里的最小 id 作批次号（稳定、可读、
|
||||
// 无需额外序列）。重试时用 COALESCE 保留原批次号，使「这批 N 条是一起发的」
|
||||
// 在多次重试后依然成立。
|
||||
//
|
||||
// 按 id 升序取前 N 条而非随机取：最早产生的投递最先发出去，积压时不会出现
|
||||
// 「新漏洞先发、老漏洞永远排在后面」的饥饿。
|
||||
func (d *DB) ClaimDigestBatch(ctx context.Context, channelID int64, lease time.Duration) ([]*NotificationDelivery, error) {
|
||||
out, err := d.claimDeliveries(ctx, lease, claimQuery{
|
||||
sql: `SELECT dd.id FROM notification_deliveries dd
|
||||
JOIN notification_channels c ON c.id = dd.channel_id
|
||||
WHERE dd.channel_id = $1 AND dd.state IN ($2,$3) AND dd.next_attempt_at <= now() AND c.enabled
|
||||
ORDER BY dd.id
|
||||
FOR UPDATE OF dd SKIP LOCKED
|
||||
LIMIT $4`,
|
||||
args: []any{channelID, NotifyStatePending, NotifyStateSending, MaxDigestBatchSize},
|
||||
}, func(tx *sql.Tx, ids []int64) error {
|
||||
batchID := ids[0]
|
||||
for _, id := range ids {
|
||||
if id < batchID {
|
||||
batchID = id
|
||||
}
|
||||
}
|
||||
ph, idArgs := placeholders(2, ids)
|
||||
_, err := tx.ExecContext(ctx, `UPDATE notification_deliveries SET batch_id = COALESCE(batch_id, $1)
|
||||
WHERE id IN (`+ph+`)`, append([]any{batchID}, idArgs...)...)
|
||||
return err
|
||||
})
|
||||
return out, err
|
||||
}
|
||||
|
||||
// claimDeliveries 执行「选取 + 置 sending 延长租约 + 读取完整行」，全在一个事务里。
|
||||
// postClaim 是可选的附加步骤（汇总批次用它写入 batch_id）。
|
||||
func (d *DB) claimDeliveries(ctx context.Context, lease time.Duration, cq claimQuery, postClaim func(*sql.Tx, []int64) error) ([]*NotificationDelivery, error) {
|
||||
tx, err := d.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer tx.Rollback() //nolint:errcheck // 提交成功后是 no-op
|
||||
|
||||
ids, err := selectForClaim(ctx, tx, cq.sql, cq.args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(ids) == 0 {
|
||||
return nil, tx.Commit()
|
||||
}
|
||||
// 置 sending 并把 next_attempt_at 推到未来：这个未来时刻即租约到期时间，
|
||||
// 「租约未到期」与「未到重试时间」因此共用同一个条件表达，不需要新增列。
|
||||
ph, idArgs := placeholders(3, ids)
|
||||
if _, err := tx.ExecContext(ctx, `UPDATE notification_deliveries
|
||||
SET state=$1, attempts=attempts+1, next_attempt_at=now()+make_interval(secs => $2)
|
||||
WHERE id IN (`+ph+`)`,
|
||||
append([]any{NotifyStateSending, lease.Seconds()}, idArgs...)...); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if postClaim != nil {
|
||||
if err := postClaim(tx, ids); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
out, err := loadDeliveriesTx(ctx, tx, ids)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, tx.Commit()
|
||||
}
|
||||
|
||||
func selectForClaim(ctx context.Context, tx *sql.Tx, query string, args ...any) ([]int64, error) {
|
||||
rows, err := tx.QueryContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var ids []int64
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ids = append(ids, id)
|
||||
}
|
||||
return ids, rows.Err()
|
||||
}
|
||||
|
||||
func loadDeliveriesTx(ctx context.Context, tx *sql.Tx, ids []int64) ([]*NotificationDelivery, error) {
|
||||
ph, args := placeholders(1, ids)
|
||||
rows, err := tx.QueryContext(ctx, joinedDeliveryQuery+` WHERE d.id IN (`+ph+`) ORDER BY d.id`, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []*NotificationDelivery{}
|
||||
for rows.Next() {
|
||||
dl, err := scanNotificationDelivery(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, dl)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// MarkDeliveriesSent 把一批投递标记为已送达。
|
||||
func (d *DB) MarkDeliveriesSent(ctx context.Context, ids []int64) error {
|
||||
ph, args := placeholders(2, ids)
|
||||
if len(args) == 0 {
|
||||
return nil
|
||||
}
|
||||
_, err := d.ExecContext(ctx, `UPDATE notification_deliveries
|
||||
SET state=$1, sent_at=now(), last_error='' WHERE id IN (`+ph+`)`, append([]any{NotifyStateSent}, args...)...)
|
||||
return err
|
||||
}
|
||||
|
||||
// RescheduleDeliveries 把一批投递退回 pending 并推后重试时间。
|
||||
//
|
||||
// 退回 pending 而不是引入新的中间状态，是为了让「还剩几次机会」只由一个地方
|
||||
// 表达（MaxNotifyAttempts），避免状态机的分支随重试策略膨胀。
|
||||
func (d *DB) RescheduleDeliveries(ctx context.Context, ids []int64, delay time.Duration, errMsg string) error {
|
||||
ph, args := placeholders(4, ids)
|
||||
if len(args) == 0 {
|
||||
return nil
|
||||
}
|
||||
_, err := d.ExecContext(ctx, `UPDATE notification_deliveries
|
||||
SET state=$1, next_attempt_at=now()+make_interval(secs => $2), last_error=$3
|
||||
WHERE id IN (`+ph+`)`,
|
||||
append([]any{NotifyStatePending, delay.Seconds(), truncateNotifyError(errMsg)}, args...)...)
|
||||
return err
|
||||
}
|
||||
|
||||
// FailDeliveries 把一批投递标记为最终失败，等待人工在投递历史里重发。
|
||||
func (d *DB) FailDeliveries(ctx context.Context, ids []int64, errMsg string) error {
|
||||
// 占位符从 $3 开始：$1 是 state、$2 是 last_error。
|
||||
ph, args := placeholders(3, ids)
|
||||
if len(args) == 0 {
|
||||
return nil
|
||||
}
|
||||
_, err := d.ExecContext(ctx, `UPDATE notification_deliveries SET state=$1, last_error=$2 WHERE id IN (`+ph+`)`,
|
||||
append([]any{NotifyStateFailed, truncateNotifyError(errMsg)}, args...)...)
|
||||
return err
|
||||
}
|
||||
|
||||
// RetryNotificationDelivery 手动重发一条投递：重置为 pending、清零重试计数、
|
||||
// 立即到期。清计数是刻意的——人工点「重发」意味着前几次失败的原因已被处理，
|
||||
// 再拿旧计数限制它没有道理。
|
||||
func (d *DB) RetryNotificationDelivery(ctx context.Context, id int64) error {
|
||||
res, err := d.ExecContext(ctx, `UPDATE notification_deliveries
|
||||
SET state=$2, attempts=0, next_attempt_at=now(), last_error=''
|
||||
WHERE id=$1 AND state IN ($3,$4)`, id, NotifyStatePending, NotifyStateFailed, NotifyStateSkipped)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := res.RowsAffected(); n == 0 {
|
||||
return fmt.Errorf("投递 %d 不存在或当前状态不允许重发", id)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// NotificationDeliveryFilter 是投递历史的查询条件。
|
||||
type NotificationDeliveryFilter struct {
|
||||
ChannelID int64
|
||||
State string
|
||||
EventKind string
|
||||
}
|
||||
|
||||
func (f NotificationDeliveryFilter) where() (string, []any) {
|
||||
var conds []string
|
||||
var args []any
|
||||
if f.ChannelID > 0 {
|
||||
args = append(args, f.ChannelID)
|
||||
conds = append(conds, fmt.Sprintf("d.channel_id=$%d", len(args)))
|
||||
}
|
||||
if f.State != "" {
|
||||
args = append(args, f.State)
|
||||
conds = append(conds, fmt.Sprintf("d.state=$%d", len(args)))
|
||||
}
|
||||
if f.EventKind != "" {
|
||||
args = append(args, f.EventKind)
|
||||
conds = append(conds, fmt.Sprintf("e.kind=$%d", len(args)))
|
||||
}
|
||||
if len(conds) == 0 {
|
||||
return "", nil
|
||||
}
|
||||
return " WHERE " + strings.Join(conds, " AND "), args
|
||||
}
|
||||
|
||||
// ListNotificationDeliveries 分页返回投递历史，新的在前。
|
||||
func (d *DB) ListNotificationDeliveries(ctx context.Context, f NotificationDeliveryFilter, page, pageSize int) ([]*NotificationDelivery, int, error) {
|
||||
if page < 1 {
|
||||
page = 1
|
||||
}
|
||||
if pageSize <= 0 || pageSize > 200 {
|
||||
pageSize = 50
|
||||
}
|
||||
where, args := f.where()
|
||||
|
||||
var total int
|
||||
if err := d.QueryRowContext(ctx, `SELECT count(*) FROM notification_deliveries d
|
||||
JOIN notification_events e ON e.id = d.event_id`+where, args...).Scan(&total); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
q := fmt.Sprintf("%s%s ORDER BY d.id DESC LIMIT $%d OFFSET $%d",
|
||||
joinedDeliveryQuery, where, len(args)+1, len(args)+2)
|
||||
rows, err := d.QueryContext(ctx, q, append(args, pageSize, (page-1)*pageSize)...)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []*NotificationDelivery{}
|
||||
for rows.Next() {
|
||||
dl, err := scanNotificationDelivery(rows)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
out = append(out, dl)
|
||||
}
|
||||
return out, total, rows.Err()
|
||||
}
|
||||
|
||||
// truncateNotifyError 把错误信息截到列可接受的长度。渠道返回的响应体可能很长
|
||||
// （通用 Webhook 打到自建服务时尤甚），不截断会让历史列表的载荷膨胀。
|
||||
func truncateNotifyError(msg string) string {
|
||||
const max = 500
|
||||
if len(msg) <= max {
|
||||
return msg
|
||||
}
|
||||
// 按字符边界回退，避免留下半个 UTF-8 字符让前端显示成乱码。
|
||||
cut := max
|
||||
for cut > 0 && !isUTF8Start(msg[cut]) {
|
||||
cut--
|
||||
}
|
||||
return msg[:cut] + "…"
|
||||
}
|
||||
|
||||
func isUTF8Start(b byte) bool { return b&0xC0 != 0x80 }
|
||||
|
||||
// placeholders 生成从 start 开始的 $n 占位串及对应参数，供 IN (...) 使用。
|
||||
// 例如 start=3, ids=[7,8] → "$3,$4", [7,8]。
|
||||
func placeholders(start int, ids []int64) (string, []any) {
|
||||
ph := make([]string, 0, len(ids))
|
||||
args := make([]any, 0, len(ids))
|
||||
for i, id := range ids {
|
||||
ph = append(ph, fmt.Sprintf("$%d", start+i))
|
||||
args = append(args, id)
|
||||
}
|
||||
return strings.Join(ph, ","), args
|
||||
}
|
||||
@@ -0,0 +1,758 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/notify"
|
||||
)
|
||||
|
||||
// 本文件的用例都会真连 PostgreSQL（无库时跳过）。这些 SQL 用到了
|
||||
// FOR UPDATE SKIP LOCKED、make_interval、JSONB、多行 IN(...) 占位符拼接，
|
||||
// 都是「编译通过但可能运行时报错」的写法，必须实跑才算验证过。
|
||||
|
||||
func notifyTestDB(t *testing.T) *DB {
|
||||
t.Helper()
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
t.Cleanup(func() { d.Close() })
|
||||
return d
|
||||
}
|
||||
|
||||
// newTestChannel 建一个渠道，测试结束自动删除。
|
||||
func newTestChannel(t *testing.T, d *DB, kind, mode string, filter string) *NotificationChannel {
|
||||
t.Helper()
|
||||
if filter == "" {
|
||||
filter = `{}`
|
||||
}
|
||||
ch := &NotificationChannel{
|
||||
Name: "测试渠道-" + t.Name(),
|
||||
Kind: kind,
|
||||
Mode: mode,
|
||||
Config: json.RawMessage(`{"webhook":"https://example.com/hook"}`),
|
||||
Filter: json.RawMessage(filter),
|
||||
RatePerMin: 100,
|
||||
}
|
||||
id, err := d.SaveNotificationChannel(context.Background(), ch)
|
||||
if err != nil {
|
||||
t.Fatalf("建渠道失败: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { d.Exec(`DELETE FROM notification_channels WHERE id=$1`, id) })
|
||||
ch.ID = id
|
||||
return ch
|
||||
}
|
||||
|
||||
// addTestEvent 直接写一条事件（不经 finding），用于测试分派与投递。
|
||||
func addTestEvent(t *testing.T, d *DB, kind string, findingID int64, snap notify.Snapshot) int64 {
|
||||
t.Helper()
|
||||
snap.Kind = kind
|
||||
snap.FindingID = findingID
|
||||
id, err := d.AddNotificationEvent(context.Background(), kind, findingID, snap)
|
||||
if err != nil {
|
||||
t.Fatalf("写事件失败: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { d.Exec(`DELETE FROM notification_events WHERE id=$1`, id) })
|
||||
return id
|
||||
}
|
||||
|
||||
func TestNotificationAssetNamesResolvesAndPreservesOrder(t *testing.T) {
|
||||
d := notifyTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 三类资产各有各的展示口径：域名、IP、URL。
|
||||
insertAsset := func(query, value string) int64 {
|
||||
t.Helper()
|
||||
var id int64
|
||||
if err := d.QueryRow(query, value).Scan(&id); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return id
|
||||
}
|
||||
domID := insertAsset(`INSERT INTO assets(type, domain) VALUES('subdomain',$1) RETURNING id`, "a.example.com")
|
||||
ipID := insertAsset(`INSERT INTO assets(type, ip) VALUES('ip',$1) RETURNING id`, "10.1.2.3")
|
||||
svcID := insertAsset(`INSERT INTO assets(type, url) VALUES('service',$1) RETURNING id`, "https://a.example.com/admin")
|
||||
t.Cleanup(func() {
|
||||
d.Exec(`DELETE FROM assets WHERE id IN ($1,$2,$3)`, domID, ipID, svcID)
|
||||
})
|
||||
|
||||
// 传入顺序刻意乱序，且含一个不存在的 id。
|
||||
got, err := d.NotificationAssetNames(ctx, []int64{svcID, 999999999, domID, ipID, svcID})
|
||||
if err != nil {
|
||||
t.Fatalf("解析资产名失败: %v", err)
|
||||
}
|
||||
want := []string{"https://a.example.com/admin", "a.example.com", "10.1.2.3"}
|
||||
if len(got) != len(want) {
|
||||
t.Fatalf("资产名数量不符，期望 %v 得到 %v", want, got)
|
||||
}
|
||||
for i := range want {
|
||||
if got[i] != want[i] {
|
||||
t.Fatalf("顺序/取值不符，期望 %v 得到 %v", want, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRecordNotificationEventTxUnwindsOnFailure 是保存点机制的核心用例：
|
||||
// 在事务里先让 notification_events 的写入必然失败（临时加一个恒 false 的约束），
|
||||
// 断言 ① 该函数报 false ② 事务没有进入 aborted 状态，后续语句仍能执行。
|
||||
//
|
||||
// 没有保存点的话，PostgreSQL 会让整个事务作废，后续任何语句都以
|
||||
// "current transaction is aborted" 失败——那正是「一个通知表的问题导致
|
||||
// 漏洞存不进库」的故障路径。
|
||||
//
|
||||
// 这里刻意用 **ROLLBACK 收尾而不是 COMMIT**：ALTER TABLE 在 PG 里是事务性的，
|
||||
// 一旦提交，那个临时约束就会永久留在 schema 里，把后续所有用例一起打挂。
|
||||
// 回滚能自动撤销 DDL，无需手工清理。断言只需要「事务还活着」，
|
||||
// 不需要真的提交。
|
||||
func TestRecordNotificationEventTxUnwindsOnFailure(t *testing.T) {
|
||||
d := notifyTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 防御性清理：若历史运行留下过这个约束，先摘掉。
|
||||
if _, err := d.Exec(`ALTER TABLE notification_events DROP CONSTRAINT IF EXISTS notify_test_never`); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
tx, err := d.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer tx.Rollback() //nolint:errcheck // 撤销临时约束，见函数注释
|
||||
|
||||
// NOT VALID：只约束此后写入的行，不去校验库里已有的历史事件
|
||||
// （否则存量行违规会导致约束加不上）。
|
||||
if _, err := tx.ExecContext(ctx, `ALTER TABLE notification_events ADD CONSTRAINT notify_test_never CHECK (false) NOT VALID`); err != nil {
|
||||
t.Fatalf("加临时约束失败: %v", err)
|
||||
}
|
||||
if RecordNotificationEventTx(ctx, tx, notify.EventFindingCreated, 1, notify.Snapshot{Severity: "high"}) {
|
||||
t.Fatal("在必然失败的约束下仍报告写入成功")
|
||||
}
|
||||
// 关键断言：事务还能用。
|
||||
var one int
|
||||
if err := tx.QueryRowContext(ctx, `SELECT 1`).Scan(&one); err != nil {
|
||||
t.Fatalf("事务已被污染（保存点未生效）: %v", err)
|
||||
}
|
||||
if err := tx.Rollback(); err != nil {
|
||||
t.Fatalf("回滚失败: %v", err)
|
||||
}
|
||||
// 确认 DDL 已随回滚撤销，不给后续用例留雷。
|
||||
var exists bool
|
||||
if err := d.QueryRow(`SELECT EXISTS(SELECT 1 FROM pg_constraint WHERE conname='notify_test_never')`).Scan(&exists); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if exists {
|
||||
t.Fatal("临时约束未被回滚撤销，会污染后续用例")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFanOutRoutesEventsByFilter(t *testing.T) {
|
||||
d := notifyTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
all := newTestChannel(t, d, notify.KindDingTalk, NotifyModeRealtime, `{}`)
|
||||
onlyCritical := newTestChannel(t, d, notify.KindDingTalk, NotifyModeRealtime, `{"min_severity":"critical"}`)
|
||||
sqlOnly := newTestChannel(t, d, notify.KindDingTalk, NotifyModeRealtime, `{"vulnclass_include":["SQL"]}`)
|
||||
|
||||
highSQL := addTestEvent(t, d, notify.EventFindingCreated, 1001, notify.Snapshot{Severity: "high", VulnClass: "SQL注入"})
|
||||
lowXSS := addTestEvent(t, d, notify.EventFindingCreated, 1002, notify.Snapshot{Severity: "low", VulnClass: "XSS"})
|
||||
criticalXSS := addTestEvent(t, d, notify.EventFindingCreated, 1003, notify.Snapshot{Severity: "critical", VulnClass: "XSS"})
|
||||
|
||||
if _, _, err := d.FanOutPendingEvents(ctx, 100); err != nil {
|
||||
t.Fatalf("分派失败: %v", err)
|
||||
}
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
eventID int64
|
||||
channel int64
|
||||
want bool
|
||||
}{
|
||||
{"全收渠道收到 high", highSQL, all.ID, true},
|
||||
{"全收渠道收到 low", lowXSS, all.ID, true},
|
||||
{"仅严重渠道跳过 high", highSQL, onlyCritical.ID, false},
|
||||
{"仅严重渠道收到 critical", criticalXSS, onlyCritical.ID, true},
|
||||
{"仅SQL渠道收到 SQL", highSQL, sqlOnly.ID, true},
|
||||
{"仅SQL渠道跳过 XSS", lowXSS, sqlOnly.ID, false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var exists bool
|
||||
if err := d.QueryRowContext(ctx, `SELECT EXISTS(SELECT 1 FROM notification_deliveries WHERE event_id=$1 AND channel_id=$2)`,
|
||||
tc.eventID, tc.channel).Scan(&exists); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if exists != tc.want {
|
||||
t.Fatalf("投递是否存在: 期望 %v 得到 %v", tc.want, exists)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// 再分派一次不应产生重复投递（fanned_out 幂等）。
|
||||
events, deliveries, err := d.FanOutPendingEvents(ctx, 100)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if events != 0 || deliveries != 0 {
|
||||
t.Fatalf("已分派的事件不应被再次处理，得到 events=%d deliveries=%d", events, deliveries)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFanOutMarksEventsWithNoMatchingChannel 覆盖「事件没命中任何渠道」的情况。
|
||||
// 这类事件必须照样被标记为已分派，否则它会永远留在待分派集合里、每个 tick 重扫。
|
||||
func TestFanOutMarksEventsWithNoMatchingChannel(t *testing.T) {
|
||||
d := notifyTestDB(t)
|
||||
ctx := context.Background()
|
||||
pick := newTestChannel(t, d, notify.KindDingTalk, NotifyModeRealtime, `{"vulnclass_include":["绝不匹配的类型"]}`)
|
||||
_ = pick
|
||||
|
||||
ev := addTestEvent(t, d, notify.EventFindingCreated, 2001, notify.Snapshot{Severity: "high", VulnClass: "XSS"})
|
||||
_, deliveries, err := d.FanOutPendingEvents(ctx, 100)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if deliveries != 0 {
|
||||
t.Fatalf("不该产生投递，得到 %d", deliveries)
|
||||
}
|
||||
var fanned bool
|
||||
if err := d.QueryRowContext(ctx, `SELECT fanned_out FROM notification_events WHERE id=$1`, ev).Scan(&fanned); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !fanned {
|
||||
t.Fatal("未命中渠道的事件也必须标记为已分派，否则会被无限重扫")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClaimRealtimeDeliveriesHonorsLeaseAndMode(t *testing.T) {
|
||||
d := notifyTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
realtime := newTestChannel(t, d, notify.KindDingTalk, NotifyModeRealtime, `{}`)
|
||||
digest := newTestChannel(t, d, notify.KindDingTalk, NotifyModeDigest, `{}`)
|
||||
|
||||
addTestEvent(t, d, notify.EventFindingCreated, 3001, notify.Snapshot{Severity: "high", VulnClass: "XSS"})
|
||||
if _, _, err := d.FanOutPendingEvents(ctx, 100); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// 实时领取只应拿到 realtime 渠道的那条，不该动 digest 渠道的。
|
||||
got, err := d.ClaimRealtimeDeliveries(ctx, realtime.ID, 10, time.Minute)
|
||||
if err != nil {
|
||||
t.Fatalf("领取失败: %v", err)
|
||||
}
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("应领到 1 条，得到 %d", len(got))
|
||||
}
|
||||
if got[0].State != NotifyStateSending || got[0].Attempts != 1 {
|
||||
t.Fatalf("领取后应为 sending 且 attempts=1，得到 state=%s attempts=%d", got[0].State, got[0].Attempts)
|
||||
}
|
||||
// 关联加载的渲染上下文必须齐全（渠道配置 + 事件快照 + finding id）。
|
||||
if got[0].Channel == nil || len(got[0].Channel.Config) == 0 {
|
||||
t.Fatal("领取结果缺少渠道配置，渲染会失败")
|
||||
}
|
||||
if got[0].FindingID != 3001 {
|
||||
t.Fatalf("finding id 未从事件带出，得到 %d", got[0].FindingID)
|
||||
}
|
||||
|
||||
// 租约未到期，第二次领取应为空——这是「同一行不会被两个 dispatcher 同时投递」
|
||||
// 的保证。
|
||||
again, err := d.ClaimRealtimeDeliveries(ctx, realtime.ID, 10, time.Minute)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(again) != 0 {
|
||||
t.Fatalf("租约期内不应重复领取，得到 %d 条", len(again))
|
||||
}
|
||||
|
||||
// digest 渠道的投递不应被实时领取碰到。
|
||||
left, err := d.ClaimRealtimeDeliveries(ctx, digest.ID, 10, time.Minute)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(left) != 0 {
|
||||
t.Fatalf("实时领取不应拿到 digest 渠道的投递，得到 %d 条", len(left))
|
||||
}
|
||||
}
|
||||
|
||||
// TestClaimExpiredLeaseRecovers 覆盖崩溃自愈：进程在投递途中挂掉会留下 sending
|
||||
// 行，租约到期后必须能被重新领起来，否则这条投递永远卡住。
|
||||
func TestClaimExpiredLeaseRecovers(t *testing.T) {
|
||||
d := notifyTestDB(t)
|
||||
ctx := context.Background()
|
||||
ch := newTestChannel(t, d, notify.KindDingTalk, NotifyModeRealtime, `{}`)
|
||||
addTestEvent(t, d, notify.EventFindingCreated, 4001, notify.Snapshot{Severity: "high"})
|
||||
if _, _, err := d.FanOutPendingEvents(ctx, 100); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
first, err := d.ClaimRealtimeDeliveries(ctx, ch.ID, 10, time.Minute)
|
||||
if err != nil || len(first) != 1 {
|
||||
t.Fatalf("首次领取失败: %v (%d 条)", err, len(first))
|
||||
}
|
||||
// 把租约手动推到过去，模拟「租约已过期」。
|
||||
if _, err := d.Exec(`UPDATE notification_deliveries SET next_attempt_at = now() - interval '1 minute' WHERE id=$1`, first[0].ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
second, err := d.ClaimRealtimeDeliveries(ctx, ch.ID, 10, time.Minute)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(second) != 1 {
|
||||
t.Fatalf("租约过期的 sending 行应可被重新领取，得到 %d 条", len(second))
|
||||
}
|
||||
if second[0].Attempts != 2 {
|
||||
t.Fatalf("重新领取应累加尝试次数，得到 %d", second[0].Attempts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClaimSkipsDisabledChannel(t *testing.T) {
|
||||
d := notifyTestDB(t)
|
||||
ctx := context.Background()
|
||||
ch := newTestChannel(t, d, notify.KindDingTalk, NotifyModeRealtime, `{}`)
|
||||
addTestEvent(t, d, notify.EventFindingCreated, 5001, notify.Snapshot{Severity: "high"})
|
||||
if _, _, err := d.FanOutPendingEvents(ctx, 100); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// 停用会把存量待发投递一起标记为 skipped。
|
||||
if err := d.SetNotificationChannelEnabled(ctx, ch.ID, false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var state string
|
||||
if err := d.QueryRow(`SELECT state FROM notification_deliveries WHERE channel_id=$1`, ch.ID).Scan(&state); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if state != NotifyStateSkipped {
|
||||
t.Fatalf("停用渠道的存量待发投递应被标记为 skipped，得到 %s", state)
|
||||
}
|
||||
got, err := d.ClaimRealtimeDeliveries(ctx, ch.ID, 10, time.Minute)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(got) != 0 {
|
||||
t.Fatalf("已停用渠道不应能被领取，得到 %d 条", len(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDigestBatchDueAndStableBatchID(t *testing.T) {
|
||||
d := notifyTestDB(t)
|
||||
ctx := context.Background()
|
||||
ch := newTestChannel(t, d, notify.KindDingTalk, NotifyModeDigest, `{}`)
|
||||
for i := 0; i < 3; i++ {
|
||||
addTestEvent(t, d, notify.EventFindingCreated, int64(6000+i), notify.Snapshot{Severity: "high"})
|
||||
}
|
||||
if _, _, err := d.FanOutPendingEvents(ctx, 100); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// 批次刚建、年龄为 0，30 分钟的周期下不该到期。
|
||||
due, err := d.DigestBatchDue(ctx, ch.ID, 30*time.Minute)
|
||||
if err != nil {
|
||||
t.Fatalf("判断批次到期失败: %v", err)
|
||||
}
|
||||
if due {
|
||||
t.Fatal("刚建立的批次不应立即到期")
|
||||
}
|
||||
|
||||
// 把三条投递的创建时间一起推老，模拟一个攒够周期的批次。
|
||||
if _, err := d.Exec(`UPDATE notification_deliveries SET created_at = now() - interval '40 minutes' WHERE channel_id=$1`, ch.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
due, err = d.DigestBatchDue(ctx, ch.ID, 30*time.Minute)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !due {
|
||||
t.Fatal("超过周期的批次应判定为到期")
|
||||
}
|
||||
|
||||
batch, err := d.ClaimDigestBatch(ctx, ch.ID, time.Minute)
|
||||
if err != nil {
|
||||
t.Fatalf("领取汇总批次失败: %v", err)
|
||||
}
|
||||
if len(batch) != 3 {
|
||||
t.Fatalf("汇总应一次领走全部 3 条，得到 %d 条", len(batch))
|
||||
}
|
||||
if batch[0].BatchID == nil {
|
||||
t.Fatal("汇总批次必须写 batch_id，否则历史里看不出这几条是一起发的")
|
||||
}
|
||||
firstBatchID := *batch[0].BatchID
|
||||
for _, dl := range batch {
|
||||
if dl.BatchID == nil || *dl.BatchID != firstBatchID {
|
||||
t.Fatalf("同一批次应共享 batch_id，得到 %v vs %d", dl.BatchID, firstBatchID)
|
||||
}
|
||||
}
|
||||
|
||||
// 让这批**整体**失败重排后再领，batch_id 必须保持原值（COALESCE 的作用）：
|
||||
// 否则一次重试就把「这批是一起发的」这个事实抹掉了。
|
||||
//
|
||||
// 必须整批重排而不是只重排一条——投递引擎发汇总消息时就是这样处理的
|
||||
// （一条消息代表整批，成败与共）。只重排一条的话，其余仍在租约期内，
|
||||
// 重领自然只拿到那一条。
|
||||
allIDs := make([]int64, 0, len(batch))
|
||||
for _, dl := range batch {
|
||||
allIDs = append(allIDs, dl.ID)
|
||||
}
|
||||
if err := d.RescheduleDeliveries(ctx, allIDs, time.Second, "模拟失败"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// 把租约推到过去，模拟退避时间已到。
|
||||
if _, err := d.Exec(`UPDATE notification_deliveries SET next_attempt_at = now() - interval '1 minute' WHERE channel_id=$1`, ch.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
reclaimed, err := d.ClaimDigestBatch(ctx, ch.ID, time.Minute)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(reclaimed) != 3 {
|
||||
t.Fatalf("重领应拿到全部 3 条，得到 %d", len(reclaimed))
|
||||
}
|
||||
if reclaimed[0].BatchID == nil || *reclaimed[0].BatchID != firstBatchID {
|
||||
t.Fatalf("重试后 batch_id 应保持原值 %d，得到 %v", firstBatchID, reclaimed[0].BatchID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeliveryStateTransitions(t *testing.T) {
|
||||
d := notifyTestDB(t)
|
||||
ctx := context.Background()
|
||||
ch := newTestChannel(t, d, notify.KindDingTalk, NotifyModeRealtime, `{}`)
|
||||
addTestEvent(t, d, notify.EventFindingCreated, 7001, notify.Snapshot{Severity: "high"})
|
||||
if _, _, err := d.FanOutPendingEvents(ctx, 100); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := d.ClaimRealtimeDeliveries(ctx, ch.ID, 10, time.Minute)
|
||||
if err != nil || len(got) != 1 {
|
||||
t.Fatalf("领取失败: %v (%d)", err, len(got))
|
||||
}
|
||||
id := got[0].ID
|
||||
|
||||
if err := d.RescheduleDeliveries(ctx, []int64{id}, time.Second, "网络抖动"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var state, lastErr string
|
||||
if err := d.QueryRow(`SELECT state, last_error FROM notification_deliveries WHERE id=$1`, id).Scan(&state, &lastErr); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if state != NotifyStatePending || lastErr != "网络抖动" {
|
||||
t.Fatalf("重排后应为 pending 并记录原因，得到 state=%s err=%q", state, lastErr)
|
||||
}
|
||||
|
||||
if err := d.FailDeliveries(ctx, []int64{id}, "重试耗尽"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := d.QueryRow(`SELECT state FROM notification_deliveries WHERE id=$1`, id).Scan(&state); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if state != NotifyStateFailed {
|
||||
t.Fatalf("应为 failed，得到 %s", state)
|
||||
}
|
||||
|
||||
// 手动重发要清零重试计数并立即到期，否则会继承旧的失败预算。
|
||||
if err := d.RetryNotificationDelivery(ctx, id); err != nil {
|
||||
t.Fatalf("重发失败: %v", err)
|
||||
}
|
||||
var attempts int
|
||||
var next time.Time
|
||||
if err := d.QueryRow(`SELECT state, attempts, next_attempt_at FROM notification_deliveries WHERE id=$1`, id).Scan(&state, &attempts, &next); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if state != NotifyStatePending || attempts != 0 {
|
||||
t.Fatalf("重发后应为 pending 且 attempts=0，得到 state=%s attempts=%d", state, attempts)
|
||||
}
|
||||
if next.After(time.Now().Add(time.Second)) {
|
||||
t.Fatal("重发应立即可领")
|
||||
}
|
||||
|
||||
// 已送达的投递不应能被重发。
|
||||
if err := d.MarkDeliveriesSent(ctx, []int64{id}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := d.RetryNotificationDelivery(ctx, id); err == nil {
|
||||
t.Fatal("已送达的投递不该允许重发")
|
||||
}
|
||||
}
|
||||
|
||||
func TestListNotificationDeliveriesPagingAndFilter(t *testing.T) {
|
||||
d := notifyTestDB(t)
|
||||
ctx := context.Background()
|
||||
ch := newTestChannel(t, d, notify.KindDingTalk, NotifyModeRealtime, `{}`)
|
||||
for i := 0; i < 5; i++ {
|
||||
addTestEvent(t, d, notify.EventFindingCreated, int64(8000+i), notify.Snapshot{Severity: "high", Name: "分页测试"})
|
||||
}
|
||||
if _, _, err := d.FanOutPendingEvents(ctx, 100); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := d.ClaimRealtimeDeliveries(ctx, ch.ID, 10, time.Minute); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
page1, total, err := d.ListNotificationDeliveries(ctx, NotificationDeliveryFilter{ChannelID: ch.ID, State: NotifyStateSending}, 1, 2)
|
||||
if err != nil {
|
||||
t.Fatalf("查询失败: %v", err)
|
||||
}
|
||||
if total != 5 {
|
||||
t.Fatalf("总数应为 5，得到 %d", total)
|
||||
}
|
||||
if len(page1) != 2 {
|
||||
t.Fatalf("每页 2 条，得到 %d", len(page1))
|
||||
}
|
||||
// 新的在前：第一页首条 id 应大于第二页首条。
|
||||
page2, _, err := d.ListNotificationDeliveries(ctx, NotificationDeliveryFilter{ChannelID: ch.ID, State: NotifyStateSending}, 2, 2)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(page2) != 2 || page2[0].ID >= page1[0].ID {
|
||||
t.Fatalf("分页顺序应为新的在前，得到 page1[0]=%d page2[0]=%d", page1[0].ID, page2[0].ID)
|
||||
}
|
||||
// 渲染上下文必须随历史一起返回，否则列表无法展示「推的是什么」。
|
||||
if page1[0].ChannelName == "" || page1[0].FindingID == 0 {
|
||||
t.Fatalf("历史项缺少展示字段: %+v", page1[0])
|
||||
}
|
||||
|
||||
// 按状态过滤：没有 pending 的。
|
||||
pending, totalPending, err := d.ListNotificationDeliveries(ctx, NotificationDeliveryFilter{ChannelID: ch.ID, State: NotifyStatePending}, 1, 50)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if totalPending != 0 || len(pending) != 0 {
|
||||
t.Fatalf("不该有 pending 投递，得到 %d 条 (total=%d)", len(pending), totalPending)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetFindingStatusWithNotifyOnlyEmitsOnRealChange(t *testing.T) {
|
||||
d := notifyTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
tk, err := d.CreateTask("通知状态变更测试", "目标", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.DeleteTask(tk.ID)
|
||||
es := d.Exploration(tk.ExplorationID)
|
||||
f, err := es.RecordFinding(ctx, RecordFindingInput{
|
||||
TaskID: tk.ID, Worker: "test", VulnClass: "SQL注入", Name: "状态变更用例",
|
||||
Severity: "high", Summary: "摘要",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { d.Exec(`DELETE FROM notification_events WHERE finding_id=$1`, f.FindingID) })
|
||||
|
||||
// 落库时已登记一条 finding_created 事件，先把它数出来作为基线。
|
||||
var base int
|
||||
if err := d.QueryRow(`SELECT count(*) FROM notification_events WHERE finding_id=$1`, f.FindingID).Scan(&base); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if base < 1 {
|
||||
t.Fatal("漏洞落库应在同一事务里登记一条推送事件")
|
||||
}
|
||||
|
||||
// 改成同一个状态：不该产生事件（避免重复提交刷出推送噪音）。
|
||||
from, found, notified, err := d.SetFindingStatusWithNotify(ctx, f.FindingID, "pending")
|
||||
if err != nil || !found {
|
||||
t.Fatalf("状态设置失败: found=%v err=%v", found, err)
|
||||
}
|
||||
if notified {
|
||||
t.Fatal("状态未变化时不应登记推送事件")
|
||||
}
|
||||
if from != "pending" {
|
||||
t.Fatalf("应返回变更前状态 pending，得到 %q", from)
|
||||
}
|
||||
|
||||
// 真正变更：应登记事件并记录 from/to。
|
||||
from, found, notified, err = d.SetFindingStatusWithNotify(ctx, f.FindingID, "fixed")
|
||||
if err != nil || !found {
|
||||
t.Fatalf("状态设置失败: found=%v err=%v", found, err)
|
||||
}
|
||||
if !notified {
|
||||
t.Fatal("状态实际变更时应登记推送事件")
|
||||
}
|
||||
if from != "pending" {
|
||||
t.Fatalf("from 应为 pending，得到 %q", from)
|
||||
}
|
||||
var snapshot []byte
|
||||
if err := d.QueryRow(`SELECT snapshot FROM notification_events WHERE finding_id=$1 AND kind=$2`,
|
||||
f.FindingID, notify.EventFindingStatusChanged).Scan(&snapshot); err != nil {
|
||||
t.Fatalf("未找到状态变更事件: %v", err)
|
||||
}
|
||||
var snap notify.Snapshot
|
||||
if err := json.Unmarshal(snapshot, &snap); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if snap.FromStatus != "pending" || snap.ToStatus != "fixed" {
|
||||
t.Fatalf("快照里的状态流转不对: %s → %s", snap.FromStatus, snap.ToStatus)
|
||||
}
|
||||
// 快照要带上渲染所需字段，否则状态变更消息会是空壳。
|
||||
if snap.VulnClass != "SQL注入" || snap.Severity != "high" || snap.Name != "状态变更用例" {
|
||||
t.Fatalf("快照缺少渲染字段: %+v", snap)
|
||||
}
|
||||
var status string
|
||||
if err := d.QueryRow(`SELECT status FROM findings WHERE id=$1`, f.FindingID).Scan(&status); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if status != "fixed" {
|
||||
t.Fatalf("状态应已更新为 fixed，得到 %s", status)
|
||||
}
|
||||
|
||||
// 不存在的漏洞：found=false，不报错。
|
||||
if _, found, _, err := d.SetFindingStatusWithNotify(ctx, 999999999, "fixed"); err != nil || found {
|
||||
t.Fatalf("不存在的漏洞应返回 found=false 且无错，得到 found=%v err=%v", found, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNotificationStatsSnapshot(t *testing.T) {
|
||||
d := notifyTestDB(t)
|
||||
ctx := context.Background()
|
||||
ch := newTestChannel(t, d, notify.KindDingTalk, NotifyModeRealtime, `{}`)
|
||||
addTestEvent(t, d, notify.EventFindingCreated, 9001, notify.Snapshot{Severity: "high"})
|
||||
if _, _, err := d.FanOutPendingEvents(ctx, 100); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
stats, err := d.NotificationStatsSnapshot(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("统计失败: %v", err)
|
||||
}
|
||||
if stats.Channels < 1 || stats.ChannelsOn < 1 {
|
||||
t.Fatalf("渠道计数不对: %+v", stats)
|
||||
}
|
||||
if stats.Pending < 1 {
|
||||
t.Fatalf("应统计到待发投递: %+v", stats)
|
||||
}
|
||||
// 刚建的投递积压年龄应接近 0，而不是负数或巨大值。
|
||||
if stats.BacklogAgeMS < 0 || stats.BacklogAgeMS > int64(time.Hour/time.Millisecond) {
|
||||
t.Fatalf("积压年龄不合法: %d ms", stats.BacklogAgeMS)
|
||||
}
|
||||
_ = ch
|
||||
}
|
||||
|
||||
func TestNotificationChannelCRUDRoundTrip(t *testing.T) {
|
||||
d := notifyTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
ch := &NotificationChannel{
|
||||
Name: "CRUD 往返",
|
||||
Kind: notify.KindEmail,
|
||||
Mode: NotifyModeDigest,
|
||||
Config: json.RawMessage(`{"host":"smtp.example.com","port":587,"from":"[email protected]","to":["[email protected]"]}`),
|
||||
Filter: json.RawMessage(`{"min_severity":"medium","on_status_change":true}`),
|
||||
RatePerMin: 42,
|
||||
}
|
||||
id, err := d.SaveNotificationChannel(ctx, ch)
|
||||
if err != nil {
|
||||
t.Fatalf("新建失败: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { d.Exec(`DELETE FROM notification_channels WHERE id=$1`, id) })
|
||||
|
||||
got, err := d.NotificationChannelByID(ctx, id)
|
||||
if err != nil {
|
||||
t.Fatalf("读取失败: %v", err)
|
||||
}
|
||||
if got.Mode != NotifyModeDigest || got.RatePerMin != 42 || got.Name != "CRUD 往返" {
|
||||
t.Fatalf("往返字段不一致: %+v", got)
|
||||
}
|
||||
if !got.IsEnabled() {
|
||||
t.Fatal("默认应为启用")
|
||||
}
|
||||
var cfg map[string]any
|
||||
if err := json.Unmarshal(got.Config, &cfg); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if cfg["host"] != "smtp.example.com" {
|
||||
t.Fatalf("配置未正确落库: %v", cfg)
|
||||
}
|
||||
var filter notify.Filter
|
||||
if err := json.Unmarshal(got.Filter, &filter); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if filter.MinSeverity != "medium" || !filter.OnStatusChange {
|
||||
t.Fatalf("过滤条件未正确落库: %+v", filter)
|
||||
}
|
||||
|
||||
// 更新后再读。
|
||||
got.Name = "改名了"
|
||||
off := false
|
||||
got.Enabled = &off
|
||||
if _, err := d.SaveNotificationChannel(ctx, got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
after, err := d.NotificationChannelByID(ctx, id)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if after.Name != "改名了" || after.IsEnabled() {
|
||||
t.Fatalf("更新未生效: %+v", after)
|
||||
}
|
||||
|
||||
// 删除后应报「不存在」而不是静默成功。
|
||||
if err := d.DeleteNotificationChannel(ctx, id); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := d.NotificationChannelByID(ctx, id); err != ErrNotificationChannelNotFound {
|
||||
t.Fatalf("期望 ErrNotificationChannelNotFound，得到 %v", err)
|
||||
}
|
||||
if err := d.DeleteNotificationChannel(ctx, id); err != ErrNotificationChannelNotFound {
|
||||
t.Fatalf("重复删除应报不存在，得到 %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSaveNotificationChannelDefaultsRate 覆盖「未指定限流时套用渠道默认值」。
|
||||
// 漏掉这一步会让新渠道默认不限流，直接把平台打到限流。
|
||||
func TestSaveNotificationChannelDefaultsRate(t *testing.T) {
|
||||
d := notifyTestDB(t)
|
||||
ctx := context.Background()
|
||||
ch := &NotificationChannel{
|
||||
Name: "默认限流", Kind: notify.KindDingTalk,
|
||||
Config: json.RawMessage(`{"webhook":"https://example.com/h"}`),
|
||||
}
|
||||
id, err := d.SaveNotificationChannel(ctx, ch)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { d.Exec(`DELETE FROM notification_channels WHERE id=$1`, id) })
|
||||
got, err := d.NotificationChannelByID(ctx, id)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.RatePerMin != 20 {
|
||||
t.Fatalf("钉钉默认限流应为 20，得到 %d", got.RatePerMin)
|
||||
}
|
||||
if got.Mode != NotifyModeRealtime {
|
||||
t.Fatalf("默认模式应为 realtime，得到 %s", got.Mode)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDeleteChannelCascadesDeliveries 锁住外键行为：渠道删除后其投递历史一并消失
|
||||
// （配置都没了，历史无从解读），但事件本身要留下——它可能还被别的渠道引用。
|
||||
func TestDeleteChannelCascadesDeliveries(t *testing.T) {
|
||||
d := notifyTestDB(t)
|
||||
ctx := context.Background()
|
||||
ch := newTestChannel(t, d, notify.KindDingTalk, NotifyModeRealtime, `{}`)
|
||||
ev := addTestEvent(t, d, notify.EventFindingCreated, 9101, notify.Snapshot{Severity: "high"})
|
||||
if _, _, err := d.FanOutPendingEvents(ctx, 100); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var before int
|
||||
if err := d.QueryRow(`SELECT count(*) FROM notification_deliveries WHERE channel_id=$1`, ch.ID).Scan(&before); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if before == 0 {
|
||||
t.Fatal("前置条件不成立：未产生投递")
|
||||
}
|
||||
if err := d.DeleteNotificationChannel(ctx, ch.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var after int
|
||||
if err := d.QueryRow(`SELECT count(*) FROM notification_deliveries WHERE channel_id=$1`, ch.ID).Scan(&after); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if after != 0 {
|
||||
t.Fatalf("渠道删除后其投递应级联删除，仍有 %d 条", after)
|
||||
}
|
||||
var evExists bool
|
||||
if err := d.QueryRow(`SELECT EXISTS(SELECT 1 FROM notification_events WHERE id=$1)`, ev).Scan(&evExists); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !evExists {
|
||||
t.Fatal("删渠道不应连带删除事件本身")
|
||||
}
|
||||
}
|
||||
+100
@@ -1284,3 +1284,103 @@ ALTER TABLE task_intercept_rules ADD COLUMN IF NOT EXISTS action TEXT NOT NULL D
|
||||
DROP TRIGGER IF EXISTS trg_task_intercept_rules_upd ON task_intercept_rules;
|
||||
CREATE TRIGGER trg_task_intercept_rules_upd BEFORE UPDATE ON task_intercept_rules
|
||||
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
|
||||
|
||||
-- =====================================================================
|
||||
-- M. 漏洞 IM 推送
|
||||
--
|
||||
-- 三张表刻意分开，核心是**爆炸半径**：写漏洞的那个事务(RecordFindingTx，
|
||||
-- 持任务行锁)只允许做一次盲 INSERT，不读渠道表、不跑用户的过滤规则。否则
|
||||
-- 一条配错的 webhook 过滤条件就能污染/中止事务，导致漏洞存不进去。
|
||||
--
|
||||
-- notification_channels 渠道实例配置(可变、含凭据、UI 管理)
|
||||
-- notification_events 事件事实(写漏洞事务内盲插，含渲染快照)
|
||||
-- notification_deliveries 投递任务(事务外 fan-out 产生，承载状态/重试/批次)
|
||||
-- =====================================================================
|
||||
|
||||
-- 渠道实例：同一 kind 可配任意多个(如「应急群」「日常群」各一个钉钉机器人)。
|
||||
-- kind 取值由 server 侧白名单校验，不加 CHECK：与 findings.status 同理，
|
||||
-- 后续加渠道不应要求改表结构。
|
||||
CREATE TABLE IF NOT EXISTS notification_channels (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
-- dingtalk 钉钉 / feishu 飞书 / wecom 企业微信 / webhook 通用 / telegram / email
|
||||
kind TEXT NOT NULL,
|
||||
enabled BOOLEAN NOT NULL DEFAULT true,
|
||||
-- 凭据(明文存储，UI 掩码回显；见 server 侧 maskChannelSecrets)。六种渠道字段差异极大，
|
||||
-- 统一 JSONB + Go 侧按 kind 严格校验，避免为每渠道加一堆 NULL 列：
|
||||
-- dingtalk {webhook,secret}
|
||||
-- feishu {webhook,secret}
|
||||
-- wecom {webhook}
|
||||
-- webhook {url,method,content_type,headers{},body_template}
|
||||
-- telegram {bot_token,chat_id,base_url}
|
||||
-- email {host,port,username,password,from,to[],tls}
|
||||
config JSONB NOT NULL DEFAULT '{}',
|
||||
-- 推送时机：realtime 命中即推 / digest 进批次按全局周期汇总成一条。
|
||||
mode TEXT NOT NULL DEFAULT 'realtime',
|
||||
-- 过滤条件，字段全部可选(缺省=不过滤)：
|
||||
-- min_severity ''|low|medium|high|critical
|
||||
-- task_ids/asset_ids 空数组=不限；非空则须交集非空
|
||||
-- vulnclass_include/exclude 关键词数组(大小写不敏感子串)；include 空=全收
|
||||
-- on_status_change bool，仅 realtime 模式有意义
|
||||
filter JSONB NOT NULL DEFAULT '{}',
|
||||
-- 每分钟投递上限；0=不限流。默认 20 对齐钉钉/企微官方硬限。
|
||||
-- 超限不丢消息，只把投递推迟到下一个 tick。
|
||||
rate_per_min INTEGER NOT NULL DEFAULT 20,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
DROP TRIGGER IF EXISTS trg_notification_channels_upd ON notification_channels;
|
||||
CREATE TRIGGER trg_notification_channels_upd BEFORE UPDATE ON notification_channels
|
||||
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
|
||||
|
||||
-- 事件事实。由 RecordFindingTx / 状态变更事务**同事务**写入，保证「漏洞落库」
|
||||
-- 与「推送任务存在」原子一致——不存在提交成功但没入队、消息永久丢失的窗口。
|
||||
-- snapshot 刻意冗余：漏洞事后会被改名/改级别/改状态，推送内容应反映「事发当时」，
|
||||
-- 且 fan-out 与渲染不必回查 findings/tasks/assets 多张表。
|
||||
-- finding 删除后事件不级联删除：与 findings 表「任务删除仍独立留存」的语义一致。
|
||||
CREATE TABLE IF NOT EXISTS notification_events (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
-- finding_created | finding_status_changed
|
||||
kind TEXT NOT NULL,
|
||||
finding_id BIGINT NOT NULL,
|
||||
snapshot JSONB NOT NULL,
|
||||
-- fan-out 幂等标记：dispatcher 按此列取待分派事件，处理完置 true。
|
||||
-- 用列而非删行，以便投递历史能回溯到事件。
|
||||
fanned_out BOOLEAN NOT NULL DEFAULT false,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_notification_events_pending
|
||||
ON notification_events(id) WHERE NOT fanned_out;
|
||||
|
||||
-- 投递任务：一条事件 × 一个启用渠道 = 一行。fan-out 在事务外做，所以渠道
|
||||
-- 后开不会补历史(与 agent_triggers 的「迟开 trigger 不补历史」语义一致，
|
||||
-- 避免启用渠道时一次性刷屏历史积压)。
|
||||
-- channel_id 级联删除：渠道配置都没了，其投递历史无意义。
|
||||
CREATE TABLE IF NOT EXISTS notification_deliveries (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
event_id BIGINT NOT NULL REFERENCES notification_events(id) ON DELETE CASCADE,
|
||||
channel_id BIGINT NOT NULL REFERENCES notification_channels(id) ON DELETE CASCADE,
|
||||
-- pending 待发 / sent 已发 / failed 重试耗尽(可手动重发) / skipped 渠道停用或批次取消
|
||||
-- pending 待发 / sending 已被某 dispatcher 领取(租约未到期) / sent 已发 /
|
||||
-- failed 重试耗尽或永久失败(可手动重发) / skipped 渠道停用。取值不加 CHECK，
|
||||
-- 与 findings.status 同理，由 server 侧白名单校验。
|
||||
state TEXT NOT NULL DEFAULT 'pending',
|
||||
attempts INTEGER NOT NULL DEFAULT 0,
|
||||
-- 兼作「下次可领取时间」与「租约到期时间」：领取时把它推到未来即构成租约，
|
||||
-- 于是「租约未到期」与「未到重试时间」共用同一个条件表达，不需要额外的
|
||||
-- lease_until 列。进程崩溃留下的 sending 行会因租约到期被下一轮重新领取。
|
||||
next_attempt_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
last_error TEXT NOT NULL DEFAULT '',
|
||||
-- digest 模式同批次共享；realtime 恒为 NULL。整批渲染成一条消息后一起置 sent。
|
||||
batch_id BIGINT,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
sent_at TIMESTAMPTZ
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_notification_deliveries_due
|
||||
ON notification_deliveries(next_attempt_at) WHERE state='pending';
|
||||
CREATE INDEX IF NOT EXISTS idx_notification_deliveries_history
|
||||
ON notification_deliveries(id DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_notification_deliveries_batch
|
||||
ON notification_deliveries(batch_id) WHERE batch_id IS NOT NULL;
|
||||
CREATE INDEX IF NOT EXISTS idx_notification_deliveries_channel
|
||||
ON notification_deliveries(channel_id, id DESC);
|
||||
Reference in new issue
Block a user