Merge branch 'upstream/master' into master and preserve domain extensions

This commit is contained in:
Bot
2026-08-31 02:46:42 +08:00
758 changed files with 89269 additions and 1366 deletions
+13 -8
View File
@@ -149,13 +149,13 @@ func checkStatus() {
role = u.Role
}
UserLock.RUnlock()
if alert.UserID != server.UserID && !role.IsAdmin() {
if alert.UserID != server.GetUserID() && !role.IsAdmin() {
continue
}
alertsStore[alert.ID][server.ID] = append(alertsStore[alert.
ID][server.ID], alert.Snapshot(AlertsCycleTransferStatsStore[alert.ID], server, DB))
// 发送通知,分为触发报警和恢复通知
max, passed := alert.Check(alertsStore[alert.ID][server.ID])
_, passed := alert.Check(alertsStore[alert.ID][server.ID])
// 保存当前服务器状态信息
curServer := model.Server{}
copier.Copy(&curServer, server)
@@ -167,7 +167,7 @@ func checkStatus() {
alertsPrevState[alert.ID][server.ID] = _RuleCheckFail
message := fmt.Sprintf("[%s] %s(%s) %s", Localizer.T("Incident"),
server.Name, IPDesensitize(server.GeoIP.IP.Join()), alert.Name)
go CronShared.SendTriggerTasks(alert.FailTriggerTasks, curServer.ID)
go CronShared.SendTriggerTasks(alert.FailTriggerTasks, curServer.ID, alert.UserID)
go NotificationShared.SendNotification(alert.NotificationGroupID, message, NotificationMuteLabel.ServerIncident(server.ID, alert.ID), &curServer)
// 清除恢复通知的静音缓存
NotificationShared.UnMuteNotification(alert.NotificationGroupID, NotificationMuteLabel.ServerIncidentResolved(server.ID, alert.ID))
@@ -177,17 +177,22 @@ func checkStatus() {
if alertsPrevState[alert.ID][server.ID] == _RuleCheckFail {
message := fmt.Sprintf("[%s] %s(%s) %s", Localizer.T("Resolved"),
server.Name, IPDesensitize(server.GeoIP.IP.Join()), alert.Name)
go CronShared.SendTriggerTasks(alert.RecoverTriggerTasks, curServer.ID)
go CronShared.SendTriggerTasks(alert.RecoverTriggerTasks, curServer.ID, alert.UserID)
go NotificationShared.SendNotification(alert.NotificationGroupID, message, NotificationMuteLabel.ServerIncidentResolved(server.ID, alert.ID), &curServer)
// 清除失败通知的静音缓存
NotificationShared.UnMuteNotification(alert.NotificationGroupID, NotificationMuteLabel.ServerIncident(server.ID, alert.ID))
}
alertsPrevState[alert.ID][server.ID] = _RuleCheckPass
}
// 清理旧数据
if max > 0 && max < len(alertsStore[alert.ID][server.ID]) {
index := len(alertsStore[alert.ID][server.ID]) - max
alertsStore[alert.ID][server.ID] = alertsStore[alert.ID][server.ID][index:]
// 清理旧数据:保留窗口由规则定义决定(各规则 Duration 的最大值),
// 而非 Check 的判定结果。window==0 表示没有任何有效规则需要回看历史
// (例如全部 Duration<=0),此时清空采样避免切片无限增长。
window := alert.RetentionWindow()
samples := alertsStore[alert.ID][server.ID]
if window <= 0 {
alertsStore[alert.ID][server.ID] = samples[:0]
} else if window < len(samples) {
alertsStore[alert.ID][server.ID] = samples[len(samples)-window:]
}
}
}
+119
View File
@@ -0,0 +1,119 @@
package singleton
import (
"testing"
"github.com/nezhahq/nezha/model"
)
// notifyDecision replays the exact send-gate from checkStatus (lines 164-186)
// for a single (alert, server) pair: given the current Check verdict and the
// previous stored state, it reports whether an incident or recovery notification
// would be dispatched and what the next stored state becomes. Kept in lockstep
// with checkStatus so the end-to-end "does it actually notify" path is testable
// without the DB/global singletons checkStatus pulls in.
func notifyDecision(triggerMode uint8, passed bool, prev uint8) (incident, recover bool, next uint8) {
if !passed {
if triggerMode == model.ModeAlwaysTrigger || prev != _RuleCheckFail {
return true, false, _RuleCheckFail
}
return false, false, _RuleCheckFail
}
if prev == _RuleCheckFail {
return false, true, _RuleCheckPass
}
return false, false, _RuleCheckPass
}
// driveCheckStatus simulates the checkStatus tick loop end-to-end: each tick
// appends one sample, runs the real Check, applies the real RetentionWindow
// trim, then runs the send-gate. It returns how many incident notifications
// would have been dispatched across all ticks and the final sample window size.
func driveCheckStatus(rule *model.AlertRule, triggerMode uint8, ticks int, sample []bool) (incidents int, finalWindow int) {
incidents, finalWindow, _ = driveCheckStatusCap(rule, triggerMode, ticks, sample)
return incidents, finalWindow
}
// driveCheckStatusCap additionally reports the peak length and capacity the
// sample slice ever reached, so tests can assert memory stays bounded.
func driveCheckStatusCap(rule *model.AlertRule, triggerMode uint8, ticks int, sample []bool) (incidents, finalLen, peakCap int) {
var samples [][]bool
prev := uint8(_RuleCheckNoData)
for i := 0; i < ticks; i++ {
samples = append(samples, append([]bool(nil), sample...))
_, passed := rule.Check(samples)
w := rule.RetentionWindow()
if w <= 0 {
samples = samples[:0]
} else if w < len(samples) {
samples = samples[len(samples)-w:]
}
if cap(samples) > peakCap {
peakCap = cap(samples)
}
incident, _, next := notifyDecision(triggerMode, passed, prev)
if incident {
incidents++
}
prev = next
}
return incidents, len(samples), peakCap
}
// TestCheckStatus_GeneralRuleFiresIncident is the end-to-end guard for the
// regression: a Duration:10 rule on a server that fails every tick must,
// after the window fills, reach passed=false and actually dispatch an incident
// notification. Before the fix the window was wiped each tick, so passed never
// became false and zero notifications were sent.
func TestCheckStatus_GeneralRuleFiresIncident(t *testing.T) {
rule := &model.AlertRule{Rules: []*model.Rule{{Type: "cpu", Duration: 10}}}
t.Run("AlwaysTrigger fires repeatedly once window fills", func(t *testing.T) {
incidents, window := driveCheckStatus(rule, model.ModeAlwaysTrigger, 30, []bool{false})
if window < 10 {
t.Fatalf("window never filled: got %d want >= 10", window)
}
if incidents == 0 {
t.Fatalf("AlwaysTrigger rule never dispatched an incident notification")
}
})
t.Run("OnetimeTrigger fires exactly once", func(t *testing.T) {
incidents, _ := driveCheckStatus(rule, model.ModeOnetimeTrigger, 30, []bool{false})
if incidents != 1 {
t.Fatalf("OnetimeTrigger must dispatch exactly one incident, got %d", incidents)
}
})
}
// TestCheckStatus_HealthyServerStaysSilent guards the other direction: a server
// passing every tick must never dispatch an incident.
func TestCheckStatus_HealthyServerStaysSilent(t *testing.T) {
rule := &model.AlertRule{Rules: []*model.Rule{{Type: "cpu", Duration: 10}}}
incidents, _ := driveCheckStatus(rule, model.ModeAlwaysTrigger, 30, []bool{true})
if incidents != 0 {
t.Fatalf("a healthy server must never trigger an incident, got %d", incidents)
}
}
// TestCheckStatus_SampleMemoryBounded pins the no-memory-leak invariant: no
// matter how many ticks run, the per-(alert,server) sample slice length and
// capacity stay bounded by the rule's retention window, never growing with
// elapsed time. Runs far more ticks than the window to expose any unbounded
// growth.
func TestCheckStatus_SampleMemoryBounded(t *testing.T) {
const duration = 10
rule := &model.AlertRule{Rules: []*model.Rule{{Type: "cpu", Duration: duration}}}
_, finalLen, peakCap := driveCheckStatusCap(rule, model.ModeAlwaysTrigger, 100000, []bool{false})
if finalLen > duration {
t.Fatalf("sample length exceeded retention window after many ticks: got %d want <= %d", finalLen, duration)
}
// append grows capacity geometrically; with length pinned at window+1 the
// backing array stabilises at a small constant. A generous 4x window bound
// catches any reintroduced unbounded growth without being flaky.
if peakCap > duration*4 {
t.Fatalf("sample capacity grew unbounded: peak cap %d exceeds 4x window %d", peakCap, duration*4)
}
}
@@ -0,0 +1,57 @@
package singleton
import (
"path/filepath"
"testing"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
"github.com/nezhahq/nezha/model"
)
func setupCleanMonitorHistoryTestDB(t *testing.T) {
t.Helper()
previousDB := DB
var err error
DB, err = gorm.Open(openSQLiteDialector(filepath.Join(t.TempDir(), "dashboard.sqlite")), &gorm.Config{})
require.NoError(t, err)
sqlDB, err := DB.DB()
require.NoError(t, err)
t.Cleanup(func() {
DB = previousDB
if err := sqlDB.Close(); err != nil {
t.Errorf("close transfer cleanup test database: %v", err)
}
})
require.NoError(t, DB.AutoMigrate(&model.Server{}, &model.Transfer{}, &model.AlertRule{}))
require.NoError(t, DB.Exec("INSERT INTO servers (id, name, uuid) VALUES (1, 'server', 'clean-monitor-history-test')").Error)
}
func TestCleanMonitorHistoryWithoutRulesDeletesAllTransfers(t *testing.T) {
setupCleanMonitorHistoryTestDB(t)
require.NoError(t, DB.Create(&model.Transfer{ServerID: 1, In: 1}).Error)
CleanMonitorHistory()
var count int64
require.NoError(t, DB.Model(&model.Transfer{}).Count(&count).Error)
require.Zero(t, count)
}
func TestCleanMonitorHistoryPreservesTransfersWhenAlertRulesCannotBeLoaded(t *testing.T) {
setupCleanMonitorHistoryTestDB(t)
require.NoError(t, DB.Create(&model.Transfer{ServerID: 1, In: 1}).Error)
require.NoError(t, DB.Exec("INSERT INTO alert_rules (id, name, rules_raw, fail_trigger_tasks_raw, recover_trigger_tasks_raw) VALUES (1, 'broken', '{', '[]', '[]')").Error)
var alerts []model.AlertRule
require.Error(t, DB.Find(&alerts).Error, "precondition: malformed rules_raw must fail AlertRule.AfterFind")
CleanMonitorHistory()
var count int64
require.NoError(t, DB.Model(&model.Transfer{}).Count(&count).Error)
require.EqualValues(t, 1, count)
}
+8
View File
@@ -1,6 +1,7 @@
package singleton
import (
"log"
"strconv"
"strings"
@@ -26,6 +27,13 @@ func InitConfigFromPath(path string) error {
if err != nil {
return err
}
rotated, err := Conf.RotateJWTSecretKeyIfNeeded(Version)
if err != nil {
return err
}
if rotated {
log.Printf("NEZHA>> Rotated jwt_secret_key for dashboard version %s", Version)
}
Conf.updateIgnoredIPNotificationID()
Conf.Oauth2Providers = utils.MapKeysToSlice(Conf.Oauth2)
+54
View File
@@ -0,0 +1,54 @@
package singleton
import (
"os"
"strings"
"testing"
"github.com/nezhahq/nezha/model"
)
func TestInitConfigFromPathRotatesJWTSecretKey(t *testing.T) {
file, err := os.CreateTemp(t.TempDir(), "nezha-config-*.yaml")
if err != nil {
t.Fatalf("create temp config: %v", err)
}
if _, err := file.WriteString("jwt_secret_key: leaked-secret\nagent_secret_key: agent-secret\njwt_secret_key_last_rotated_version: v2.0.12\n"); err != nil {
t.Fatalf("write temp config: %v", err)
}
if err := file.Close(); err != nil {
t.Fatalf("close temp config: %v", err)
}
originalConf := Conf
originalVersion := Version
originalTemplates := FrontendTemplates
Version = "v2.0.13"
FrontendTemplates = nil
t.Cleanup(func() {
Conf = originalConf
Version = originalVersion
FrontendTemplates = originalTemplates
})
if err := InitConfigFromPath(file.Name()); err != nil {
t.Fatalf("init config: %v", err)
}
if Conf.JWTSecretKey == "leaked-secret" {
t.Fatal("jwt_secret_key was not rotated")
}
if Conf.JWTSecretKeyLastRotatedVersion != model.JWTSecretKeyRotationBaselineVersion {
t.Fatalf("jwt secret key marker = %q, want %q", Conf.JWTSecretKeyLastRotatedVersion, model.JWTSecretKeyRotationBaselineVersion)
}
saved, err := os.ReadFile(file.Name())
if err != nil {
t.Fatalf("read saved config: %v", err)
}
if strings.Contains(string(saved), "leaked-secret") {
t.Fatalf("saved config still contains leaked jwt_secret_key: %s", saved)
}
if !strings.Contains(string(saved), "jwt_secret_key_last_rotated_version: v2.0.13") {
t.Fatalf("saved config did not persist jwt secret key marker: %s", saved)
}
}
+175 -9
View File
@@ -5,6 +5,8 @@ import (
"fmt"
"slices"
"strings"
"sync"
"time"
"github.com/jinzhu/copier"
@@ -15,9 +17,25 @@ import (
pb "github.com/nezhahq/nezha/proto"
)
const alertTriggerCronResultAuthorizationTTL = 24 * time.Hour
type CronClass struct {
class[uint64, *model.Cron]
*cron.Cron
pendingAlertTriggerTasksMu sync.Mutex
pendingAlertTriggerTasks map[uint64]map[uint64][]time.Time
closeOnce sync.Once
}
// Close stops the scheduler and joins every job before a test restores globals.
// The embedded cron.Stop only exposes the completion context; callers must await it.
func (c *CronClass) Close() {
if c == nil || c.Cron == nil {
return
}
c.closeOnce.Do(func() {
<-c.Cron.Stop().Done()
})
}
func NewCronClass() *CronClass {
@@ -64,7 +82,8 @@ func NewCronClass() *CronClass {
list: list,
sortedList: sortedList,
},
Cron: cronx,
Cron: cronx,
pendingAlertTriggerTasks: make(map[uint64]map[uint64][]time.Time),
}
}
@@ -78,6 +97,7 @@ func (c *CronClass) Update(cr *model.Cron) {
delete(c.list, cr.ID)
c.list[cr.ID] = cr
c.listMu.Unlock()
c.deleteAlertTriggerCronResultAuthorizations([]uint64{cr.ID})
c.sortList()
}
@@ -92,6 +112,7 @@ func (c *CronClass) Delete(idList []uint64) {
delete(c.list, id)
}
c.listMu.Unlock()
c.deleteAlertTriggerCronResultAuthorizations(idList)
c.sortList()
}
@@ -110,11 +131,11 @@ func (c *CronClass) sortList() {
c.sortedList = sortedList
}
func (c *CronClass) SendTriggerTasks(taskIDs []uint64, triggerServer uint64) {
func (c *CronClass) SendTriggerTasks(taskIDs []uint64, triggerServer uint64, triggerOwner uint64) {
c.listMu.RLock()
var cronLists []*model.Cron
for _, taskID := range taskIDs {
if c, ok := c.list[taskID]; ok {
if c, ok := c.list[taskID]; ok && cronCanBeTriggeredByOwner(c, triggerOwner) {
cronLists = append(cronLists, c)
}
}
@@ -126,6 +147,117 @@ func (c *CronClass) SendTriggerTasks(taskIDs []uint64, triggerServer uint64) {
}
}
func cronCanBeTriggeredByOwner(cr *model.Cron, triggerOwner uint64) bool {
return cr.UserID == triggerOwner || userIsAdmin(triggerOwner)
}
func CanReportCronResult(cr *model.Cron, reporter *model.Server) bool {
if cr == nil || reporter == nil || !cronCanSendToServer(cr, reporter) {
return false
}
if cr.Cover == model.CronCoverAll {
return !slices.Contains(cr.Servers, reporter.ID)
}
if cr.Cover == model.CronCoverIgnoreAll {
return slices.Contains(cr.Servers, reporter.ID)
}
if cr.Cover == model.CronCoverAlertTrigger {
return CronShared != nil && CronShared.consumeAlertTriggerCronResult(cr.ID, reporter.ID)
}
return false
}
func (c *CronClass) reserveAlertTriggerCronResult(cronID uint64, serverID uint64) {
c.pendingAlertTriggerTasksMu.Lock()
defer c.pendingAlertTriggerTasksMu.Unlock()
now := time.Now()
c.pruneExpiredAlertTriggerCronResultsLocked(now)
if c.pendingAlertTriggerTasks == nil {
c.pendingAlertTriggerTasks = make(map[uint64]map[uint64][]time.Time)
}
if c.pendingAlertTriggerTasks[cronID] == nil {
c.pendingAlertTriggerTasks[cronID] = make(map[uint64][]time.Time)
}
c.pendingAlertTriggerTasks[cronID][serverID] = append(c.pendingAlertTriggerTasks[cronID][serverID], now.Add(alertTriggerCronResultAuthorizationTTL))
}
func (c *CronClass) revokeAlertTriggerCronResult(cronID uint64, serverID uint64) {
c.pendingAlertTriggerTasksMu.Lock()
defer c.pendingAlertTriggerTasksMu.Unlock()
serverTasks := c.pendingAlertTriggerTasks[cronID]
expiresAtList := serverTasks[serverID]
if len(expiresAtList) == 0 {
return
}
expiresAtList = expiresAtList[:len(expiresAtList)-1]
if len(expiresAtList) == 0 {
delete(serverTasks, serverID)
} else {
serverTasks[serverID] = expiresAtList
}
if len(serverTasks) == 0 {
delete(c.pendingAlertTriggerTasks, cronID)
}
}
func (c *CronClass) consumeAlertTriggerCronResult(cronID uint64, serverID uint64) bool {
c.pendingAlertTriggerTasksMu.Lock()
defer c.pendingAlertTriggerTasksMu.Unlock()
c.pruneExpiredAlertTriggerCronResultsLocked(time.Now())
return c.consumeAlertTriggerCronResultLocked(cronID, serverID)
}
func (c *CronClass) consumeAlertTriggerCronResultLocked(cronID uint64, serverID uint64) bool {
serverTasks := c.pendingAlertTriggerTasks[cronID]
expiresAtList := serverTasks[serverID]
if len(expiresAtList) == 0 {
return false
}
expiresAtList = expiresAtList[1:]
if len(expiresAtList) == 0 {
delete(serverTasks, serverID)
} else {
serverTasks[serverID] = expiresAtList
}
if len(serverTasks) == 0 {
delete(c.pendingAlertTriggerTasks, cronID)
}
return true
}
func (c *CronClass) pruneExpiredAlertTriggerCronResultsLocked(now time.Time) {
for cronID, serverTasks := range c.pendingAlertTriggerTasks {
for serverID, expiresAtList := range serverTasks {
validExpiresAtList := expiresAtList[:0]
for _, expiresAt := range expiresAtList {
if expiresAt.After(now) {
validExpiresAtList = append(validExpiresAtList, expiresAt)
}
}
if len(validExpiresAtList) == 0 {
delete(serverTasks, serverID)
} else {
serverTasks[serverID] = validExpiresAtList
}
}
if len(serverTasks) == 0 {
delete(c.pendingAlertTriggerTasks, cronID)
}
}
}
func (c *CronClass) deleteAlertTriggerCronResultAuthorizations(cronIDs []uint64) {
c.pendingAlertTriggerTasksMu.Lock()
defer c.pendingAlertTriggerTasksMu.Unlock()
for _, cronID := range cronIDs {
delete(c.pendingAlertTriggerTasks, cronID)
}
}
func ManualTrigger(cr *model.Cron) {
CronTrigger(cr)()
}
@@ -141,12 +273,21 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() {
return
}
if s, ok := ServerShared.Get(triggerServer[0]); ok {
if s.TaskStream != nil {
s.TaskStream.Send(&pb.Task{
if !cronCanSendToServer(cr, s) {
return
}
if s.GetTaskStream() != nil {
cronShared := CronShared
if cronShared != nil {
cronShared.reserveAlertTriggerCronResult(cr.ID, s.ID)
}
if err := s.SendTask(&pb.Task{
Id: cr.ID,
Data: cr.Command,
Type: model.TaskTypeCommand,
})
}); err != nil && cronShared != nil {
cronShared.revokeAlertTriggerCronResult(cr.ID, s.ID)
}
} else {
// 保存当前服务器状态信息
curServer := model.Server{}
@@ -157,15 +298,24 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() {
return
}
for _, s := range ServerShared.Range {
// 先在锁内快照 server 列表再逐个 SendTaskServerShared.Range 会在整个
// 回调期间持 listMu.RLock,而 SendTask 走阻塞 gRPC,一个卡死的 agent
// 会让需要写锁的 server 编辑/删除被拖死。GetList 克隆后即释放锁。
for _, s := range ServerShared.GetList() {
if s == nil {
continue
}
if !cronCanSendToServer(cr, s) {
continue
}
if cr.Cover == model.CronCoverAll && crIgnoreMap[s.ID] {
continue
}
if cr.Cover == model.CronCoverIgnoreAll && !crIgnoreMap[s.ID] {
continue
}
if s.TaskStream != nil {
s.TaskStream.Send(&pb.Task{
if s.GetTaskStream() != nil {
_ = s.SendTask(&pb.Task{
Id: cr.ID,
Data: cr.Command,
Type: model.TaskTypeCommand,
@@ -179,3 +329,19 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() {
}
}
}
func cronCanSendToServer(cr *model.Cron, server *model.Server) bool {
return cr.UserID == server.GetUserID() || userIsAdmin(cr.UserID)
}
func userIsAdmin(userID uint64) bool {
if userID == 0 {
return true
}
UserLock.RLock()
defer UserLock.RUnlock()
userInfo, ok := UserInfoMap[userID]
return ok && userInfo.Role.IsAdmin()
}
@@ -0,0 +1,97 @@
package singleton
import (
"context"
"testing"
"time"
"github.com/robfig/cron/v3"
"github.com/stretchr/testify/require"
)
const lifecycleTestTimeout = time.Second
func TestCronClassClose_waitsForRunningJobs(t *testing.T) {
// Given
started := make(chan struct{})
release := make(chan struct{})
events := make(chan string, 9)
cronClass := &CronClass{Cron: cron.New(cron.WithSeconds())}
_, err := cronClass.AddFunc("@every 1ns", func() {
defer func() { events <- "job" }()
close(started)
<-release
})
require.NoError(t, err)
cronClass.Start()
<-started
// When
closed := make(chan struct{})
for range 8 {
go func() {
cronClass.Close()
events <- "close"
closed <- struct{}{}
}()
}
// Then
close(release)
firstEvent := awaitCronLifecycleEvent(t, events, "cron lifecycle did not complete")
if firstEvent != "job" {
t.Fatalf("Close returned before the running cron job returned: first event=%q", firstEvent)
}
for range 8 {
awaitCronLifecycleSignal(t, closed, "concurrent Close call did not return")
}
for range 8 {
if event := awaitCronLifecycleEvent(t, events, "concurrent Close call did not complete"); event != "close" {
t.Fatalf("unexpected cron lifecycle event: %q", event)
}
}
}
func TestCronClassClose_isIdempotentAndNilSafe(t *testing.T) {
cronClass := &CronClass{Cron: cron.New(cron.WithSeconds())}
cronClass.Start()
closed := make(chan struct{})
for range 8 {
go func() {
cronClass.Close()
closed <- struct{}{}
}()
}
for range 8 {
awaitCronLifecycleSignal(t, closed, "concurrent Close call did not return")
}
cronClass.Close()
var nilCronClass *CronClass
nilCronClass.Close()
(&CronClass{}).Close()
}
func awaitCronLifecycleSignal(t *testing.T, signal <-chan struct{}, message string) {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), lifecycleTestTimeout)
defer cancel()
select {
case <-signal:
case <-ctx.Done():
t.Fatal(message)
}
}
func awaitCronLifecycleEvent(t *testing.T, events <-chan string, message string) string {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), lifecycleTestTimeout)
defer cancel()
select {
case event := <-events:
return event
case <-ctx.Done():
t.Fatal(message)
return ""
}
}
+20 -1
View File
@@ -3,6 +3,7 @@ package singleton
import (
"cmp"
"fmt"
"log"
"slices"
"github.com/libdns/cloudflare"
@@ -56,12 +57,30 @@ func (c *DDNSClass) Delete(idList []uint64) {
c.sortList()
}
func (c *DDNSClass) GetDDNSProvidersFromProfiles(profileId []uint64, ip *model.IP) ([]*ddns2.Provider, error) {
// profileOwnedByRealAdmin reports whether uid is a genuine admin user that
// may share its DDNS profiles globally. userIsAdmin(0) returns true as a
// "system resource" shortcut, but a profile with UserID==0 is a migration /
// default-value artifact, not an admin grant — sharing it with foreign server
// owners reopens GHSA-39g2-8x68-pmx8. A real admin always has a non-zero ID.
func profileOwnedByRealAdmin(uid uint64) bool {
return uid != 0 && userIsAdmin(uid)
}
// GHSA-39g2-8x68-pmx8: bind-time CheckPermission 对「不存在的 profile ID」放行,
// 攻击者可预绑定将来才会被受害者创建的自增 ID。worker 解析时必须按 ownerUID
// 重新校验归属,跳过非 server owner(且非管理员)所有的 profile。
func (c *DDNSClass) GetDDNSProvidersFromProfiles(profileId []uint64, ip *model.IP, ownerUID uint64) ([]*ddns2.Provider, error) {
profiles := make([]*model.DDNSProfile, 0, len(profileId))
c.listMu.RLock()
for _, id := range profileId {
if profile, ok := c.list[id]; ok {
if profile.UserID != ownerUID && !profileOwnedByRealAdmin(profile.UserID) {
// Fail-closed skip: an admin may bind a member-owned profile,
// but worker-time only runs same-owner or real-admin profiles.
log.Printf("NEZHA>> Skipping DDNS profile %d (owner %d) for server owner %d: not owned by server owner or a real admin", profile.ID, profile.UserID, ownerUID)
continue
}
profiles = append(profiles, profile)
} else {
c.listMu.RUnlock()
+121
View File
@@ -0,0 +1,121 @@
package singleton
import (
"testing"
"github.com/nezhahq/nezha/model"
)
// newDDNSClassForTest builds a DDNSClass backed by an in-memory profile map,
// mirroring the production cache layout without touching the database.
func newDDNSClassForTest(profiles ...*model.DDNSProfile) *DDNSClass {
list := make(map[uint64]*model.DDNSProfile, len(profiles))
for _, p := range profiles {
list[p.ID] = p
}
return &DDNSClass{
class: class[uint64, *model.DDNSProfile]{
list: list,
sortedList: profiles,
},
}
}
// GHSA-39g2-8x68-pmx8: a server owned by the attacker must not be able to
// drive a DDNS update through a DDNS profile owned by another (victim) user.
// The worker-time resolution must skip foreign-owned profiles.
func TestGetDDNSProvidersSkipsForeignOwnedProfile(t *testing.T) {
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember}, // attacker / server owner
200: {Role: model.RoleMember}, // victim / profile owner
})
victimProfile := &model.DDNSProfile{
Common: model.Common{ID: 1, UserID: 200},
Provider: model.ProviderDummy,
Name: "victim-profile",
AccessSecret: "victim-secret",
}
dc := newDDNSClassForTest(victimProfile)
providers, err := dc.GetDDNSProvidersFromProfiles([]uint64{1}, &model.IP{}, 100)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(providers) != 0 {
t.Fatalf("expected foreign-owned profile to be skipped, got %d provider(s)", len(providers))
}
}
// A server owner using their own DDNS profile must still resolve normally.
func TestGetDDNSProvidersAllowsOwnedProfile(t *testing.T) {
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
})
ownProfile := &model.DDNSProfile{
Common: model.Common{ID: 5, UserID: 100},
Provider: model.ProviderDummy,
Name: "own-profile",
}
dc := newDDNSClassForTest(ownProfile)
providers, err := dc.GetDDNSProvidersFromProfiles([]uint64{5}, &model.IP{}, 100)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(providers) != 1 {
t.Fatalf("expected owned profile to resolve, got %d provider(s)", len(providers))
}
}
// An admin-owned profile may be shared across servers (admin resources are
// global), so an admin profile resolves regardless of the server owner.
func TestGetDDNSProvidersAllowsAdminOwnedProfile(t *testing.T) {
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin}, // admin / profile owner
100: {Role: model.RoleMember},
})
adminProfile := &model.DDNSProfile{
Common: model.Common{ID: 9, UserID: 1},
Provider: model.ProviderDummy,
Name: "admin-profile",
}
dc := newDDNSClassForTest(adminProfile)
providers, err := dc.GetDDNSProvidersFromProfiles([]uint64{9}, &model.IP{}, 100)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(providers) != 1 {
t.Fatalf("expected admin-owned profile to resolve, got %d provider(s)", len(providers))
}
}
// GHSA-39g2-8x68-pmx8 (UserID==0 variant): userIsAdmin(0) returns true as a
// "system resource" shortcut, but a DDNS profile with UserID==0 is not a real
// admin grant — it is a migration/default-value artifact. A foreign server
// owner must NOT be able to drive an update through such a profile, so the
// worker must skip a UserID==0 profile that the caller does not own.
func TestGetDDNSProvidersSkipsUnownedZeroUserProfile(t *testing.T) {
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember}, // attacker / server owner
})
orphanProfile := &model.DDNSProfile{
Common: model.Common{ID: 3, UserID: 0},
Provider: model.ProviderDummy,
Name: "orphan-profile",
AccessSecret: "orphan-secret",
}
dc := newDDNSClassForTest(orphanProfile)
providers, err := dc.GetDDNSProvidersFromProfiles([]uint64{3}, &model.IP{}, 100)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(providers) != 0 {
t.Fatalf("expected UserID==0 foreign profile to be skipped, got %d provider(s)", len(providers))
}
}
+18 -4
View File
@@ -2,20 +2,34 @@
name: "OfficialAdmin"
repository: "https://github.com/nezhahq/admin-frontend"
author: "nezhahq"
version: "v2.0.6"
version: "v2.3.4"
is_admin: true
is_official: true
- path: "user-dist"
name: "Official"
repository: "https://github.com/hamster1963/nezha-dash-v2"
author: "hamster1963"
version: "v2.0.3"
version: "v2.4.2"
is_official: true
- path: "nezha-pixel-dist"
name: "Nezha-Pixel"
repository: "https://github.com/karllao/nezha-pixel"
author: "karllao"
version: "v1.6.0"
# Third-party user themes consume the opaque Server.PublicNote field. Theme
# maintainers must validate URL schemes (including after decoding) before using
# values such as customData.orderLink in href or window.open; the Dashboard
# backend and admin frontend do not execute those fields.
- path: "nazhua-dist"
name: "Nazhua"
repository: "https://github.com/hi2shark/nazhua"
author: "hi2hi"
version: "v0.9.1"
author: "hi2shark"
version: "v1.2.0"
- path: "aobobo-dist"
name: "Aobobo"
repository: "https://github.com/hi2shark/aobobo"
author: "hi2shark"
version: "v1.5.1"
- path: "nezha-ascii-dist"
name: "Nezha-ASCII"
repository: "https://github.com/hamster1963/nezha-ascii"
+55
View File
@@ -0,0 +1,55 @@
package singleton
import (
"log"
"time"
"github.com/nezhahq/nezha/model"
)
const (
JWTSessionGCSchedule = "@every 10m"
JWTSessionRevokedRetention = 24 * time.Hour
JWTSessionExpiredGrace = 1 * time.Hour
)
func StartJWTSessionGC() error {
if _, err := CronShared.AddFunc(JWTSessionGCSchedule, RunJWTSessionGC); err != nil {
return err
}
RunJWTSessionGC()
return nil
}
func RunJWTSessionGC() {
if DB == nil {
return
}
now := time.Now()
if err := DB.
Where("expires_at < ?", now.Add(-JWTSessionExpiredGrace)).
Delete(&model.JWTSession{}).Error; err != nil {
log.Printf("NEZHA>> JWTSession GC delete expired failed: %v", err)
}
if err := DB.
Where("revoked_at IS NOT NULL AND revoked_at < ?", now.Add(-JWTSessionRevokedRetention)).
Delete(&model.JWTSession{}).Error; err != nil {
log.Printf("NEZHA>> JWTSession GC delete revoked failed: %v", err)
}
}
func RevokeJWTSession(keyID string) error {
now := time.Now()
return DB.Model(&model.JWTSession{}).
Where("key_id = ? AND revoked_at IS NULL", keyID).
Update("revoked_at", &now).Error
}
func RevokeJWTSessionsByUser(userID uint64) error {
now := time.Now()
return DB.Model(&model.JWTSession{}).
Where("user_id = ? AND revoked_at IS NULL", userID).
Update("revoked_at", &now).Error
}
+70
View File
@@ -2,12 +2,81 @@ package singleton
import (
"cmp"
"net"
"net/netip"
"slices"
"strings"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/utils"
)
// GHSA-x6fg-52vr-hj4w: NAT 是 commonHandler,任意认证成员可创建。newHTTPandGRPCMux
// 在分发 dashboard/gRPC 之前先按 r.Host 命中 NAT,故成员若把 Domain 设成 dashboard
// 自身 host 即可抢占全局路由(disabled 触发 DoSenabled 把请求隧道到攻击者 agent)。
// 把 dashboard 的 InstallHost、ListenHost 以及运维声明的 ReservedHosts 列为
// 保留 hostcreate/update 时拒绝,启动建表时丢弃,确保补丁前已植入的恶意记录
// 在升级后不再生效。每个 host 拆成 hostname 后比较(忽略端口与大小写),反代/
// 默认端口下的端口变体也拦得住。反代部署时进程看不到对外域名,运维把它配进
// Conf.ReservedHosts(逗号分隔)即可让此处覆盖到公网入口。
func IsReservedDashboardHost(domain string) bool {
if Conf == nil {
return false
}
target := splitDashboardHostname(domain)
if target == "" {
return false
}
hosts := []string{Conf.InstallHost, Conf.DashboardHost, Conf.ListenHost}
hosts = append(hosts, strings.Split(Conf.ReservedHosts, ",")...)
for _, host := range hosts {
if reserved := splitDashboardHostname(host); reserved != "" && reserved == target {
return true
}
}
return false
}
// splitDashboardHostname 归一化为小写 hostname,对 bracketed IPv6[::1]、
// [::1]:8008)与裸 host:port 一视同仁,避免 candidate 与 reserved 解析形态不
// 一致导致漏拦。两类等价形态也必须收敛,否则 guard 放行而 r.Host 精确命中
// 仍能劫持路由:
// - DNS absolute name 的尾点(panel.example.com. 与 panel.example.com 指向同
// 一主机),去掉单个尾点;
// - IP literal 的压缩/展开写法(::1 与 0:0:0:0:0:0:0:1),用 netip 归一到
// 规范文本。
func splitDashboardHostname(host string) string {
host = strings.ToLower(strings.TrimSpace(host))
if host == "" {
return ""
}
if h, _, err := net.SplitHostPort(host); err == nil && h != "" {
host = h
} else {
host = strings.Trim(host, "[]")
}
host = strings.TrimSuffix(host, ".")
if addr, err := netip.ParseAddr(host); err == nil {
return addr.String()
}
return host
}
// filterReservedNATProfiles 丢弃 Domain 命中 dashboard 保留 host 的 NAT 记录,
// 让 NewNATClass 启动建表时不把补丁前植入的劫持记录加载进路由表。
func filterReservedNATProfiles(in []*model.NAT) []*model.NAT {
out := in[:0]
for _, profile := range in {
if profile == nil || IsReservedDashboardHost(profile.Domain) {
continue
}
out = append(out, profile)
}
return out
}
type NATClass struct {
class[string, *model.NAT]
@@ -18,6 +87,7 @@ func NewNATClass() *NATClass {
var sortedList []*model.NAT
DB.Find(&sortedList)
sortedList = filterReservedNATProfiles(sortedList)
list := make(map[string]*model.NAT, len(sortedList))
idToDomain := make(map[uint64]string, len(sortedList))
for _, profile := range sortedList {
+139
View File
@@ -0,0 +1,139 @@
package singleton
import (
"testing"
"github.com/nezhahq/nezha/model"
)
func withReservedHostConf(t *testing.T, c *model.Config) {
t.Helper()
original := Conf
Conf = &ConfigClass{Config: c}
t.Cleanup(func() { Conf = original })
}
// GHSA-x6fg-52vr-hj4w: the reserved-host check is the single source of truth
// for both the create/update guard and the startup cache filter. It must
// reject any NAT domain whose hostname collides with the dashboard's own
// InstallHost / ListenHost, regardless of port or case.
func TestIsReservedDashboardHost(t *testing.T) {
withReservedHostConf(t, &model.Config{
ConfigDashboard: model.ConfigDashboard{InstallHost: "dashboard.example:8008"},
ListenHost: "10.0.0.5",
ListenPort: 8008,
})
cases := []struct {
name string
domain string
want bool
}{
{"exact install host", "dashboard.example:8008", true},
{"install host case-insensitive", "Dashboard.Example:8008", true},
{"install host without port", "dashboard.example", true},
{"install host arbitrary port", "dashboard.example:8443", true},
{"listen host and port", "10.0.0.5:8008", true},
{"listen host bare", "10.0.0.5", true},
{"unrelated domain", "tunnel.member.example", false},
{"empty domain", "", false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
if got := IsReservedDashboardHost(tc.domain); got != tc.want {
t.Fatalf("IsReservedDashboardHost(%q) = %v, want %v", tc.domain, got, tc.want)
}
})
}
}
// GHSA-x6fg-52vr-hj4w (reverse-proxy coverage): InstallHost/ListenHost alone
// cannot cover a dashboard reached through a reverse proxy on a public domain
// that the dashboard process never sees. ReservedHosts lets the operator
// declare those extra hostnames (comma-separated) so members still cannot
// register a NAT domain that collides with the public entry point.
func TestIsReservedDashboardHostHonoursReservedHostsList(t *testing.T) {
withReservedHostConf(t, &model.Config{
ConfigDashboard: model.ConfigDashboard{
InstallHost: "internal.example:8008",
ReservedHosts: "panel.example.com, Admin.Example.COM:443 , ",
},
})
reserved := []string{
"panel.example.com",
"panel.example.com:8443",
"admin.example.com",
"ADMIN.EXAMPLE.COM:443",
"internal.example",
}
for _, d := range reserved {
if !IsReservedDashboardHost(d) {
t.Errorf("IsReservedDashboardHost(%q) = false, want true (declared reserved host)", d)
}
}
if IsReservedDashboardHost("tunnel.member.example") {
t.Error("unrelated member domain must not be reserved")
}
if IsReservedDashboardHost("") {
t.Error("empty domain must not be reserved")
}
}
// The startup cache must not load a NAT record whose domain is reserved, so a
// malicious record planted before the patch cannot keep hijacking dashboard
// routing after upgrade. filterReservedNATProfiles is the gate NewNATClass
// runs over the DB result set.
func TestFilterReservedNATProfilesDropsReserved(t *testing.T) {
withReservedHostConf(t, &model.Config{
ConfigDashboard: model.ConfigDashboard{InstallHost: "dashboard.example:8008"},
})
in := []*model.NAT{
{Common: model.Common{ID: 1}, Domain: "dashboard.example", Enabled: true},
{Common: model.Common{ID: 2}, Domain: "tunnel.member.example", Enabled: true},
{Common: model.Common{ID: 3}, Domain: "Dashboard.Example:9999", Enabled: false},
}
out := filterReservedNATProfiles(in)
if len(out) != 1 {
t.Fatalf("expected only the non-reserved profile to survive, got %d", len(out))
}
if out[0].Domain != "tunnel.member.example" {
t.Fatalf("surviving profile must be the member tunnel, got %q", out[0].Domain)
}
}
// GHSA-x6fg-52vr-hj4w (canonical-host coverage): the routing match is an exact
// lookup on r.Host, so a member who registers a NAT Domain that is a DNS/IP
// *equivalent* of the dashboard host — but a different literal string — still
// hijacks the matching r.Host. The guard must collapse the trailing DNS dot and
// the IPv6 compressed/expanded forms, or these variants slip past create/update.
func TestIsReservedDashboardHostCollapsesEquivalentForms(t *testing.T) {
withReservedHostConf(t, &model.Config{
ConfigDashboard: model.ConfigDashboard{
InstallHost: "panel.example.com",
ReservedHosts: "[::1]:8008",
},
})
reserved := []string{
"panel.example.com.", // trailing dot, no port
"panel.example.com.:8008", // trailing dot with port
"PANEL.EXAMPLE.COM.", // trailing dot, mixed case
"[0:0:0:0:0:0:0:1]:8008", // IPv6 expanded form of ::1
"::1", // IPv6 compressed, bare
"[::1]", // IPv6 compressed, bracketed
}
for _, d := range reserved {
if !IsReservedDashboardHost(d) {
t.Errorf("IsReservedDashboardHost(%q) = false, want true (equivalent of reserved host)", d)
}
}
if IsReservedDashboardHost("tunnel.member.example.") {
t.Error("unrelated member domain with trailing dot must not be reserved")
}
}
@@ -0,0 +1,940 @@
package singleton
import (
"context"
"net/http/httptest"
"slices"
"testing"
"time"
"github.com/gin-gonic/gin"
"github.com/patrickmn/go-cache"
"github.com/robfig/cron/v3"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"github.com/nezhahq/nezha/model"
pb "github.com/nezhahq/nezha/proto"
"google.golang.org/grpc/metadata"
)
type capturedTaskStream struct {
tasks chan *pb.Task
}
func newCapturedTaskStream() *capturedTaskStream {
return &capturedTaskStream{tasks: make(chan *pb.Task, 4)}
}
func (s *capturedTaskStream) Send(task *pb.Task) error {
s.tasks <- task
return nil
}
func (s *capturedTaskStream) Recv() (*pb.TaskResult, error) { return nil, context.Canceled }
func (s *capturedTaskStream) SetHeader(metadata.MD) error { return nil }
func (s *capturedTaskStream) SendHeader(metadata.MD) error { return nil }
func (s *capturedTaskStream) SetTrailer(metadata.MD) {}
func (s *capturedTaskStream) Context() context.Context { return context.Background() }
func (s *capturedTaskStream) SendMsg(any) error { return nil }
func (s *capturedTaskStream) RecvMsg(any) error { return context.Canceled }
// withTaskStream attaches a TaskStream to a freshly constructed Server using the
// new atomic accessor. The field itself is unexported (see Fix #12) precisely
// because direct struct-literal access invited torn interface reads on hot
// paths — tests use this helper rather than reaching in, mirroring production
// callsites.
func withTaskStream(s *model.Server, stream pb.NezhaService_RequestTaskServer) *model.Server {
s.SetTaskStream(stream)
return s
}
func replaceServerSharedForSecurityTest(t *testing.T, servers ...*model.Server) {
t.Helper()
original := ServerShared
serverClass := &ServerClass{
class: class[uint64, *model.Server]{
list: make(map[uint64]*model.Server),
},
uuidToID: make(map[string]uint64),
}
for _, server := range servers {
serverClass.list[server.ID] = server
}
ServerShared = serverClass
t.Cleanup(func() { ServerShared = original })
}
func replaceUserInfoMapForSecurityTest(t *testing.T, users map[uint64]model.UserInfo) {
t.Helper()
UserLock.Lock()
original := UserInfoMap
UserInfoMap = users
UserLock.Unlock()
t.Cleanup(func() {
UserLock.Lock()
UserInfoMap = original
UserLock.Unlock()
})
}
func TestCronTriggerSkipsServersOwnedByOtherUsers(t *testing.T) {
firstStream := newCapturedTaskStream()
secondStream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, firstStream),
withTaskStream(&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "admin-server"}, secondStream),
)
cronTask := &model.Cron{
Common: model.Common{ID: 99, UserID: 100},
Command: "id",
Cover: model.CronCoverAll,
Servers: []uint64{},
}
CronTrigger(cronTask)()
assertTaskCommand(t, firstStream, "id")
assertNoTask(t, secondStream)
}
func TestSendTriggerTasksSkipsCronOwnedByAnotherUser(t *testing.T) {
attackerStream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "attacker-server"}, attackerStream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin},
200: {Role: model.RoleMember},
})
adminCron := &model.Cron{
Common: model.Common{ID: 42, UserID: 1},
Command: "admin-maintenance",
Cover: model.CronCoverAlertTrigger,
}
cronClass := &CronClass{
class: class[uint64, *model.Cron]{
list: map[uint64]*model.Cron{adminCron.ID: adminCron},
},
}
cronClass.SendTriggerTasks([]uint64{adminCron.ID}, 7, 200)
assertNoTask(t, attackerStream)
}
func assertTaskCommand(t *testing.T, stream *capturedTaskStream, expectedCommand string) {
t.Helper()
select {
case task := <-stream.tasks:
if task.GetType() != model.TaskTypeCommand {
t.Fatalf("expected command task type, got %v", task.GetType())
}
if task.GetData() != expectedCommand {
t.Fatalf("expected command %q, got %q", expectedCommand, task.GetData())
}
case <-time.After(time.Second):
t.Fatalf("expected command %q to be sent", expectedCommand)
}
}
func assertNoTask(t *testing.T, stream *capturedTaskStream) {
t.Helper()
select {
case task := <-stream.tasks:
t.Fatalf("expected no task to be sent, got command %q", task.GetData())
case <-time.After(50 * time.Millisecond):
}
}
func TestCronTriggerSendsToMemberOwnedServer(t *testing.T) {
memberStream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, memberStream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
})
cronTask := &model.Cron{
Common: model.Common{ID: 99, UserID: 100},
Command: "id",
Cover: model.CronCoverAll,
}
CronTrigger(cronTask)()
assertTaskCommand(t, memberStream, "id")
}
func TestCronTriggerAdminCronFansOutAcrossOwners(t *testing.T) {
first := newCapturedTaskStream()
second := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, first),
withTaskStream(&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "admin-server"}, second),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin},
100: {Role: model.RoleMember},
200: {Role: model.RoleAdmin},
})
cronTask := &model.Cron{
Common: model.Common{ID: 99, UserID: 1},
Command: "maintenance",
Cover: model.CronCoverAll,
}
CronTrigger(cronTask)()
assertTaskCommand(t, first, "maintenance")
assertTaskCommand(t, second, "maintenance")
}
func TestCronTriggerLegacyZeroOwnerFansOut(t *testing.T) {
first := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, first),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
})
cronTask := &model.Cron{
Common: model.Common{ID: 99, UserID: 0},
Command: "legacy",
Cover: model.CronCoverAll,
}
CronTrigger(cronTask)()
assertTaskCommand(t, first, "legacy")
}
func TestCronTriggerSkipsServersWhenOwnerNotKnown(t *testing.T) {
stream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, stream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
})
cronTask := &model.Cron{
Common: model.Common{ID: 99, UserID: 999},
Command: "ghost",
Cover: model.CronCoverAll,
}
CronTrigger(cronTask)()
assertNoTask(t, stream)
}
func TestSendTriggerTasksAllowsSelfOwnedCron(t *testing.T) {
stream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "member-server"}, stream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
})
memberCron := &model.Cron{
Common: model.Common{ID: 42, UserID: 200},
Command: "member-task",
Cover: model.CronCoverAlertTrigger,
}
cronClass := &CronClass{
class: class[uint64, *model.Cron]{
list: map[uint64]*model.Cron{memberCron.ID: memberCron},
},
}
cronClass.SendTriggerTasks([]uint64{memberCron.ID}, 7, 200)
assertTaskCommand(t, stream, "member-task")
}
func TestSendTriggerTasksAllowsAdminCallerToTriggerAny(t *testing.T) {
stream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 9, UserID: 100}, Name: "any-server"}, stream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin},
100: {Role: model.RoleMember},
})
memberCron := &model.Cron{
Common: model.Common{ID: 42, UserID: 100},
Command: "member-task",
Cover: model.CronCoverAlertTrigger,
}
cronClass := &CronClass{
class: class[uint64, *model.Cron]{
list: map[uint64]*model.Cron{memberCron.ID: memberCron},
},
}
cronClass.SendTriggerTasks([]uint64{memberCron.ID}, 9, 1)
assertTaskCommand(t, stream, "member-task")
}
func TestSendTriggerTasksIgnoresUnknownTaskIDs(t *testing.T) {
stream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "member-server"}, stream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
200: {Role: model.RoleMember},
})
cronClass := &CronClass{
class: class[uint64, *model.Cron]{
list: map[uint64]*model.Cron{},
},
}
cronClass.SendTriggerTasks([]uint64{12345}, 7, 200)
cronClass.SendTriggerTasks(nil, 7, 200)
assertNoTask(t, stream)
}
func TestSendTriggerTasksMixedCronIDsOnlyFiresAllowed(t *testing.T) {
stream := newCapturedTaskStream()
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "member-server"}, stream),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin},
200: {Role: model.RoleMember},
})
memberCron := &model.Cron{
Common: model.Common{ID: 7, UserID: 200},
Command: "member-task",
Cover: model.CronCoverAlertTrigger,
}
adminCron := &model.Cron{
Common: model.Common{ID: 8, UserID: 1},
Command: "admin-task",
Cover: model.CronCoverAlertTrigger,
}
cronClass := &CronClass{
class: class[uint64, *model.Cron]{
list: map[uint64]*model.Cron{memberCron.ID: memberCron, adminCron.ID: adminCron},
},
}
cronClass.SendTriggerTasks([]uint64{memberCron.ID, adminCron.ID}, 7, 200)
select {
case task := <-stream.tasks:
if task.GetData() != "member-task" {
t.Fatalf("expected member-task, got %q", task.GetData())
}
case <-time.After(time.Second):
t.Fatalf("expected member-task to be sent")
}
assertNoTask(t, stream)
}
func TestAlertTriggerCronResultAuthorizationConsumesOneDispatch(t *testing.T) {
cronClass := &CronClass{}
cronClass.reserveAlertTriggerCronResult(42, 7)
cronClass.reserveAlertTriggerCronResult(42, 7)
if !cronClass.consumeAlertTriggerCronResult(42, 7) {
t.Fatal("expected first alert-trigger authorization to be consumed")
}
if !cronClass.consumeAlertTriggerCronResult(42, 7) {
t.Fatal("expected second alert-trigger authorization to be consumed")
}
if cronClass.consumeAlertTriggerCronResult(42, 7) {
t.Fatal("expected alert-trigger authorization to be consumed only once per dispatch")
}
}
func TestAlertTriggerCronResultAuthorizationExpires(t *testing.T) {
cronClass := &CronClass{
pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{
42: {7: {time.Now().Add(-time.Second)}},
},
}
if cronClass.consumeAlertTriggerCronResult(42, 7) {
t.Fatal("expired alert-trigger authorization must not be accepted")
}
if len(cronClass.pendingAlertTriggerTasks) != 0 {
t.Fatal("expired alert-trigger authorization must be pruned")
}
}
func TestAlertTriggerCronResultAuthorizationRevokeRemovesLatestDispatch(t *testing.T) {
existingAuthorizationExpiresAt := time.Now().Add(time.Hour)
cronClass := &CronClass{
pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{
42: {7: {existingAuthorizationExpiresAt}},
},
}
cronClass.reserveAlertTriggerCronResult(42, 7)
cronClass.revokeAlertTriggerCronResult(42, 7)
authorizations := cronClass.pendingAlertTriggerTasks[42][7]
if len(authorizations) != 1 {
t.Fatalf("expected one previous alert-trigger authorization to remain, got %d", len(authorizations))
}
if !authorizations[0].Equal(existingAuthorizationExpiresAt) {
t.Fatal("send failure rollback must remove the newest reserved authorization")
}
}
func TestCronClassUpdatePrunesAlertTriggerCronResultAuthorization(t *testing.T) {
cronClass := &CronClass{
Cron: cron.New(cron.WithSeconds()),
class: class[uint64, *model.Cron]{
list: map[uint64]*model.Cron{42: {Common: model.Common{ID: 42}}},
},
pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{
42: {7: {time.Now().Add(time.Hour)}},
},
}
cronClass.Update(&model.Cron{Common: model.Common{ID: 42}})
if len(cronClass.pendingAlertTriggerTasks) != 0 {
t.Fatal("cron update must prune old alert-trigger result authorizations")
}
}
func TestCronClassDeletePrunesAlertTriggerCronResultAuthorization(t *testing.T) {
cronClass := &CronClass{
Cron: cron.New(cron.WithSeconds()),
class: class[uint64, *model.Cron]{
list: map[uint64]*model.Cron{42: {Common: model.Common{ID: 42}}},
},
pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{
42: {7: {time.Now().Add(time.Hour)}},
},
}
cronClass.Delete([]uint64{42})
if len(cronClass.pendingAlertTriggerTasks) != 0 {
t.Fatal("cron delete must prune alert-trigger result authorizations")
}
}
// CanReportCronResult is the cron-side dual of canReportServiceResult: it gates
// agent-reported TaskTypeCommand results to only the cron/server pairs the
// dashboard actually fanned the task out to. Without these inbound checks any
// authenticated agent could fabricate a TaskResult for an arbitrary cron ID and
// poison LastResult / fire success/failure notifications belonging to another
// tenant. The tests below pin each Cover branch end-to-end against the dispatch
// logic in CronTrigger so the two sides stay symmetric.
func TestCanReportCronResultRejectsNilCronOrReporter(t *testing.T) {
cr := &model.Cron{Common: model.Common{ID: 7, UserID: 100}, Cover: model.CronCoverAll}
reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}}
if CanReportCronResult(nil, reporter) {
t.Fatal("nil cron must be rejected — would dereference inside cover branches")
}
if CanReportCronResult(cr, nil) {
t.Fatal("nil reporter must be rejected")
}
}
func TestCanReportCronResultRejectsForeignReporter(t *testing.T) {
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
200: {Role: model.RoleMember},
})
cr := &model.Cron{
Common: model.Common{ID: 7, UserID: 100},
Cover: model.CronCoverAll,
}
foreign := &model.Server{Common: model.Common{ID: 1, UserID: 200}}
if CanReportCronResult(cr, foreign) {
t.Fatal("foreign-user reporter must be rejected: CronTrigger never dispatched to it")
}
}
func TestCanReportCronResultCronCoverAllRejectsReporterInDenyList(t *testing.T) {
cr := &model.Cron{
Common: model.Common{ID: 7, UserID: 100},
Cover: model.CronCoverAll,
Servers: []uint64{1},
}
reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}}
if CanReportCronResult(cr, reporter) {
t.Fatal("CronCoverAll treats Servers as deny-list; reporter in the list must be rejected")
}
}
func TestCanReportCronResultCronCoverAllAcceptsReporterNotInDenyList(t *testing.T) {
cr := &model.Cron{
Common: model.Common{ID: 7, UserID: 100},
Cover: model.CronCoverAll,
Servers: []uint64{99},
}
reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}}
if !CanReportCronResult(cr, reporter) {
t.Fatal("CronCoverAll with reporter NOT in Servers must accept — CronTrigger dispatches to it")
}
}
func TestCanReportCronResultCronCoverIgnoreAllAcceptsReporterInAllowList(t *testing.T) {
cr := &model.Cron{
Common: model.Common{ID: 7, UserID: 100},
Cover: model.CronCoverIgnoreAll,
Servers: []uint64{1},
}
reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}}
if !CanReportCronResult(cr, reporter) {
t.Fatal("CronCoverIgnoreAll treats Servers as allow-list; reporter in the list must be accepted")
}
}
func TestCanReportCronResultCronCoverIgnoreAllRejectsReporterOutsideAllowList(t *testing.T) {
cr := &model.Cron{
Common: model.Common{ID: 7, UserID: 100},
Cover: model.CronCoverIgnoreAll,
Servers: []uint64{99},
}
reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}}
if CanReportCronResult(cr, reporter) {
t.Fatal("CronCoverIgnoreAll with reporter NOT in Servers must reject — CronTrigger never dispatched to it")
}
}
// failingTaskStream simulates a TaskStream whose Send always errors. CronTrigger
// uses this signal to revoke a reserved alert-trigger authorization, so the
// agent can't later attach to the cron via CanReportCronResult based on a
// dispatch that never actually reached the wire.
type failingTaskStream struct {
capturedTaskStream
sendErr error
}
func newFailingTaskStream(err error) *failingTaskStream {
return &failingTaskStream{
capturedTaskStream: capturedTaskStream{tasks: make(chan *pb.Task, 4)},
sendErr: err,
}
}
func (s *failingTaskStream) Send(task *pb.Task) error {
s.tasks <- task
return s.sendErr
}
func TestCronTriggerRevokesAlertTriggerAuthorizationOnSendFailure(t *testing.T) {
failing := newFailingTaskStream(context.Canceled)
replaceServerSharedForSecurityTest(t,
withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 100}, Name: "broken-server"}, failing),
)
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
100: {Role: model.RoleMember},
})
originalCronShared := CronShared
t.Cleanup(func() { CronShared = originalCronShared })
CronShared = &CronClass{
class: class[uint64, *model.Cron]{list: map[uint64]*model.Cron{}},
pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{},
}
cr := &model.Cron{
Common: model.Common{ID: 42, UserID: 100},
Cover: model.CronCoverAlertTrigger,
}
CronTrigger(cr, 7)()
// drain the dispatched task — Send error is what we care about, not the payload
select {
case <-failing.tasks:
case <-time.After(time.Second):
t.Fatal("expected CronTrigger to call Send before reacting to the error")
}
if CronShared.consumeAlertTriggerCronResult(42, 7) {
t.Fatal("Send failure must revoke the reserved alert-trigger authorization; otherwise a foreign agent could later report a result for a dispatch that never reached the wire")
}
if len(CronShared.pendingAlertTriggerTasks) != 0 {
t.Fatalf("expected pendingAlertTriggerTasks to be empty after revoke, got %d entries", len(CronShared.pendingAlertTriggerTasks))
}
}
func TestClassCheckPermission(t *testing.T) {
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin},
200: {Role: model.RoleMember},
})
sharedClass := &ServerClass{
class: class[uint64, *model.Server]{
list: map[uint64]*model.Server{
1: {Common: model.Common{ID: 1, UserID: 200}},
2: {Common: model.Common{ID: 2, UserID: 1}},
},
},
uuidToID: map[string]uint64{},
}
memberCtx, _ := gin.CreateTestContext(httptest.NewRecorder())
memberCtx.Set(model.CtxKeyAuthorizedUser, &model.User{
Common: model.Common{ID: 200},
Role: model.RoleMember,
})
adminCtx, _ := gin.CreateTestContext(httptest.NewRecorder())
adminCtx.Set(model.CtxKeyAuthorizedUser, &model.User{
Common: model.Common{ID: 1},
Role: model.RoleAdmin,
})
if !sharedClass.CheckPermission(memberCtx, slices.Values([]uint64{1})) {
t.Fatal("expected member to access own resource")
}
if sharedClass.CheckPermission(memberCtx, slices.Values([]uint64{2})) {
t.Fatal("expected member to be denied foreign resource")
}
if !sharedClass.CheckPermission(memberCtx, slices.Values([]uint64{})) {
t.Fatal("expected empty iterator to be allowed")
}
if !sharedClass.CheckPermission(memberCtx, slices.Values([]uint64{999})) {
t.Fatal("expected unknown id to be ignored (vacuous true)")
}
if !sharedClass.CheckPermission(adminCtx, slices.Values([]uint64{1, 2})) {
t.Fatal("expected admin to access any resource")
}
}
func TestServiceMonitorResultSkipsReporterOutsideServiceCover(t *testing.T) {
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "covered-server"},
&model.Server{Common: model.Common{ID: 2, UserID: 100}, Name: "uncovered-server"},
)
addServiceMonitorSecurityService(t, ss, &model.Service{
Common: model.Common{ID: 10, UserID: 100},
Name: "selected-only-service",
Type: model.TaskTypeTCPPing,
Target: "example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
})
ss.Dispatch(serviceMonitorResult(2, 10, model.TaskTypeTCPPing, true))
ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true))
waitForServiceHistory(t, 10, 1)
assertNoServiceHistory(t, 10, 2)
}
func TestServiceMonitorResultSkipsCoveredReporterOwnedByAnotherUser(t *testing.T) {
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "owner-server"},
&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "foreign-server"},
)
addServiceMonitorSecurityService(t, ss, &model.Service{
Common: model.Common{ID: 10, UserID: 100},
Name: "owner-only-service",
Type: model.TaskTypeTCPPing,
Target: "example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true, 2: true},
})
ss.Dispatch(serviceMonitorResult(2, 10, model.TaskTypeTCPPing, true))
ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true))
waitForServiceHistory(t, 10, 1)
assertNoServiceHistory(t, 10, 2)
}
func TestServiceMonitorResultSkipsMismatchedTaskType(t *testing.T) {
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "owner-server"},
)
addServiceMonitorSecurityService(t, ss, &model.Service{
Common: model.Common{ID: 10, UserID: 100},
Name: "http-service",
Type: model.TaskTypeHTTPGet,
Target: "https://example.invalid",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
})
ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, false))
ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeHTTPGet, true))
waitForTodayStats(t, ss, 10, 1, 0)
}
func TestServiceMonitorResultSkipsUnknownReporter(t *testing.T) {
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "owner-server"},
)
addServiceMonitorSecurityService(t, ss, &model.Service{
Common: model.Common{ID: 10, UserID: 100},
Name: "known-reporter-service",
Type: model.TaskTypeTCPPing,
Target: "example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
})
ss.Dispatch(serviceMonitorResult(999, 10, model.TaskTypeTCPPing, true))
ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true))
waitForServiceHistory(t, 10, 1)
assertNoServiceHistory(t, 10, 999)
}
func TestServiceMonitorResultAllowsCoveredReporterOwnedByServiceOwner(t *testing.T) {
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "owner-server"},
)
addServiceMonitorSecurityService(t, ss, &model.Service{
Common: model.Common{ID: 10, UserID: 100},
Name: "owner-service",
Type: model.TaskTypeTCPPing,
Target: "example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
})
ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true))
waitForServiceHistory(t, 10, 1)
}
func TestServiceMonitorResultAllowsCoveredReporterForAdminOwnedService(t *testing.T) {
replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
1: {Role: model.RoleAdmin},
200: {Role: model.RoleMember},
})
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "member-server"},
)
addServiceMonitorSecurityService(t, ss, &model.Service{
Common: model.Common{ID: 10, UserID: 1},
Name: "admin-service",
Type: model.TaskTypeTCPPing,
Target: "example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{2: true},
})
ss.Dispatch(serviceMonitorResult(2, 10, model.TaskTypeTCPPing, true))
waitForServiceHistory(t, 10, 2)
}
func newServiceMonitorSecurityHarness(t *testing.T, servers ...*model.Server) *ServiceSentinel {
t.Helper()
originalDB := DB
originalConf := Conf
originalCache := Cache
originalCronShared := CronShared
originalServerShared := ServerShared
originalServiceSentinelShared := ServiceSentinelShared
originalNotificationShared := NotificationShared
originalTSDBShared := TSDBShared
originalLoc := Loc
var sqlDBClose func() error
t.Cleanup(func() {
DB = originalDB
Conf = originalConf
Cache = originalCache
CronShared = originalCronShared
ServerShared = originalServerShared
ServiceSentinelShared = originalServiceSentinelShared
NotificationShared = originalNotificationShared
TSDBShared = originalTSDBShared
Loc = originalLoc
if sqlDBClose != nil {
_ = sqlDBClose()
}
})
db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{})
if err != nil {
t.Fatal(err)
}
sqlDB, err := db.DB()
if err != nil {
t.Fatal(err)
}
sqlDB.SetMaxOpenConns(1)
sqlDBClose = sqlDB.Close
DB = db
if err := DB.AutoMigrate(
model.Server{},
model.Service{},
model.ServiceHistory{},
model.Notification{},
model.NotificationGroup{},
model.NotificationGroupNotification{},
); err != nil {
t.Fatal(err)
}
Conf = &ConfigClass{Config: &model.Config{AvgPingCount: 1}}
Cache = cache.New(time.Minute, time.Minute)
CronShared = &CronClass{
Cron: cron.New(cron.WithSeconds()),
class: class[uint64, *model.Cron]{list: map[uint64]*model.Cron{}},
}
NotificationShared = &NotificationClass{
class: class[uint64, *model.Notification]{list: map[uint64]*model.Notification{}},
groupToIDList: map[uint64]map[uint64]*model.Notification{},
idToGroupList: map[uint64]map[uint64]struct{}{},
groupList: map[uint64]string{},
}
TSDBShared = nil
Loc = time.UTC
serverClass := &ServerClass{
class: class[uint64, *model.Server]{
list: make(map[uint64]*model.Server),
},
uuidToID: make(map[string]uint64),
}
for _, server := range servers {
serverClass.list[server.ID] = server
}
ServerShared = serverClass
bus := make(chan *model.Service, 1)
ss, err := NewServiceSentinel(bus)
if err != nil {
t.Fatal(err)
}
ServiceSentinelShared = ss
// LIFO Cleanup ordering: this Close() runs BEFORE the earlier t.Cleanup that
// restores Conf/Cache/CronShared/NotificationShared/TSDBShared, so the
// worker has fully exited before we swap those globals out. Skipping this
// step causes `go test -race` to flag the write-vs-read between the
// teardown and the still-running worker.
t.Cleanup(func() { ss.Close() })
return ss
}
func addServiceMonitorSecurityService(t *testing.T, ss *ServiceSentinel, service *model.Service) {
t.Helper()
if err := DB.Create(service).Error; err != nil {
t.Fatal(err)
}
if err := ss.Update(service); err != nil {
t.Fatal(err)
}
}
func serviceMonitorResult(reporter, serviceID uint64, taskType uint8, successful bool) ReportData {
return ReportData{
Reporter: reporter,
Data: &pb.TaskResult{
Id: serviceID,
Type: uint64(taskType),
Delay: 12,
Data: "service monitor result",
Successful: successful,
},
}
}
func waitForServiceHistory(t *testing.T, serviceID, serverID uint64) {
t.Helper()
deadline := time.After(time.Second)
for {
var count int64
if err := DB.Model(&model.ServiceHistory{}).
Where("service_id = ? AND server_id = ?", serviceID, serverID).
Count(&count).Error; err != nil {
t.Fatal(err)
}
if count > 0 {
return
}
select {
case <-deadline:
t.Fatalf("expected service history for service %d from server %d", serviceID, serverID)
default:
time.Sleep(10 * time.Millisecond)
}
}
}
func assertNoServiceHistory(t *testing.T, serviceID, serverID uint64) {
t.Helper()
var count int64
if err := DB.Model(&model.ServiceHistory{}).
Where("service_id = ? AND server_id = ?", serviceID, serverID).
Count(&count).Error; err != nil {
t.Fatal(err)
}
if count != 0 {
t.Fatalf("expected no service history for service %d from server %d, got %d", serviceID, serverID, count)
}
}
func waitForTodayStats(t *testing.T, ss *ServiceSentinel, serviceID uint64, wantUp, wantDown uint64) {
t.Helper()
deadline := time.After(time.Second)
for {
ss.serviceResponseDataStoreLock.RLock()
stats := ss.serviceStatusToday[serviceID]
var up, down uint64
if stats != nil {
up = stats.Up
down = stats.Down
}
ss.serviceResponseDataStoreLock.RUnlock()
if up == wantUp && down == wantDown {
return
}
if down > wantDown {
t.Fatalf("expected service %d down count %d, got %d", serviceID, wantDown, down)
}
select {
case <-deadline:
t.Fatalf("expected service %d stats up=%d down=%d", serviceID, wantUp, wantDown)
default:
time.Sleep(10 * time.Millisecond)
}
}
}
+78 -7
View File
@@ -6,6 +6,7 @@ import (
"log"
"slices"
"strings"
"sync"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/ddns"
@@ -15,6 +16,10 @@ import (
type ServerClass struct {
class[uint64, *model.Server]
// lifecycleMu serializes changes to the authoritative server entries with
// synchronous ServiceSentinel report processing.
lifecycleMu sync.RWMutex
uuidToID map[string]uint64
sortedListForGuest []*model.Server
@@ -30,18 +35,67 @@ func NewServerClass() *ServerClass {
var servers []model.Server
DB.Find(&servers)
for _, s := range servers {
innerS := s
model.InitServer(&innerS)
sc.list[innerS.ID] = &innerS
for i := range servers {
innerS := &servers[i]
model.InitServer(innerS)
sc.list[innerS.ID] = innerS
sc.uuidToID[innerS.UUID] = innerS.ID
}
sc.sortList()
model.OwnerServerIDsLookup = sc.ownerServerIDs
model.AllServerIDsLookup = sc.allServerIDs
model.OwnerIsAdminLookup = ownerIsAdmin
return sc
}
func (c *ServerClass) ownerServerIDs(ownerUID uint64) []uint64 {
var ids []uint64
c.Range(func(id uint64, s *model.Server) bool {
if s != nil && s.GetUserID() == ownerUID {
ids = append(ids, id)
}
return true
})
return ids
}
func (c *ServerClass) allServerIDs() []uint64 {
var ids []uint64
c.Range(func(id uint64, s *model.Server) bool {
if s != nil {
ids = append(ids, id)
}
return true
})
return ids
}
func ownerIsAdmin(ownerUID uint64) bool {
return userIsAdmin(ownerUID)
}
func (c *ServerClass) lockLifecycleRead() {
c.lifecycleMu.RLock()
}
func (c *ServerClass) unlockLifecycleRead() {
c.lifecycleMu.RUnlock()
}
func (c *ServerClass) lockLifecycleWrite() {
c.lifecycleMu.Lock()
}
func (c *ServerClass) unlockLifecycleWrite() {
c.lifecycleMu.Unlock()
}
func (c *ServerClass) Update(s *model.Server, uuid string) {
c.lockLifecycleWrite()
defer c.unlockLifecycleWrite()
c.listMu.Lock()
c.list[s.ID] = s
@@ -61,11 +115,17 @@ func (c *ServerClass) Update(s *model.Server, uuid string) {
}
func (c *ServerClass) Delete(idList []uint64) {
c.lockLifecycleWrite()
defer c.unlockLifecycleWrite()
c.listMu.Lock()
for _, id := range idList {
serverUUID := c.list[id].UUID
delete(c.uuidToID, serverUUID)
s, ok := c.list[id]
if !ok {
continue
}
delete(c.uuidToID, s.UUID)
delete(c.list, id)
}
@@ -74,6 +134,17 @@ func (c *ServerClass) Delete(idList []uint64) {
c.sortList()
}
// setUserID updates in-memory ownership under the server lifecycle lock so a
// transfer cannot change authorization during synchronous report processing.
func (c *ServerClass) setUserID(id, userID uint64) {
c.lockLifecycleWrite()
defer c.unlockLifecycleWrite()
if s, ok := c.Get(id); ok && s != nil {
s.SetUserID(userID)
}
}
func (c *ServerClass) GetSortedListForGuest() []*model.Server {
c.sortedListMu.RLock()
defer c.sortedListMu.RUnlock()
@@ -93,7 +164,7 @@ func (c *ServerClass) UpdateDDNS(server *model.Server, ip *model.IP) error {
confServers := strings.Split(Conf.DNSServers, ",")
ctx := context.WithValue(context.Background(), ddns.DNSServerKey{}, utils.IfOr(confServers[0] != "", confServers, utils.DNSServers))
providers, err := DDNSShared.GetDDNSProvidersFromProfiles(server.DDNSProfiles, utils.IfOr(ip != nil, ip, &server.GeoIP.IP))
providers, err := DDNSShared.GetDDNSProvidersFromProfiles(server.DDNSProfiles, utils.IfOr(ip != nil, ip, &server.GeoIP.IP), server.GetUserID())
if err != nil {
return err
}
@@ -0,0 +1,32 @@
package singleton
import (
"testing"
"github.com/nezhahq/nezha/model"
)
func TestServerClassDeleteMissingIDNoPanic(t *testing.T) {
c := &ServerClass{
class: class[uint64, *model.Server]{
list: map[uint64]*model.Server{
1: {Common: model.Common{ID: 1}, UUID: "uuid-1"},
},
},
uuidToID: map[string]uint64{"uuid-1": 1},
}
c.Delete([]uint64{999999})
if _, ok := c.list[1]; !ok {
t.Fatalf("existing server 1 must remain after deleting a non-existent id")
}
c.Delete([]uint64{1, 424242})
if _, ok := c.list[1]; ok {
t.Fatalf("server 1 should be removed")
}
if _, ok := c.uuidToID["uuid-1"]; ok {
t.Fatalf("uuid mapping for server 1 should be removed")
}
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,37 @@
package singleton
import (
"testing"
"github.com/stretchr/testify/require"
"github.com/nezhahq/nezha/model"
)
func TestServiceSentinelUpdateRejectsNonProbeTaskTypes(t *testing.T) {
ss := &ServiceSentinel{}
require.Error(t, ss.Update(nil))
for _, taskType := range []uint8{0, model.TaskTypeCommand, model.TaskTypeApplyConfig, model.TaskTypeExec, 255} {
require.Error(t, ss.Update(&model.Service{Type: taskType}), "type %d must not be scheduled", taskType)
}
}
func TestServiceSentinelQuarantinesInvalidPersistedTypes(t *testing.T) {
ss := newServiceMonitorSecurityHarness(t)
insert := `INSERT INTO services
(id, user_id, name, type, target, duration, cover, skip_servers_raw, fail_trigger_tasks_raw, recover_trigger_tasks_raw)
VALUES (?, 100, ?, ?, 'example.invalid:443', 3600, ?, '{}', '[]', '[]')`
require.NoError(t, DB.Exec(insert, 91, "legacy-command", model.TaskTypeCommand, model.ServiceCoverIgnoreAll).Error)
require.NoError(t, DB.Exec(insert, 92, "legacy-apply-config", model.TaskTypeApplyConfig, model.ServiceCoverIgnoreAll).Error)
require.NoError(t, DB.Exec(insert, 93, "valid-probe", model.TaskTypeTCPPing, model.ServiceCoverIgnoreAll).Error)
require.NoError(t, ss.loadServiceHistory())
_, commandLoaded := ss.Get(91)
_, applyConfigLoaded := ss.Get(92)
valid, validLoaded := ss.Get(93)
require.False(t, commandLoaded)
require.False(t, applyConfigLoaded)
require.True(t, validLoaded)
require.Equal(t, uint8(model.TaskTypeTCPPing), valid.Type)
}
+338 -209
View File
@@ -74,8 +74,11 @@ type ServiceSentinel struct {
serviceCurrentStatusData map[uint64]*serviceTaskStatus // 当前任务结果缓存
serviceResponseDataStore map[uint64]serviceResponseData // 当前数据
serviceResponsePing map[uint64]map[uint64]*pingStore // [service_id] -> ClientID -> delay
tlsCertCache map[uint64]string
serviceResponsePing map[uint64]map[uint64]*pingStore // guarded by serviceResponseDataStoreLock; [service_id] -> ClientID -> delay
tlsCertCache map[uint64]string // guarded by serviceResponseDataStoreLock
serviceReportValidatedHook func(uint64)
loadStatsResponseLockedHook func()
serviceReportBeforeTLSSideEffectsHook func(uint64)
servicesLock sync.RWMutex
serviceListLock sync.RWMutex
@@ -85,6 +88,16 @@ type ServiceSentinel struct {
// 30天数据缓存
monthlyStatusLock sync.Mutex
monthlyStatus map[uint64]*serviceResponseItem
// closeOnce + workerWG together let Close() wait for the worker goroutine
// to fully exit. Without this, a test that swaps ServiceSentinelShared back
// to its original value in t.Cleanup races against the still-running
// worker, which keeps reading globals like Conf/CronShared/NotificationShared.
// Production never calls Close() — the process exits while the worker is
// still running and that is fine — but tests must drain the worker before
// restoring globals.
closeOnce sync.Once
workerWG sync.WaitGroup
}
// NewServiceSentinel 创建服务监控器
@@ -113,7 +126,11 @@ func NewServiceSentinel(serviceSentinelDispatchBus chan<- *model.Service) (*Serv
ss.loadTodayStats(today)
// 启动服务监控器
go ss.worker()
ss.workerWG.Add(1)
go func() {
defer ss.workerWG.Done()
ss.worker()
}()
// 每日将游标往后推一天
_, err = CronShared.AddFunc("0 0 0 * * *", ss.refreshMonthlyServiceStatus)
@@ -192,7 +209,15 @@ func (ss *ServiceSentinel) loadServiceHistory() error {
return err
}
validServices := services[:0]
for _, service := range services {
if err := model.ValidateServiceMonitorType(uint64(service.Type)); err != nil {
// Existing databases may contain values written before Service.Type was
// constrained. Quarantine them in the database for operator review, but
// never register a cron job that could dispatch a privileged Agent task.
log.Printf("NEZHA>> quarantining service %d: %v", service.ID, err)
continue
}
task := service
// 通过cron定时将服务监控任务传递给任务调度管道
service.CronJobID, err = CronShared.AddFunc(task.CronSpec(), func() {
@@ -205,7 +230,9 @@ func (ss *ServiceSentinel) loadServiceHistory() error {
ss.serviceCurrentStatusData[service.ID] = new(serviceTaskStatus)
ss.serviceCurrentStatusData[service.ID].result = make([]*pb.TaskResult, 0, _CurrentStatusSize)
ss.serviceStatusToday[service.ID] = &_TodayStatsOfService{}
validServices = append(validServices, service)
}
services = validServices
ss.serviceList = services
sortServices(ss.serviceList)
@@ -322,6 +349,13 @@ func (ss *ServiceSentinel) loadTodayStats(today time.Time) {
}
func (ss *ServiceSentinel) Update(m *model.Service) error {
if m == nil {
return fmt.Errorf("service is nil")
}
if err := model.ValidateServiceMonitorType(uint64(m.Type)); err != nil {
return err
}
ss.serviceResponseDataStoreLock.Lock()
defer ss.serviceResponseDataStoreLock.Unlock()
ss.monthlyStatusLock.Lock()
@@ -372,11 +406,21 @@ func (ss *ServiceSentinel) Delete(ids []uint64) {
for _, id := range ids {
delete(ss.serviceCurrentStatusData, id)
delete(ss.serviceResponseDataStore, id)
delete(ss.serviceResponsePing, id)
delete(ss.tlsCertCache, id)
delete(ss.serviceStatusToday, id)
// 停掉定时任务
CronShared.Remove(ss.services[id].CronJobID)
// GHSA-jx78-55p5-rwv5 (Finding 2): guard against a caller supplying an id
// that does not exist in the in-memory registry. CheckPermission returns
// vacuously true for unknown ids, so the controller layer cannot prevent
// this. Without the guard, ss.services[id] is nil and the .CronJobID
// field access panics, aborting the Delete loop before the remaining valid
// ids are cleaned from memory — their service records were already deleted
// from the database, producing zombie services.
if svc := ss.services[id]; svc != nil {
CronShared.Remove(svc.CronJobID)
}
delete(ss.services, id)
delete(ss.monthlyStatus, id)
@@ -384,12 +428,15 @@ func (ss *ServiceSentinel) Delete(ids []uint64) {
}
func (ss *ServiceSentinel) LoadStats() map[uint64]*serviceResponseItem {
ss.servicesLock.RLock()
defer ss.servicesLock.RUnlock()
ss.serviceResponseDataStoreLock.RLock()
defer ss.serviceResponseDataStoreLock.RUnlock()
if ss.loadStatsResponseLockedHook != nil {
ss.loadStatsResponseLockedHook()
}
ss.monthlyStatusLock.Lock()
defer ss.monthlyStatusLock.Unlock()
ss.servicesLock.RLock()
defer ss.servicesLock.RUnlock()
// 刷新最新一天的数据
for k := range ss.services {
@@ -424,11 +471,6 @@ func (ss *ServiceSentinel) CopyStats() map[uint64]model.ServiceResponseItem {
sri := make(map[uint64]model.ServiceResponseItem)
for k, service := range stats {
if !service.service.EnableShowInService {
delete(stats, k)
continue
}
service.ServiceName = service.service.Name
sri[k] = service.ServiceResponseItem
}
@@ -472,223 +514,297 @@ func (ss *ServiceSentinel) CheckPermission(c *gin.Context, idList iter.Seq[uint6
return true
}
func canReportServiceResult(service *model.Service, reporter *model.Server, taskType uint64) bool {
if service == nil || reporter == nil || uint64(service.Type) != taskType {
return false
}
switch service.Cover {
case model.ServiceCoverAll:
if service.SkipServers[reporter.ID] {
return false
}
case model.ServiceCoverIgnoreAll:
if !service.SkipServers[reporter.ID] {
return false
}
default:
return false
}
return service.UserID == reporter.GetUserID() || userIsAdmin(service.UserID)
}
// Close shuts down the ServiceSentinel worker goroutine and waits for it to
// exit. It is idempotent and safe to call more than once.
//
// Why this exists: the worker reads multiple package-level globals during
// each report (Conf, CronShared via notifyCheck, NotificationShared via
// UnMuteNotification, ServerShared, TSDBShared). A test fixture that swaps
// those globals out in t.Cleanup MUST first call Close() — otherwise the
// cleanup write races the still-running worker's read and `go test -race`
// fires (see security_regression_test.go newServiceMonitorSecurityHarness).
// Production never calls Close because the process exits with the worker
// still running, which is fine.
func (ss *ServiceSentinel) Close() {
ss.closeOnce.Do(func() {
close(ss.serviceReportChannel)
ss.workerWG.Wait()
})
}
// worker 服务监控的实际工作流程
//
// IMPORTANT: this loop reads several package-level globals (Conf, CronShared,
// NotificationShared, ServerShared, TSDBShared). Any test that replaces those
// globals via t.Cleanup must first call ServiceSentinel.Close() so the worker
// drains and exits before the swap, otherwise the race detector trips. See
// the Close() comment above for the full rationale.
func (ss *ServiceSentinel) worker() {
// 从服务状态汇报管道获取汇报的服务数据
for r := range ss.serviceReportChannel {
css, _ := ss.Get(r.Data.GetId())
if css == nil || css.ID == 0 {
log.Printf("NEZHA>> Incorrect service monitor report %+v", r)
continue
}
css = nil
mh := r.Data
if mh.Type == model.TaskTypeTCPPing || mh.Type == model.TaskTypeICMPPing {
// TCP/ICMP Ping 使用平均值计算后再写入
serviceTcpMap, ok := ss.serviceResponsePing[mh.GetId()]
if !ok {
serviceTcpMap = make(map[uint64]*pingStore)
ss.serviceResponsePing[mh.GetId()] = serviceTcpMap
}
ts, ok := serviceTcpMap[r.Reporter]
if !ok {
ts = &pingStore{}
}
ts.count++
ts.ping = (ts.ping*float64(ts.count-1) + float64(mh.Delay)) / float64(ts.count)
if mh.Successful {
ts.successCount++
}
if ts.count == Conf.AvgPingCount {
if TSDBEnabled() {
if err := TSDBShared.WriteServiceMetrics(&tsdb.ServiceMetrics{
ServiceID: mh.GetId(),
ServerID: r.Reporter,
Timestamp: time.Now(),
Delay: ts.ping,
Successful: ts.successCount*2 >= ts.count,
}); err != nil {
log.Printf("NEZHA>> Failed to save service monitor metrics to TSDB: %v", err)
}
} else {
if err := DB.Create(&model.ServiceHistory{
ServiceID: mh.GetId(),
AvgDelay: ts.ping,
Data: mh.Data,
ServerID: r.Reporter,
}).Error; err != nil {
log.Printf("NEZHA>> Failed to save service monitor metrics: %v", err)
}
serverShared := ServerShared
func() {
defer func() {
if recovered := recover(); recovered != nil {
log.Printf("NEZHA>> Service monitor report processing panicked: %v", recovered)
}
ts.count = 0
ts.ping = 0
ts.successCount = 0
}
serviceTcpMap[r.Reporter] = ts
} else {
}()
ss.processReport(r, serverShared)
}()
}
}
func (ss *ServiceSentinel) processReport(r ReportData, serverShared *ServerClass) {
serverShared.lockLifecycleRead()
defer serverShared.unlockLifecycleRead()
cs, _ := ss.Get(r.Data.GetId())
reporter, _ := serverShared.Get(r.Reporter)
// 入站结果必须匹配出站任务派发边界,避免 agent 伪造其他服务 ID 写入监控状态。
if !canReportServiceResult(cs, reporter, r.Data.GetType()) {
log.Printf("NEZHA>> Incorrect service monitor report %+v", r)
return
}
if ss.serviceReportValidatedHook != nil {
ss.serviceReportValidatedHook(r.Data.GetId())
}
mh := r.Data
m := serverShared.GetList()
// Serialize Delete and Update before this accepted report causes any side effect.
ss.serviceResponseDataStoreLock.Lock()
defer ss.serviceResponseDataStoreLock.Unlock()
serviceStatusToday := ss.serviceStatusToday[mh.GetId()]
serviceCurrentStatusData := ss.serviceCurrentStatusData[mh.GetId()]
currentService, serviceExists := ss.Get(mh.GetId())
if serviceStatusToday == nil || serviceCurrentStatusData == nil || !serviceExists ||
!canReportServiceResult(currentService, reporter, mh.GetType()) {
return
}
cs = currentService
if mh.Type == model.TaskTypeTCPPing || mh.Type == model.TaskTypeICMPPing {
// TCP/ICMP Ping 使用平均值计算后再写入
serviceTcpMap, ok := ss.serviceResponsePing[mh.GetId()]
if !ok {
serviceTcpMap = make(map[uint64]*pingStore)
ss.serviceResponsePing[mh.GetId()] = serviceTcpMap
}
ts, ok := serviceTcpMap[r.Reporter]
if !ok {
ts = &pingStore{}
}
ts.count++
ts.ping = (ts.ping*float64(ts.count-1) + float64(mh.Delay)) / float64(ts.count)
if mh.Successful {
ts.successCount++
}
if ts.count == Conf.AvgPingCount {
if TSDBEnabled() {
if err := TSDBShared.WriteServiceMetrics(&tsdb.ServiceMetrics{
ServiceID: mh.GetId(),
ServerID: r.Reporter,
Timestamp: time.Now(),
Delay: float64(mh.Delay),
Successful: mh.Successful,
Delay: ts.ping,
Successful: ts.successCount*2 >= ts.count,
}); err != nil {
log.Printf("NEZHA>> Failed to save service monitor metrics to TSDB: %v", err)
}
}
}
ss.serviceResponseDataStoreLock.Lock()
// 写入当天状态
if mh.Successful {
ss.serviceStatusToday[mh.GetId()].Delay = (ss.serviceStatusToday[mh.
GetId()].Delay*float64(ss.serviceStatusToday[mh.GetId()].Up) +
float64(mh.Delay)) / float64(ss.serviceStatusToday[mh.GetId()].Up+1)
ss.serviceStatusToday[mh.GetId()].Up++
} else {
ss.serviceStatusToday[mh.GetId()].Down++
}
currentTime := time.Now()
if ss.serviceCurrentStatusData[mh.GetId()].t.IsZero() {
ss.serviceCurrentStatusData[mh.GetId()].t = currentTime
}
// 写入当前数据
if ss.serviceCurrentStatusData[mh.GetId()].t.Before(currentTime) {
ss.serviceCurrentStatusData[mh.GetId()].t = currentTime.Add(30 * time.Second)
ss.serviceCurrentStatusData[mh.GetId()].result = append(ss.serviceCurrentStatusData[mh.GetId()].result, mh)
}
// 更新当前状态
ss.serviceResponseDataStore[mh.GetId()] = serviceResponseData{}
// 永远是最新的 30 个数据的状态 [01:00, 02:00, 03:00] -> [04:00, 02:00, 03: 00]
for _, cs := range ss.serviceCurrentStatusData[mh.GetId()].result {
if cs.GetId() > 0 {
rd := ss.serviceResponseDataStore[mh.GetId()]
if cs.Successful {
rd.Up++
rd.Delay = (rd.Delay*float64(rd.Up-1) + float64(cs.Delay)) / float64(rd.Up)
} else {
rd.Down++
}
ss.serviceResponseDataStore[mh.GetId()] = rd
}
}
// 计算在线率,
var stateCode uint8
{
upPercent := uint64(0)
rd := ss.serviceResponseDataStore[mh.GetId()]
if rd.Down+rd.Up > 0 {
upPercent = rd.Up * 100 / (rd.Down + rd.Up)
}
stateCode = GetStatusCode(upPercent)
}
if len(ss.serviceCurrentStatusData[mh.GetId()].result) == _CurrentStatusSize {
ss.serviceCurrentStatusData[mh.GetId()].t = currentTime
if !TSDBEnabled() {
rd := ss.serviceResponseDataStore[mh.GetId()]
} else {
if err := DB.Create(&model.ServiceHistory{
ServiceID: mh.GetId(),
AvgDelay: rd.Delay,
AvgDelay: ts.ping,
Data: mh.Data,
Up: rd.Up,
Down: rd.Down,
ServerID: r.Reporter,
}).Error; err != nil {
log.Printf("NEZHA>> Failed to save service monitor metrics: %v", err)
}
}
ss.serviceCurrentStatusData[mh.GetId()].result = ss.serviceCurrentStatusData[mh.GetId()].result[:0]
ts.count = 0
ts.ping = 0
ts.successCount = 0
}
cs, _ := ss.Get(mh.GetId())
m := ServerShared.GetList()
// 延迟报警
if mh.Delay > 0 {
delayCheck(&r, m, cs, mh)
}
// 状态变更报警+触发任务执行
if stateCode == StatusDown || stateCode != ss.serviceCurrentStatusData[mh.GetId()].lastStatus {
lastStatus := ss.serviceCurrentStatusData[mh.GetId()].lastStatus
// 存储新的状态值
ss.serviceCurrentStatusData[mh.GetId()].lastStatus = stateCode
notifyCheck(&r, m, cs, mh, lastStatus, stateCode)
}
ss.serviceResponseDataStoreLock.Unlock()
// TLS 证书报警
var errMsg string
if strings.HasPrefix(mh.Data, "SSL证书错误:") {
// i/o timeout、connection timeout、EOF 错误
if !strings.HasSuffix(mh.Data, "timeout") &&
!strings.HasSuffix(mh.Data, "EOF") &&
!strings.HasSuffix(mh.Data, "timed out") {
errMsg = mh.Data
if cs.Notify {
muteLabel := NotificationMuteLabel.ServiceTLS(mh.GetId(), "network")
go NotificationShared.SendNotification(cs.NotificationGroupID, Localizer.Tf("[TLS] Fetch cert info failed, Reporter: %s, Error: %s", cs.Name, errMsg), muteLabel)
}
serviceTcpMap[r.Reporter] = ts
} else {
if TSDBEnabled() {
if err := TSDBShared.WriteServiceMetrics(&tsdb.ServiceMetrics{
ServiceID: mh.GetId(),
ServerID: r.Reporter,
Timestamp: time.Now(),
Delay: float64(mh.Delay),
Successful: mh.Successful,
}); err != nil {
log.Printf("NEZHA>> Failed to save service monitor metrics to TSDB: %v", err)
}
} else {
// 清除网络错误静音缓存
NotificationShared.UnMuteNotification(cs.NotificationGroupID, NotificationMuteLabel.ServiceTLS(mh.GetId(), "network"))
}
}
var newCert = strings.Split(mh.Data, "|")
if len(newCert) > 1 {
enableNotify := cs.Notify
// 写入当天状态
if mh.Successful {
serviceStatusToday.Delay = (serviceStatusToday.Delay*float64(serviceStatusToday.Up) +
float64(mh.Delay)) / float64(serviceStatusToday.Up+1)
serviceStatusToday.Up++
} else {
serviceStatusToday.Down++
}
// 首次获取证书信息时,缓存证书信息
if ss.tlsCertCache[mh.GetId()] == "" {
ss.tlsCertCache[mh.GetId()] = mh.Data
currentTime := time.Now()
if serviceCurrentStatusData.t.IsZero() {
serviceCurrentStatusData.t = currentTime
}
// 写入当前数据
if serviceCurrentStatusData.t.Before(currentTime) {
serviceCurrentStatusData.t = currentTime.Add(30 * time.Second)
serviceCurrentStatusData.result = append(serviceCurrentStatusData.result, mh)
}
// 更新当前状态
ss.serviceResponseDataStore[mh.GetId()] = serviceResponseData{}
// 永远是最新的 30 个数据的状态 [01:00, 02:00, 03:00] -> [04:00, 02:00, 03: 00]
for _, cs := range serviceCurrentStatusData.result {
if cs.GetId() > 0 {
rd := ss.serviceResponseDataStore[mh.GetId()]
if cs.Successful {
rd.Up++
rd.Delay = (rd.Delay*float64(rd.Up-1) + float64(cs.Delay)) / float64(rd.Up)
} else {
rd.Down++
}
ss.serviceResponseDataStore[mh.GetId()] = rd
}
}
// 计算在线率,
var stateCode uint8
{
upPercent := uint64(0)
rd := ss.serviceResponseDataStore[mh.GetId()]
if rd.Down+rd.Up > 0 {
upPercent = rd.Up * 100 / (rd.Down + rd.Up)
}
stateCode = GetStatusCode(upPercent)
}
if len(serviceCurrentStatusData.result) == _CurrentStatusSize {
serviceCurrentStatusData.t = currentTime
if !TSDBEnabled() {
rd := ss.serviceResponseDataStore[mh.GetId()]
if err := DB.Create(&model.ServiceHistory{
ServiceID: mh.GetId(),
AvgDelay: rd.Delay,
Data: mh.Data,
Up: rd.Up,
Down: rd.Down,
}).Error; err != nil {
log.Printf("NEZHA>> Failed to save service monitor metrics: %v", err)
}
}
serviceCurrentStatusData.result = serviceCurrentStatusData.result[:0]
}
// 延迟报警
if mh.Delay > 0 {
delayCheck(&r, m, cs, mh)
}
// 状态变更报警+触发任务执行
if stateCode == StatusDown || stateCode != serviceCurrentStatusData.lastStatus {
lastStatus := serviceCurrentStatusData.lastStatus
// 存储新的状态值
serviceCurrentStatusData.lastStatus = stateCode
notifyCheck(&r, m, cs, mh, lastStatus, stateCode)
}
// TLS 证书报警
if ss.serviceReportBeforeTLSSideEffectsHook != nil {
ss.serviceReportBeforeTLSSideEffectsHook(mh.GetId())
}
var errMsg string
if strings.HasPrefix(mh.Data, "SSL证书错误:") {
// i/o timeout、connection timeout、EOF 错误
if !strings.HasSuffix(mh.Data, "timeout") &&
!strings.HasSuffix(mh.Data, "EOF") &&
!strings.HasSuffix(mh.Data, "timed out") {
errMsg = mh.Data
if cs.Notify {
muteLabel := NotificationMuteLabel.ServiceTLS(mh.GetId(), "network")
go NotificationShared.SendNotification(cs.NotificationGroupID, Localizer.Tf("[TLS] Fetch cert info failed, Reporter: %s, Error: %s", cs.Name, errMsg), muteLabel)
}
}
} else {
// 清除网络错误静音缓存
NotificationShared.UnMuteNotification(cs.NotificationGroupID, NotificationMuteLabel.ServiceTLS(mh.GetId(), "network"))
var newCert = strings.Split(mh.Data, "|")
if len(newCert) > 1 {
enableNotify := cs.Notify
// 首次获取证书信息时,缓存证书信息
if ss.tlsCertCache[mh.GetId()] == "" {
ss.tlsCertCache[mh.GetId()] = mh.Data
}
oldCert := strings.Split(ss.tlsCertCache[mh.GetId()], "|")
isCertChanged := false
expiresOld, _ := time.Parse("2006-01-02 15:04:05 -0700 MST", oldCert[1])
expiresNew, _ := time.Parse("2006-01-02 15:04:05 -0700 MST", newCert[1])
// 证书变更时,更新缓存
if oldCert[0] != newCert[0] && !expiresNew.Equal(expiresOld) {
isCertChanged = true
ss.tlsCertCache[mh.GetId()] = mh.Data
}
notificationGroupID := cs.NotificationGroupID
serviceName := cs.Name
// 需要发送提醒
if enableNotify {
// 证书过期提醒
if expiresNew.Before(time.Now().AddDate(0, 0, 7)) {
expiresTimeStr := expiresNew.Format("2006-01-02 15:04:05")
errMsg = Localizer.Tf(
"The TLS certificate will expire within seven days. Expiration time: %s",
expiresTimeStr,
)
// 静音规则: 服务id+证书过期时间
// 用于避免多个监测点对相同证书同时报警
muteLabel := NotificationMuteLabel.ServiceTLS(mh.GetId(), fmt.Sprintf("expire_%s", expiresTimeStr))
go NotificationShared.SendNotification(notificationGroupID, fmt.Sprintf("[TLS] %s %s", serviceName, errMsg), muteLabel)
}
oldCert := strings.Split(ss.tlsCertCache[mh.GetId()], "|")
isCertChanged := false
expiresOld, _ := time.Parse("2006-01-02 15:04:05 -0700 MST", oldCert[1])
expiresNew, _ := time.Parse("2006-01-02 15:04:05 -0700 MST", newCert[1])
// 证书变更提醒
if isCertChanged {
errMsg = Localizer.Tf(
"TLS certificate changed, old: issuer %s, expires at %s; new: issuer %s, expires at %s",
oldCert[0], expiresOld.Format("2006-01-02 15:04:05"), newCert[0], expiresNew.Format("2006-01-02 15:04:05"))
// 证书变更时,更新缓存
if oldCert[0] != newCert[0] && !expiresNew.Equal(expiresOld) {
isCertChanged = true
ss.tlsCertCache[mh.GetId()] = mh.Data
}
notificationGroupID := cs.NotificationGroupID
serviceName := cs.Name
// 需要发送提醒
if enableNotify {
// 证书过期提醒
if expiresNew.Before(time.Now().AddDate(0, 0, 7)) {
expiresTimeStr := expiresNew.Format("2006-01-02 15:04:05")
errMsg = Localizer.Tf(
"The TLS certificate will expire within seven days. Expiration time: %s",
expiresTimeStr,
)
// 静音规则: 服务id+证书过期时间
// 用于避免多个监测点对相同证书同时报警
muteLabel := NotificationMuteLabel.ServiceTLS(mh.GetId(), fmt.Sprintf("expire_%s", expiresTimeStr))
go NotificationShared.SendNotification(notificationGroupID, fmt.Sprintf("[TLS] %s %s", serviceName, errMsg), muteLabel)
}
// 证书变更提醒
if isCertChanged {
errMsg = Localizer.Tf(
"TLS certificate changed, old: issuer %s, expires at %s; new: issuer %s, expires at %s",
oldCert[0], expiresOld.Format("2006-01-02 15:04:05"), newCert[0], expiresNew.Format("2006-01-02 15:04:05"))
// 证书变更后会自动更新缓存,所以不需要静音
go NotificationShared.SendNotification(notificationGroupID, fmt.Sprintf("[TLS] %s %s", serviceName, errMsg), "")
}
// 证书变更后会自动更新缓存,所以不需要静音
go NotificationShared.SendNotification(notificationGroupID, fmt.Sprintf("[TLS] %s %s", serviceName, errMsg), "")
}
}
}
@@ -700,17 +816,25 @@ func delayCheck(r *ReportData, m map[uint64]*model.Server, ss *model.Service, mh
return
}
// GHSA-jx78-55p5-rwv5 (incomplete fix of GHSA-qjpp-gffx-2wm9): the server
// map snapshot m is taken outside serviceResponseDataStoreLock and
// ServerShared has its own independent lock, so a concurrent batch-delete of
// the reporter's server can remove the entry between the pre-lock validation
// and this point. Guard against the nil pointer before using the server.
reporterServer := m[r.Reporter]
if reporterServer == nil {
return
}
notificationGroupID := ss.NotificationGroupID
minMuteLabel := NotificationMuteLabel.ServiceLatencyMin(mh.GetId())
maxMuteLabel := NotificationMuteLabel.ServiceLatencyMax(mh.GetId())
if mh.Delay > ss.MaxLatency {
// 延迟超过最大值
reporterServer := m[r.Reporter]
msg := Localizer.Tf("[Latency] %s %2f > %2f, Reporter: %s", ss.Name, mh.Delay, ss.MaxLatency, reporterServer.Name)
go NotificationShared.SendNotification(notificationGroupID, msg, minMuteLabel)
} else if mh.Delay < ss.MinLatency {
// 延迟低于最小值
reporterServer := m[r.Reporter]
msg := Localizer.Tf("[Latency] %s %2f < %2f, Reporter: %s", ss.Name, mh.Delay, ss.MinLatency, reporterServer.Name)
go NotificationShared.SendNotification(notificationGroupID, msg, maxMuteLabel)
} else {
@@ -722,10 +846,16 @@ func delayCheck(r *ReportData, m map[uint64]*model.Server, ss *model.Service, mh
func notifyCheck(r *ReportData, m map[uint64]*model.Server,
ss *model.Service, mh *pb.TaskResult, lastStatus, stateCode uint8) {
// GHSA-jx78-55p5-rwv5: guard against concurrent server deletion (same TOCTOU
// class as the 2026-07-21 fix, a few dozen lines lower in the same worker).
// ServerShared has its own lock; m is a snapshot taken outside
// serviceResponseDataStoreLock, so the server may have been removed between
// the pre-lock validation and here.
reporterServer := m[r.Reporter]
// 判断是否需要发送通知
isNeedSendNotification := ss.Notify && (lastStatus != 0 || stateCode == StatusDown)
if isNeedSendNotification {
reporterServer := m[r.Reporter]
if isNeedSendNotification && reporterServer != nil {
notificationGroupID := ss.NotificationGroupID
notificationMsg := Localizer.Tf("[%s] %s Reporter: %s, Error: %s", StatusCodeToString(stateCode), ss.Name, reporterServer.Name, mh.Data)
muteLabel := NotificationMuteLabel.ServiceStateChanged(mh.GetId())
@@ -740,14 +870,13 @@ func notifyCheck(r *ReportData, m map[uint64]*model.Server,
// 判断是否需要触发任务
isNeedTriggerTask := ss.EnableTriggerTask && lastStatus != 0
if isNeedTriggerTask {
reporterServer := m[r.Reporter]
if isNeedTriggerTask && reporterServer != nil {
if stateCode == StatusGood && lastStatus != stateCode {
// 当前状态正常 前序状态非正常时 触发恢复任务
go CronShared.SendTriggerTasks(ss.RecoverTriggerTasks, reporterServer.ID)
go CronShared.SendTriggerTasks(ss.RecoverTriggerTasks, reporterServer.ID, ss.UserID)
} else if lastStatus == StatusGood && lastStatus != stateCode {
// 前序状态正常 当前状态非正常时 触发失败任务
go CronShared.SendTriggerTasks(ss.FailTriggerTasks, reporterServer.ID)
go CronShared.SendTriggerTasks(ss.FailTriggerTasks, reporterServer.ID, ss.UserID)
}
}
}
@@ -0,0 +1,646 @@
package singleton
import (
"context"
"fmt"
"os"
"os/exec"
"strings"
"sync"
"testing"
"time"
"github.com/nezhahq/nezha/model"
)
// Regression markers for Finding 1 and Finding 2 of GHSA-jx78-55p5-rwv5
// (incomplete fix of GHSA-qjpp-gffx-2wm9).
const (
concurrentServerDeleteSuccessMarker = "ghsa-jx78-55p5-rwv5-finding1-no-crash"
deleteUnknownIDSuccessMarker = "ghsa-jx78-55p5-rwv5-finding2-no-zombie"
)
const serviceSentinelLifecycleSuccessMarker = "service-sentinel-stale-report-lifecycle-success"
func TestServiceSentinelReporterDeleteWaitsForSynchronousReportProcessing(t *testing.T) {
ctx, cancel := context.WithTimeout(t.Context(), time.Second)
defer cancel()
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"},
)
service := &model.Service{
Common: model.Common{ID: 10, UserID: 1},
Name: "lifecycle-service",
Type: model.TaskTypeTCPPing,
Target: "lifecycle.example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
}
addServiceMonitorSecurityService(t, ss, service)
reportValidated := make(chan struct{})
releaseReport := make(chan struct{})
var releaseOnce sync.Once
release := func() { releaseOnce.Do(func() { close(releaseReport) }) }
ss.serviceReportValidatedHook = func(serviceID uint64) {
if serviceID == service.ID {
close(reportValidated)
<-releaseReport
}
}
t.Cleanup(func() {
release()
ss.Close()
})
ss.Dispatch(serviceMonitorResult(1, service.ID, model.TaskTypeTCPPing, true))
select {
case <-reportValidated:
case <-ctx.Done():
t.Fatal(ctx.Err())
}
deleteDone := make(chan struct{})
go func() {
ServerShared.Delete([]uint64{1})
close(deleteDone)
}()
select {
case <-deleteDone:
t.Fatal("server deletion returned before the accepted report completed")
case <-time.After(25 * time.Millisecond):
}
release()
select {
case <-deleteDone:
case <-ctx.Done():
t.Fatal(ctx.Err())
}
var historyCount int64
if err := DB.Model(&model.ServiceHistory{}).
Where("service_id = ? AND server_id = ?", service.ID, 1).
Count(&historyCount).Error; err != nil {
t.Fatal(err)
}
if historyCount != 1 {
t.Fatalf("expected report side effects before deletion returned, got %d history rows", historyCount)
}
if _, ok := ServerShared.Get(1); ok {
t.Fatal("expected reporter to be deleted after the report completed")
}
}
func TestServiceSentinelWorkerRejectsReportAfterReporterDeletion(t *testing.T) {
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"},
)
service := &model.Service{
Common: model.Common{ID: 10, UserID: 1},
Name: "deleted-reporter-service",
Type: model.TaskTypeTCPPing,
Target: "deleted-reporter.example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
}
addServiceMonitorSecurityService(t, ss, service)
ServerShared.Delete([]uint64{1})
ss.Dispatch(serviceMonitorResult(1, service.ID, model.TaskTypeTCPPing, true))
ss.Close()
var historyCount int64
if err := DB.Model(&model.ServiceHistory{}).
Where("service_id = ? AND server_id = ?", service.ID, 1).
Count(&historyCount).Error; err != nil {
t.Fatal(err)
}
if historyCount != 0 {
t.Fatalf("expected no history after reporter deletion, got %d rows", historyCount)
}
ss.serviceResponseDataStoreLock.RLock()
_, pingCached := ss.serviceResponsePing[service.ID]
_, responseCached := ss.serviceResponseDataStore[service.ID]
stats := ss.serviceStatusToday[service.ID]
ss.serviceResponseDataStoreLock.RUnlock()
if pingCached {
t.Fatal("expected no ping cache side effect after reporter deletion")
}
if responseCached {
t.Fatal("expected no response cache side effect after reporter deletion")
}
if stats == nil || stats.Up != 0 || stats.Down != 0 {
t.Fatalf("expected no stats side effect after reporter deletion, got %+v", stats)
}
}
func TestServiceSentinelWorkerRecoversPerReportPanic(t *testing.T) {
ctx, cancel := context.WithTimeout(t.Context(), time.Second)
defer cancel()
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"},
)
panicService := &model.Service{
Common: model.Common{ID: 10, UserID: 1},
Name: "panic-service",
Type: model.TaskTypeHTTPGet,
Target: "https://panic.example.invalid",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
}
validService := &model.Service{
Common: model.Common{ID: 20, UserID: 1},
Name: "valid-service",
Type: model.TaskTypeTCPPing,
Target: "valid.example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
}
addServiceMonitorSecurityService(t, ss, panicService)
addServiceMonitorSecurityService(t, ss, validService)
ss.serviceReportBeforeTLSSideEffectsHook = func(serviceID uint64) {
if serviceID == panicService.ID {
panic("test service report panic")
}
}
ss.Dispatch(serviceMonitorResult(1, panicService.ID, model.TaskTypeHTTPGet, true))
ss.Dispatch(serviceMonitorResult(1, validService.ID, model.TaskTypeTCPPing, true))
waitForServiceHistory(t, validService.ID, 1)
ss.Close()
if !ss.serviceResponseDataStoreLock.TryLock() {
t.Fatal("panic leaked the service response lock")
}
ss.serviceResponseDataStoreLock.Unlock()
deleteDone := make(chan struct{})
go func() {
ServerShared.Delete([]uint64{1})
close(deleteDone)
}()
select {
case <-deleteDone:
case <-ctx.Done():
t.Fatal("panic leaked a lifecycle lock: " + ctx.Err().Error())
}
}
func TestServiceSentinelWorkerIgnoresStaleReportAfterDeletion(t *testing.T) {
if os.Getenv("NEZHA_SERVICE_SENTINEL_LIFECYCLE_CHILD") == "1" {
testServiceSentinelWorkerIgnoresStaleReportAfterDeletionChild(t)
return
}
// Given
ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
defer cancel()
child := exec.CommandContext(ctx, os.Args[0], "-test.run=^TestServiceSentinelWorkerIgnoresStaleReportAfterDeletion$")
child.Env = append(os.Environ(), "NEZHA_SERVICE_SENTINEL_LIFECYCLE_CHILD=1")
// When
output, err := child.CombinedOutput()
// Then
if ctx.Err() != nil {
t.Fatalf("service sentinel lifecycle child timed out: %v\n%s", ctx.Err(), output)
}
if err != nil {
t.Fatalf("service sentinel lifecycle child failed: %v\n%s", err, output)
}
if !strings.Contains(string(output), serviceSentinelLifecycleSuccessMarker) {
t.Fatalf("service sentinel lifecycle child did not report success:\n%s", output)
}
}
func testServiceSentinelWorkerIgnoresStaleReportAfterDeletionChild(t *testing.T) {
// Given
ctx, cancel := context.WithTimeout(t.Context(), time.Second)
defer cancel()
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"},
)
for _, service := range []*model.Service{
{
Common: model.Common{ID: 10, UserID: 1},
Name: "stale-service",
Type: model.TaskTypeTCPPing,
Target: "stale.example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
},
{
Common: model.Common{ID: 20, UserID: 1},
Name: "valid-service",
Type: model.TaskTypeTCPPing,
Target: "valid.example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
},
} {
addServiceMonitorSecurityService(t, ss, service)
}
acceptedStaleReport := make(chan struct{})
releaseWorker := make(chan struct{})
var releaseOnce sync.Once
releaseWorkerHook := func() {
releaseOnce.Do(func() { close(releaseWorker) })
}
ss.serviceReportValidatedHook = func(serviceID uint64) {
if serviceID == 10 {
close(acceptedStaleReport)
<-releaseWorker
}
}
t.Cleanup(func() {
releaseWorkerHook()
ss.Close()
})
// When
ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true))
select {
case <-acceptedStaleReport:
case <-ctx.Done():
t.Fatal(ctx.Err())
}
ss.Delete([]uint64{10})
releaseWorkerHook()
ss.Dispatch(serviceMonitorResult(1, 20, model.TaskTypeTCPPing, true))
ss.Close()
// Then
var staleHistoryCount int64
if err := DB.Model(&model.ServiceHistory{}).
Where("service_id = ? AND server_id = ?", 10, 1).
Count(&staleHistoryCount).Error; err != nil {
t.Fatal(err)
}
if staleHistoryCount != 0 {
t.Fatalf("expected stale service to write zero per-reporter history rows, got %d", staleHistoryCount)
}
var validHistoryCount int64
if err := DB.Model(&model.ServiceHistory{}).
Where("service_id = ? AND server_id = ?", 20, 1).
Count(&validHistoryCount).Error; err != nil {
t.Fatal(err)
}
if validHistoryCount != 1 {
t.Fatalf("expected exactly one valid service history row, got %d", validHistoryCount)
}
ss.serviceResponseDataStoreLock.RLock()
_, stalePingCached := ss.serviceResponsePing[10]
validStats := ss.serviceStatusToday[20]
ss.serviceResponseDataStoreLock.RUnlock()
if stalePingCached {
t.Fatal("expected stale service ping cache to be deleted")
}
if validStats == nil || validStats.Up != 1 || validStats.Down != 0 {
t.Fatalf("expected valid service stats up=1 down=0, got %+v", validStats)
}
if _, err := fmt.Fprintln(os.Stdout, serviceSentinelLifecycleSuccessMarker); err != nil {
t.Fatal(err)
}
}
func TestServiceSentinelWorkerRevalidatesReportAfterUpdate(t *testing.T) {
// Given
ctx, cancel := context.WithTimeout(t.Context(), time.Second)
defer cancel()
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"},
)
service := &model.Service{
Common: model.Common{ID: 10, UserID: 1},
Name: "updatable-service",
Type: model.TaskTypeTCPPing,
Target: "updatable.example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
}
addServiceMonitorSecurityService(t, ss, service)
ss.serviceResponseDataStoreLock.Lock()
ss.serviceStatusToday[service.ID] = &_TodayStatsOfService{Up: 7, Down: 3, Delay: 12.5}
ss.serviceResponseDataStoreLock.Unlock()
acceptedReport := make(chan struct{})
releaseWorker := make(chan struct{})
var releaseOnce sync.Once
releaseWorkerHook := func() {
releaseOnce.Do(func() { close(releaseWorker) })
}
ss.serviceReportValidatedHook = func(serviceID uint64) {
if serviceID == service.ID {
close(acceptedReport)
<-releaseWorker
}
}
t.Cleanup(func() {
releaseWorkerHook()
ss.Close()
})
// When
ss.Dispatch(serviceMonitorResult(1, service.ID, model.TaskTypeTCPPing, true))
select {
case <-acceptedReport:
case <-ctx.Done():
t.Fatal(ctx.Err())
}
updatedService := *service
updatedService.Name = "updated-service"
updatedService.SkipServers = map[uint64]bool{}
if err := ss.Update(&updatedService); err != nil {
t.Fatal(err)
}
releaseWorkerHook()
ss.Close()
// Then
var historyCount int64
if err := DB.Model(&model.ServiceHistory{}).
Where("service_id = ? AND server_id = ?", service.ID, 1).
Count(&historyCount).Error; err != nil {
t.Fatal(err)
}
if historyCount != 0 {
t.Fatalf("expected updated service to write zero per-reporter history rows, got %d", historyCount)
}
ss.serviceResponseDataStoreLock.RLock()
_, pingCached := ss.serviceResponsePing[service.ID]
stats := ss.serviceStatusToday[service.ID]
ss.serviceResponseDataStoreLock.RUnlock()
if pingCached {
t.Fatal("expected updated service report to leave no ping cache entry")
}
if stats == nil || stats.Up != 7 || stats.Down != 3 || stats.Delay != 12.5 {
t.Fatalf("expected existing service stats to remain unchanged, got %+v", stats)
}
currentService, ok := ss.Get(service.ID)
if !ok || currentService.Name != updatedService.Name || currentService.SkipServers[1] {
t.Fatalf("expected updated service configuration, got %+v", currentService)
}
}
func TestServiceSentinelLoadStatsFollowsLifecycleLockOrder(t *testing.T) {
// Given
ss := &ServiceSentinel{
serviceStatusToday: make(map[uint64]*_TodayStatsOfService),
serviceResponseDataStore: make(map[uint64]serviceResponseData),
services: make(map[uint64]*model.Service),
monthlyStatus: make(map[uint64]*serviceResponseItem),
}
ss.loadStatsResponseLockedHook = func() {
if ss.serviceResponseDataStoreLock.TryLock() {
ss.serviceResponseDataStoreLock.Unlock()
t.Fatal("LoadStats invoked the hook before acquiring the response read lock")
}
if !ss.monthlyStatusLock.TryLock() {
t.Fatal("LoadStats acquired monthlyStatusLock before the response lock hook")
}
ss.monthlyStatusLock.Unlock()
if !ss.servicesLock.TryLock() {
t.Fatal("LoadStats acquired servicesLock before the response lock hook")
}
ss.servicesLock.Unlock()
}
// When / Then
ss.LoadStats()
}
func TestServiceSentinelWorkerHoldsResponseLockDuringTLSSideEffects(t *testing.T) {
// Given
ctx, cancel := context.WithTimeout(t.Context(), time.Second)
defer cancel()
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"},
)
service := &model.Service{
Common: model.Common{ID: 10, UserID: 1},
Name: "tls-service",
Type: model.TaskTypeHTTPGet,
Target: "https://tls.example.invalid",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
}
addServiceMonitorSecurityService(t, ss, service)
tlsSideEffectsReady := make(chan struct{})
releaseTLSSideEffects := make(chan struct{})
var releaseOnce sync.Once
releaseTLSSideEffectsHook := func() {
releaseOnce.Do(func() { close(releaseTLSSideEffects) })
}
ss.serviceReportBeforeTLSSideEffectsHook = func(serviceID uint64) {
if serviceID == service.ID {
close(tlsSideEffectsReady)
<-releaseTLSSideEffects
}
}
t.Cleanup(func() {
releaseTLSSideEffectsHook()
ss.Close()
})
report := serviceMonitorResult(1, service.ID, model.TaskTypeHTTPGet, true)
report.Data.Data = "issuer|2030-01-02 15:04:05 +0000 UTC"
// When
ss.Dispatch(report)
select {
case <-tlsSideEffectsReady:
case <-ctx.Done():
t.Fatal(ctx.Err())
}
responseLockAcquired := ss.serviceResponseDataStoreLock.TryLock()
if responseLockAcquired {
ss.serviceResponseDataStoreLock.Unlock()
t.Fatal("worker released the response lock before TLS side effects")
}
releaseTLSSideEffectsHook()
ss.Close()
// Then
ss.serviceResponseDataStoreLock.RLock()
cachedCertificate := ss.tlsCertCache[service.ID]
ss.serviceResponseDataStoreLock.RUnlock()
if cachedCertificate != report.Data.Data {
t.Fatalf("expected TLS cache %q, got %q", report.Data.Data, cachedCertificate)
}
}
// TestServiceSentinelWorkerSurvivesConcurrentReporterServerDelete is a
// regression test for GHSA-jx78-55p5-rwv5 Finding 1 (incomplete fix of
// GHSA-qjpp-gffx-2wm9).
//
// The vulnerability: after the 2026-07-21 fix, the worker re-validates the
// service under serviceResponseDataStoreLock, but then takes a fresh snapshot
// m := ServerShared.GetList() with no guard. A concurrent batch-delete of the
// reporter's own server removes it between the pre-lock validation and the
// GetList call, so m[r.Reporter] is nil. delayCheck and notifyCheck then
// dereference m[r.Reporter].Name unconditionally — SIGSEGV.
//
// The subprocess-isolation pattern is used because the pre-fix code path
// panicked (nil pointer dereference in an unrecovered goroutine), which would
// crash the whole test binary rather than simply failing a single test.
func TestServiceSentinelWorkerSurvivesConcurrentReporterServerDelete(t *testing.T) {
if os.Getenv("NEZHA_SENTINEL_CONCURRENT_DELETE_CHILD") == "1" {
testServiceSentinelWorkerSurvivesConcurrentReporterServerDeleteChild(t)
return
}
ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
defer cancel()
child := exec.CommandContext(ctx, os.Args[0],
"-test.run=^TestServiceSentinelWorkerSurvivesConcurrentReporterServerDelete$",
"-test.v",
)
child.Env = append(os.Environ(), "NEZHA_SENTINEL_CONCURRENT_DELETE_CHILD=1")
output, err := child.CombinedOutput()
if ctx.Err() != nil {
t.Fatalf("child process timed out: %v\n%s", ctx.Err(), output)
}
if err != nil {
t.Fatalf("child process crashed (likely nil deref in delayCheck/notifyCheck): %v\n%s", err, output)
}
if !strings.Contains(string(output), concurrentServerDeleteSuccessMarker) {
t.Fatalf("child did not print success marker:\n%s", output)
}
}
func testServiceSentinelWorkerSurvivesConcurrentReporterServerDeleteChild(t *testing.T) {
// Given: a reporter server and a service with latency-alerting enabled so
// that delayCheck (the vulnerable sink at line 785) is exercised on every
// dispatch. MaxLatency=1 ensures delay=12 always exceeds the threshold and
// the notification branch (not just the mute-clear branch) is taken.
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"},
)
service := &model.Service{
Common: model.Common{ID: 10, UserID: 1},
Name: "latency-service",
Type: model.TaskTypeTCPPing,
Target: "example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
LatencyNotify: true,
MaxLatency: 1,
}
addServiceMonitorSecurityService(t, ss, service)
reportProcessing := make(chan struct{})
releaseWorker := make(chan struct{})
var releaseOnce sync.Once
releaseWorkerFn := func() { releaseOnce.Do(func() { close(releaseWorker) }) }
// serviceReportValidatedHook runs while the report holds the lifecycle read
// lock. Deletion must therefore run in another goroutine and wait until this
// hook releases; attempting Delete here would try to upgrade the RWMutex.
ss.serviceReportValidatedHook = func(serviceID uint64) {
if serviceID == service.ID {
close(reportProcessing)
<-releaseWorker
}
}
t.Cleanup(func() {
releaseWorkerFn()
ss.Close()
})
// When
ss.Dispatch(serviceMonitorResult(1, service.ID, model.TaskTypeTCPPing, true))
select {
case <-reportProcessing:
case <-t.Context().Done():
t.Fatal(t.Context().Err())
}
deleteDone := make(chan struct{})
go func() {
ServerShared.Delete([]uint64{1})
close(deleteDone)
}()
releaseWorkerFn()
select {
case <-deleteDone:
case <-t.Context().Done():
t.Fatal(t.Context().Err())
}
ss.Close()
// Then: no crash; the worker handled the nil reporter gracefully.
if _, err := fmt.Fprintln(os.Stdout, concurrentServerDeleteSuccessMarker); err != nil {
t.Fatal(err)
}
}
// TestServiceSentinelDeleteWithUnknownIDDoesNotLeaveZombies is a regression
// test for GHSA-jx78-55p5-rwv5 Finding 2 (low severity).
//
// The vulnerability: ServiceSentinel.Delete iterates the caller-supplied id
// slice and does CronShared.Remove(ss.services[id].CronJobID) without checking
// whether id is present in ss.services. CheckPermission returns vacuously
// true for unknown ids, so the controller layer cannot block this path.
// ss.services[unknownID] returns nil, and .CronJobID panics. Because the
// panic aborts the loop, every id ordered AFTER the bogus one is never removed
// from the in-memory registry even though its database row was already deleted,
// producing zombie services that keep dispatching cron probes.
func TestServiceSentinelDeleteWithUnknownIDDoesNotLeaveZombies(t *testing.T) {
// Given: one legitimate service (ID 10) registered in the sentinel.
ss := newServiceMonitorSecurityHarness(t,
&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"},
)
service := &model.Service{
Common: model.Common{ID: 10, UserID: 1},
Name: "real-service",
Type: model.TaskTypeTCPPing,
Target: "example.invalid:443",
Duration: 3600,
Cover: model.ServiceCoverIgnoreAll,
SkipServers: map[uint64]bool{1: true},
}
addServiceMonitorSecurityService(t, ss, service)
// When: Delete is called with a bogus ID first, then the real service ID.
// Before the fix this panicked on ss.services[99999].CronJobID and left
// service 10 as a zombie.
ss.Delete([]uint64{99999, service.ID})
// Then: the real service must be fully removed from every in-memory map.
ss.serviceResponseDataStoreLock.RLock()
_, todayPresent := ss.serviceStatusToday[service.ID]
_, pingPresent := ss.serviceResponsePing[service.ID]
ss.serviceResponseDataStoreLock.RUnlock()
ss.servicesLock.RLock()
_, servicePresent := ss.services[service.ID]
ss.servicesLock.RUnlock()
ss.monthlyStatusLock.Lock()
_, monthlyPresent := ss.monthlyStatus[service.ID]
ss.monthlyStatusLock.Unlock()
if todayPresent {
t.Error("zombie: serviceStatusToday still contains the deleted service")
}
if pingPresent {
t.Error("zombie: serviceResponsePing still contains the deleted service")
}
if servicePresent {
t.Error("zombie: services map still contains the deleted service")
}
if monthlyPresent {
t.Error("zombie: monthlyStatus still contains the deleted service")
}
if _, err := fmt.Fprintln(os.Stdout, deleteUnknownIDSuccessMarker); err != nil {
t.Fatal(err)
}
}
+40 -8
View File
@@ -11,7 +11,6 @@ import (
"github.com/gin-gonic/gin"
"github.com/patrickmn/go-cache"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"sigs.k8s.io/yaml"
@@ -34,6 +33,10 @@ var (
NotificationShared *NotificationClass
NATShared *NATClass
CronShared *CronClass
// ServerTransferShared is initialized in LoadSingleton AFTER ServerShared
// (so the in-memory pending index can write back into ServerShared.UserID
// on transitions) and AFTER initUser (so PushIfOnline can read secrets
// from UserInfoMap).
)
//go:embed frontend-templates.yaml
@@ -59,6 +62,7 @@ func LoadSingleton(bus chan<- *model.Service) (err error) {
NotificationShared = NewNotificationClass()
ServerShared = NewServerClass()
CronShared = NewCronClass()
ServerTransferShared = NewServerTransferClass()
// 最后初始化 ServiceSentinel
ServiceSentinelShared, err = NewServiceSentinel(bus)
if err == nil {
@@ -79,7 +83,7 @@ func InitFrontendTemplates() error {
// InitDBFromPath 从给出的文件路径中加载数据库
func InitDBFromPath(path string) error {
var err error
DB, err = gorm.Open(sqlite.Open(path), &gorm.Config{
DB, err = gorm.Open(openSQLiteDialector(path), &gorm.Config{
CreateBatchSize: 200,
})
if err != nil {
@@ -92,11 +96,22 @@ func InitDBFromPath(path string) error {
model.Notification{}, model.AlertRule{}, model.Service{}, model.NotificationGroupNotification{},
model.Cron{}, model.Transfer{}, model.ServerGroupServer{},
model.NAT{}, model.DDNSProfile{}, model.NotificationGroupNotification{},
model.WAF{}, model.Oauth2Bind{}, model.Domain{})
model.WAF{}, model.Oauth2Bind{}, model.Domain{}, model.ServerTransfer{}, model.JWTSession{},
model.APIToken{}, model.MCPAuditLog{})
if err != nil {
return err
}
// 旧 mcp:* scope 与 nezha:* 并行了一段时间,HasScope 通过别名让 mcp:fs:write
// 静默扩到 REST nezha:server:write。统一命名后这里把残留旧 scope 一次性
// 归一化(或在仅剩危险旧 scope 时整张 PAT 删除),保证运行时不再依赖别名。
if rewritten, deleted, mErr := model.MigrateLegacyMCPScopes(DB); mErr != nil {
log.Printf("NEZHA>> MigrateLegacyMCPScopes failed: %v", mErr)
} else if rewritten > 0 || deleted > 0 {
log.Printf("NEZHA>> Migrated legacy mcp:* api token scopes: rewritten=%d deleted=%d", rewritten, deleted)
}
return nil
}
@@ -114,16 +129,15 @@ func RecordTransferHourlyUsage(servers ...*model.Server) {
}
for server := range slist {
_, _, deltaIn, deltaOut := server.TransferDeltaAndAdvance()
tx := model.Transfer{
ServerID: server.ID,
In: utils.SubUintChecked(server.State.NetInTransfer, server.PrevTransferInSnapshot),
Out: utils.SubUintChecked(server.State.NetOutTransfer, server.PrevTransferOutSnapshot),
In: deltaIn,
Out: deltaOut,
}
if tx.In == 0 && tx.Out == 0 {
continue
}
server.PrevTransferInSnapshot = server.State.NetInTransfer
server.PrevTransferOutSnapshot = server.State.NetOutTransfer
tx.CreatedAt = nowTrimSeconds
txs = append(txs, tx)
}
@@ -134,6 +148,13 @@ func RecordTransferHourlyUsage(servers ...*model.Server) {
log.Printf("NEZHA>> Saved traffic metrics to database. Affected %d row(s), Error: %v", len(txs), DB.Create(txs).Error)
}
func PersistTransfer(transfer model.Transfer) error {
if transfer.In == 0 && transfer.Out == 0 {
return nil
}
return DB.Create(&transfer).Error
}
// CleanMonitorHistory 清理流量记录(TSDB 有自己的保留策略)
func CleanMonitorHistory() {
// 清理已被删除的服务器的流量记录
@@ -143,7 +164,10 @@ func CleanMonitorHistory() {
specialServerKeep := make(map[uint64]time.Time)
var specialServerIDs []uint64
var alerts []model.AlertRule
DB.Find(&alerts)
if err := DB.Find(&alerts).Error; err != nil {
log.Printf("NEZHA>> Failed to load alert rules while cleaning transfer history: %v", err)
return
}
for _, alert := range alerts {
for _, rule := range alert.Rules {
// 是不是流量记录规则
@@ -171,6 +195,14 @@ func CleanMonitorHistory() {
for id, couldRemove := range specialServerKeep {
DB.Unscoped().Delete(&model.Transfer{}, "server_id = ? AND datetime(`created_at`) < datetime(?)", id, couldRemove)
}
if len(specialServerIDs) == 0 {
if allServerKeep.IsZero() {
DB.Unscoped().Session(&gorm.Session{AllowGlobalUpdate: true}).Delete(&model.Transfer{})
} else {
DB.Unscoped().Delete(&model.Transfer{}, "datetime(`created_at`) < datetime(?)", allServerKeep)
}
return
}
if allServerKeep.IsZero() {
DB.Unscoped().Delete(&model.Transfer{}, "server_id NOT IN (?)", specialServerIDs)
} else {
@@ -0,0 +1,83 @@
//go:build agentcompat && linux
package singleton
import (
"errors"
"path/filepath"
"strings"
"testing"
)
func TestSQLiteAttributionErrorsDoNotExposeDatabasePath(t *testing.T) {
// Given
databasePath := filepath.Join(t.TempDir(), "private-dashboard.sqlite")
unsupportedDSN := "file:" + databasePath + "?mode=memory"
missingJournal := filepath.Join(t.TempDir(), "missing-journal")
// When
memoryDatabase, dsnErr := openSQLiteAttributionTestDB(unsupportedDSN)
if dsnErr == nil {
dsnErr = memoryDatabase.Ping()
}
if memoryDatabase != nil {
t.Cleanup(func() {
if closeErr := memoryDatabase.Close(); closeErr != nil {
t.Error(closeErr)
}
})
}
_, journalErr := sqliteAttributionOpenJournalDescriptor(missingJournal)
// Then
if !errors.Is(dsnErr, ErrSQLiteAttributionUnsupportedDSN) {
t.Fatal("file URI error does not wrap the typed unsupported DSN error")
}
if !errors.Is(journalErr, ErrSQLiteAttributionJournalIdentity) {
t.Fatal("journal error does not wrap the typed journal identity error")
}
if strings.Contains(dsnErr.Error(), databasePath) || strings.Contains(dsnErr.Error(), unsupportedDSN) {
t.Fatal("unsupported DSN error exposes its database path")
}
if strings.Contains(journalErr.Error(), missingJournal) {
t.Fatal("journal identity error exposes its database path")
}
}
func TestSQLiteAttributionAcceptsOnDiskFileURIAndRejectsInMemoryDSN(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
databasePath := filepath.Join(t.TempDir(), "dashboard.sqlite")
// When
fileDatabase, fileErr := openSQLiteAttributionTestDB("file:" + databasePath)
if fileDatabase != nil {
t.Cleanup(func() {
if closeErr := fileDatabase.Close(); closeErr != nil {
t.Error(closeErr)
}
})
}
if fileErr == nil {
fileErr = fileDatabase.Ping()
}
memoryDatabase, memoryErr := openSQLiteAttributionTestDB(":memory:")
if memoryDatabase != nil {
t.Cleanup(func() {
if closeErr := memoryDatabase.Close(); closeErr != nil {
t.Error(closeErr)
}
})
}
if memoryErr == nil {
memoryErr = memoryDatabase.Ping()
}
// Then
if fileErr != nil {
t.Fatal("file URI for an on-disk database was rejected")
}
if !errors.Is(memoryErr, ErrSQLiteAttributionUnsupportedDSN) {
t.Fatal("in-memory DSN does not return the typed unsupported DSN error")
}
}
@@ -0,0 +1,180 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"database/sql"
"database/sql/driver"
"errors"
"testing"
"golang.org/x/sys/unix"
)
func TestSQLiteAttributionConnectionCloseFinalizesActiveWriteTransaction(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
databasePath := sqliteAttributionTestDatabasePath(t)
rawConnection, err := sqliteAttributionDriver{}.Open(databasePath)
if err != nil {
t.Fatal(err)
}
connection := rawConnection.(*sqliteAttributionConnection)
if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil {
t.Fatal(err)
}
enableSQLiteAttribution()
transaction, err := connection.Begin()
if err != nil {
t.Fatal(err)
}
statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)")
if err != nil {
t.Fatal(err)
}
if _, err := statement.Exec([]driver.Value{"close-active"}); err != nil {
t.Fatal(err)
}
if err := statement.Close(); err != nil {
t.Fatal(err)
}
identity, descriptor, active := sqliteAttributionTransactionState(t, connection)
if !active {
t.Fatal("active transaction is missing before Close")
}
// When
firstCloseErr := connection.Close()
secondCloseErr := connection.Close()
commitErr := transaction.Commit()
rollbackErr := transaction.Rollback()
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
standard, openErr := sql.Open("sqlite3", databasePath)
if openErr != nil {
t.Fatal(openErr)
}
defer standard.Close()
var count int
countErr := standard.QueryRow("SELECT COUNT(*) FROM settings").Scan(&count)
tracker := sqliteAttributionTracker.Load()
tracker.mu.Lock()
_, active = tracker.transactions[identity]
tracker.mu.Unlock()
// Then
if firstCloseErr != nil || secondCloseErr != nil {
t.Fatalf("Close errors = %v / %v", firstCloseErr, secondCloseErr)
}
if !errors.Is(commitErr, driver.ErrBadConn) {
t.Fatalf("Commit after Close error = %v, want driver.ErrBadConn", commitErr)
}
if !errors.Is(rollbackErr, driver.ErrBadConn) {
t.Fatalf("Rollback after Close error = %v, want driver.ErrBadConn", rollbackErr)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Fatalf("journal descriptor after Close = %v, want EBADF", descriptorErr)
}
if countErr != nil || count != 0 {
t.Fatalf("closed active transaction persisted count=%d err=%v", count, countErr)
}
if active {
t.Fatal("Close left the tracker transaction active")
}
}
func TestSQLiteAttributionConnectionCloseWakesSelectedCommit(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
databasePath := sqliteAttributionTestDatabasePath(t)
rawConnection, err := sqliteAttributionDriver{}.Open(databasePath)
if err != nil {
t.Fatal(err)
}
connection := rawConnection.(*sqliteAttributionConnection)
if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil {
t.Fatal(err)
}
enableSQLiteAttribution()
transaction, err := connection.Begin()
if err != nil {
t.Fatal(err)
}
statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)")
if err != nil {
t.Fatal(err)
}
if _, err := statement.Exec([]driver.Value{"close-wakes-commit"}); err != nil {
t.Fatal(err)
}
if err := statement.Close(); err != nil {
t.Fatal(err)
}
_, descriptor, active := sqliteAttributionTransactionState(t, connection)
if !active {
t.Fatal("active transaction is missing before selected Commit")
}
tracker := sqliteAttributionTracker.Load()
session, err := tracker.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
finalizing := make(chan error, 1)
commit := make(chan error, 1)
go func() {
_, waitErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing)
finalizing <- waitErr
}()
// When
go func() { commit <- transaction.Commit() }()
if finalizingErr := <-finalizing; finalizingErr != nil {
t.Fatalf("Commit finalization wait error = %v", finalizingErr)
}
closeErr := connection.Close()
commitErr := <-commit
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
standard, err := sql.Open("sqlite3", databasePath)
if err != nil {
t.Fatal(err)
}
defer standard.Close()
var count int
if err := standard.QueryRow("SELECT COUNT(*) FROM settings").Scan(&count); err != nil {
t.Fatal(err)
}
// Then
if closeErr != nil {
t.Fatal(closeErr)
}
var holdErr *SQLiteHoldError
if !errors.As(commitErr, &holdErr) || !errors.Is(commitErr, ErrSQLiteHoldAborted) {
t.Fatalf("Commit error after Close = %v", commitErr)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Fatalf("journal descriptor after Close = %v, want EBADF", descriptorErr)
}
if count != 0 {
t.Fatalf("Close while Commit waited persisted %d rows", count)
}
}
type sqliteAttributionBadConnRows struct{}
func (sqliteAttributionBadConnRows) Columns() []string { return []string{"value"} }
func (sqliteAttributionBadConnRows) Close() error { return nil }
func (sqliteAttributionBadConnRows) Next([]driver.Value) error { return driver.ErrBadConn }
func TestSQLiteAttributionReadonlyRowsReturnBadConnWithoutPanic(t *testing.T) {
// Given
rows := &sqliteAttributionRows{rows: sqliteAttributionBadConnRows{}}
// When
err := rows.Next(make([]driver.Value, 1))
// Then
if !errors.Is(err, driver.ErrBadConn) {
t.Fatalf("readonly rows Next error = %v, want driver.ErrBadConn", err)
}
}
@@ -0,0 +1,102 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"errors"
"testing"
"golang.org/x/sys/unix"
)
func TestSQLiteAttributionCommitReleaseLinearizesBeforeContextCancellation(t *testing.T) {
// Given
transactionContext, cancel := context.WithCancel(context.Background())
defer cancel()
connection, transaction, session, databasePath := sqliteAttributionHeldTransaction(t, "released-before-cancel", transactionContext)
identity, descriptor, active := sqliteAttributionTransactionState(t, connection)
if !active {
t.Fatal("released transaction is not active")
}
// When
commit := sqliteAttributionStartHeldCommit(t, transaction, session)
if err := sqliteAttributionTracker.Load().ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
cancel()
commitErr := <-commit
terminal, terminalErr := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing)
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
// Then
if commitErr != nil {
t.Fatalf("released Commit error after context cancellation = %v", commitErr)
}
if terminalErr != nil || !terminal.Released || !terminal.Selected || !terminal.Finalizing {
t.Fatalf("released terminal=%+v err=%v", terminal, terminalErr)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Fatalf("journal descriptor after released Commit = %v, want EBADF", descriptorErr)
}
if count := sqliteAttributionPersistedCount(t, databasePath); count != 1 {
t.Fatalf("released Commit persisted %d rows", count)
}
if _, _, active := sqliteAttributionTransactionState(t, connection); active {
t.Fatal("released Commit left the connection transaction active")
}
if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) {
t.Fatal("released Commit left the tracker transaction active")
}
}
func TestSQLiteAttributionCommitCancellationAbortsBeforeRelease(t *testing.T) {
// Given
transactionContext, cancel := context.WithCancel(context.Background())
defer cancel()
connection, transaction, session, databasePath := sqliteAttributionHeldTransaction(t, "cancelled-before-release", transactionContext)
identity, descriptor, active := sqliteAttributionTransactionState(t, connection)
if !active {
t.Fatal("cancelled transaction is not active")
}
// When
commit := sqliteAttributionStartHeldCommit(t, transaction, session)
cancel()
commitErr := <-commit
releaseErr := sqliteAttributionTracker.Load().ReleaseSQLiteHold(session)
_, terminalErr := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing)
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
tracker := sqliteAttributionTracker.Load()
tracker.mu.Lock()
terminal := tracker.terminal
tracker.mu.Unlock()
// Then
var holdErr *SQLiteHoldError
if !errors.As(commitErr, &holdErr) || !errors.Is(commitErr, ErrSQLiteHoldAborted) || !errors.Is(commitErr, context.Canceled) {
t.Fatalf("cancelled Commit error = %v", commitErr)
}
if !errors.Is(releaseErr, ErrSQLiteHoldStaleSession) {
t.Fatalf("release after cancellation-owned abort = %v", releaseErr)
}
if !errors.Is(terminalErr, ErrSQLiteHoldAborted) {
t.Fatalf("cancelled terminal wait error = %v, want aborted", terminalErr)
}
if terminal == nil || terminal.released {
t.Fatalf("cancelled terminal state = %+v, want aborted and unreleased", terminal)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Fatalf("journal descriptor after cancelled Commit = %v, want EBADF", descriptorErr)
}
if count := sqliteAttributionPersistedCount(t, databasePath); count != 0 {
t.Fatalf("cancelled Commit persisted %d rows", count)
}
if _, _, active := sqliteAttributionTransactionState(t, connection); active {
t.Fatal("cancelled Commit left the connection transaction active")
}
if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) {
t.Fatal("cancelled Commit left the tracker transaction active")
}
}
@@ -0,0 +1,154 @@
//go:build agentcompat && linux
package singleton
import (
"database/sql/driver"
"errors"
"sync"
"sync/atomic"
"testing"
)
type sqliteAttributionBlockingTx struct {
commitStarted chan struct{}
allowCommit chan struct{}
commitCalls atomic.Int32
rollbackCalls atomic.Int32
lifecycleMu *sync.Mutex
lockFree atomic.Bool
}
func (transaction *sqliteAttributionBlockingTx) Commit() error {
transaction.commitCalls.Add(1)
// Probe before publishing entry so losing terminal calls cannot contend with this boundary check.
if transaction.lifecycleMu.TryLock() {
transaction.lockFree.Store(true)
transaction.lifecycleMu.Unlock()
}
close(transaction.commitStarted)
<-transaction.allowCommit
return nil
}
func (transaction *sqliteAttributionBlockingTx) Rollback() error {
transaction.rollbackCalls.Add(1)
return nil
}
func TestSQLiteAttributionTransactionCompletionRunsRawCommitExactlyOnce(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
identity := sqliteHoldTestTransaction(201)
if err := tracker.BeginSQLiteTransaction(identity); err != nil {
t.Fatal(err)
}
raw := &sqliteAttributionBlockingTx{commitStarted: make(chan struct{}), allowCommit: make(chan struct{})}
state := &sqliteAttributionTransaction{
transaction: identity,
raw: raw,
tracker: tracker,
journalFD: -1,
done: make(chan struct{}),
}
var rawCloseCalls atomic.Int32
connection := &sqliteAttributionConnection{transaction: state}
connection.closeRawConnection = func() error {
if !connection.lifecycleMu.TryLock() {
return errors.New("raw connection Close ran while lifecycle lock was held")
}
connection.lifecycleMu.Unlock()
rawCloseCalls.Add(1)
return nil
}
raw.lifecycleMu = &connection.lifecycleMu
owner := &sqliteAttributionTx{connection: connection, state: state}
secondCommit := &sqliteAttributionTx{connection: connection, state: state}
commitResult := make(chan error, 1)
secondCommitResult := make(chan error, 1)
rollbackResult := make(chan error, 1)
closeResult := make(chan error, 1)
secondCommitStarted := make(chan struct{})
rollbackStarted := make(chan struct{})
closeStarted := make(chan struct{})
// When
go func() { commitResult <- owner.Commit() }()
<-raw.commitStarted
go func() {
close(secondCommitStarted)
secondCommitResult <- secondCommit.Commit()
}()
go func() {
close(rollbackStarted)
rollbackResult <- owner.Rollback()
}()
go func() {
close(closeStarted)
closeResult <- connection.Close()
}()
<-secondCommitStarted
<-rollbackStarted
<-closeStarted
for _, result := range []<-chan error{secondCommitResult, rollbackResult} {
select {
case err := <-result:
t.Fatalf("completion loser returned before raw Commit completion: %v", err)
default:
}
}
select {
case <-state.done:
t.Fatal("completion published before raw Commit was allowed to finish")
default:
}
if calls := rawCloseCalls.Load(); calls != 0 {
t.Fatalf("raw connection Close calls while raw Commit was in flight = %d, want 0", calls)
}
if !raw.lockFree.Load() {
t.Fatal("raw Commit ran while lifecycle lock was held")
}
if !sqliteAttributionTrackerTransactionActive(tracker, identity) {
t.Fatal("tracker transaction became inactive while raw Commit was blocked")
}
close(raw.allowCommit)
commitErr := <-commitResult
secondCommitErr := <-secondCommitResult
rollbackErr := <-rollbackResult
closeErr := <-closeResult
// Then
if commitErr != nil {
t.Fatalf("raw Commit error = %v", commitErr)
}
for _, loserErr := range []error{secondCommitErr, rollbackErr} {
if !errors.Is(loserErr, driver.ErrBadConn) {
t.Fatalf("completion loser error = %v, want driver.ErrBadConn", loserErr)
}
}
if closeErr != nil && !errors.Is(closeErr, driver.ErrBadConn) {
t.Fatalf("connection Close error = %v, want nil or driver.ErrBadConn", closeErr)
}
select {
case <-state.done:
default:
t.Fatal("completion signal did not close after raw Commit finished")
}
if calls := raw.commitCalls.Load(); calls != 1 {
t.Fatalf("raw Commit calls = %d, want 1", calls)
}
if calls := rawCloseCalls.Load(); calls != 1 {
t.Fatalf("raw connection Close calls after raw Commit completed = %d, want 1", calls)
}
if calls := raw.rollbackCalls.Load(); calls != 0 {
t.Fatalf("raw Rollback calls while raw Commit owned completion = %d, want 0", calls)
}
if _, _, active := sqliteAttributionTransactionState(t, connection); active {
t.Fatal("completed transaction remained attached to the connection")
}
if sqliteAttributionTrackerTransactionActive(tracker, identity) {
t.Fatal("completed transaction remained active in the tracker")
}
}
var _ driver.Tx = (*sqliteAttributionBlockingTx)(nil)
@@ -0,0 +1,26 @@
//go:build agentcompat && linux
package singleton
import (
"errors"
)
func (connection *sqliteAttributionConnection) Close() error {
connection.lifecycleMu.Lock()
state := connection.transaction
connection.lifecycleMu.Unlock()
if state == nil {
return connection.closeRaw()
}
// database/sql may close Conn before Tx reaches Rollback; converge the attribution state here.
rollbackErr := state.rollback(connection)
return errors.Join(rollbackErr, connection.closeRaw())
}
func (connection *sqliteAttributionConnection) closeRaw() error {
if connection.closeRawConnection != nil {
return connection.closeRawConnection()
}
return connection.connection.Close()
}
@@ -0,0 +1,52 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"database/sql/driver"
"errors"
)
func (connection *sqliteAttributionConnection) Exec(query string, values []driver.Value) (driver.Result, error) {
statement, err := connection.Prepare(query)
if err != nil {
return nil, err
}
defer statement.Close()
return statement.Exec(values)
}
func (connection *sqliteAttributionConnection) ExecContext(ctx context.Context, query string, values []driver.NamedValue) (driver.Result, error) {
statement, err := connection.PrepareContext(ctx, query)
if err != nil {
return nil, err
}
defer statement.Close()
return statement.(driver.StmtExecContext).ExecContext(ctx, values)
}
func (connection *sqliteAttributionConnection) Query(query string, values []driver.Value) (driver.Rows, error) {
statement, err := connection.Prepare(query)
if err != nil {
return nil, err
}
wrapped := statement.(*sqliteAttributionStatement)
return wrapped.queryOwned(func() (driver.Rows, error) { return wrapped.statement.Query(values) }, statement)
}
func (connection *sqliteAttributionConnection) QueryContext(ctx context.Context, query string, values []driver.NamedValue) (driver.Rows, error) {
statement, err := connection.PrepareContext(ctx, query)
if err != nil {
return nil, err
}
return connection.queryContextStatement(ctx, statement.(*sqliteAttributionStatement), values, statement)
}
func (connection *sqliteAttributionConnection) queryContextStatement(ctx context.Context, statement *sqliteAttributionStatement, values []driver.NamedValue, owner driver.Stmt) (driver.Rows, error) {
contextStatement, ok := statement.statement.(driver.StmtQueryContext)
if !ok {
return nil, errors.Join(driver.ErrSkip, owner.Close())
}
return statement.queryOwned(func() (driver.Rows, error) { return contextStatement.QueryContext(ctx, values) }, owner)
}
@@ -0,0 +1,149 @@
//go:build agentcompat && linux
package singleton
import (
"errors"
"runtime"
"strings"
"golang.org/x/sys/unix"
)
func (connection *sqliteAttributionConnection) beforeWrite(classification sqliteAttributionClassification) error {
if !sqliteAttributionEnabled.Load() || classification.readonly {
return nil
}
if !classification.valid() {
return &SQLiteAttributionError{Cause: ErrSQLiteAttributionUnsupportedWrite}
}
connection.lifecycleMu.Lock()
state := connection.transaction
if state == nil {
connection.lifecycleMu.Unlock()
return &SQLiteAttributionError{Cause: ErrSQLiteAttributionUnboundWrite}
}
poison := state.poison
connection.lifecycleMu.Unlock()
if poison != nil {
return poison
}
connection.execution = &sqliteAttributionExecution{classification: classification, origin: sqliteAttributionOrigin()}
return nil
}
func (connection *sqliteAttributionConnection) discardExecution() { connection.execution = nil }
func (connection *sqliteAttributionConnection) publishExecution() error {
execution := connection.execution
connection.execution = nil
if execution == nil {
return nil
}
connection.lifecycleMu.Lock()
state := connection.transaction
if state == nil {
connection.lifecycleMu.Unlock()
return &SQLiteAttributionError{Cause: ErrSQLiteAttributionUnboundWrite}
}
if execution.hook.mismatch || !execution.hook.seen {
err := &SQLiteAttributionError{Cause: ErrSQLiteAttributionUnsupportedWrite}
state.poison = err
connection.lifecycleMu.Unlock()
return err
}
if state.journalFD < 0 {
descriptor, err := sqliteAttributionOpenJournalDescriptor(connection.journal)
if err != nil {
state.poison = err
connection.lifecycleMu.Unlock()
return err
}
state.journalFD = descriptor
}
journal, err := sqliteAttributionJournalIdentityFromDescriptor(state.journalFD)
if err != nil {
state.poison = err
connection.lifecycleMu.Unlock()
return err
}
transaction := state.transaction
tracker := state.tracker
connection.lifecycleMu.Unlock()
origin := execution.origin
origin.Operation = execution.classification.operation
origin.Table = execution.classification.table
err = tracker.RecordSQLiteWrite(transaction, SQLiteWriteObservation{
Origin: origin,
Update: SQLiteUpdateObservation{Operation: execution.classification.operation, Table: execution.classification.table, Journal: journal},
})
if err != nil {
return connection.poison(err)
}
return nil
}
func (connection *sqliteAttributionConnection) poison(err error) error {
connection.lifecycleMu.Lock()
state := connection.transaction
if state != nil {
state.poison = err
}
connection.lifecycleMu.Unlock()
return err
}
// The adapter owns the explicit main transaction and DELETE journal; path-only stat is racy, so retain this O_PATH descriptor through finalization.
func sqliteAttributionOpenJournalDescriptor(path string) (int, error) {
descriptor, err := unix.Open(path, unix.O_PATH|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0)
if err != nil {
return -1, &SQLiteAttributionError{Cause: errors.Join(ErrSQLiteAttributionJournalIdentity, err)}
}
return descriptor, nil
}
func sqliteAttributionCloseJournalDescriptor(descriptor int) error {
if descriptor < 0 {
return nil
}
if err := unix.Close(descriptor); err != nil {
return &SQLiteAttributionError{Cause: errors.Join(ErrSQLiteAttributionJournalIdentity, err)}
}
return nil
}
func sqliteAttributionJournalIdentityFromDescriptor(descriptor int) (SQLiteJournalIdentity, error) {
var status unix.Statx_t
mask := uint32(unix.STATX_BASIC_STATS | unix.STATX_BTIME | unix.STATX_MNT_ID)
if err := unix.Statx(descriptor, "", unix.AT_EMPTY_PATH|unix.AT_STATX_SYNC_AS_STAT, int(mask), &status); err != nil {
return SQLiteJournalIdentity{}, &SQLiteAttributionError{Cause: errors.Join(ErrSQLiteAttributionJournalIdentity, err)}
}
required := uint32(unix.STATX_MNT_ID | unix.STATX_BTIME)
if status.Mask&required != required {
return SQLiteJournalIdentity{}, &SQLiteAttributionError{Cause: ErrSQLiteAttributionJournalIdentity}
}
return SQLiteJournalIdentity{MountID: status.Mnt_id, DeviceMajor: status.Dev_major, DeviceMinor: status.Dev_minor, Inode: status.Ino, BirthSeconds: status.Btime.Sec, BirthNanoseconds: status.Btime.Nsec}, nil
}
func sqliteAttributionOrigin() SQLiteExecutionOrigin {
programCounters := make([]uintptr, 16)
count := runtime.Callers(3, programCounters)
programCounters = programCounters[:count]
frames := runtime.CallersFrames(programCounters)
frame, more := frames.Next()
for more && !strings.Contains(frame.Function, "github.com/nezhahq/nezha/") {
frame, more = frames.Next()
}
return SQLiteExecutionOrigin{StackHash: sqliteAttributionStackHash(programCounters), FirstNezhaFrame: frame.Function}
}
func sqliteAttributionStackHash(programCounters []uintptr) uint64 {
const offsetBasis uint64 = 14695981039346656037
const prime uint64 = 1099511628211
hash := offsetBasis
for _, programCounter := range programCounters {
hash ^= uint64(programCounter)
hash *= prime
}
return hash
}
@@ -0,0 +1,186 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"errors"
"path/filepath"
"testing"
"time"
"github.com/nezhahq/nezha/model"
"gorm.io/gorm"
)
func TestSQLiteAttributionHoldFacadeEnablesOnArmAndDisablesOnAbort(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
t.Cleanup(resetSQLiteAttributionForTest)
// When
receipt, armErr := ArmNextSQLiteHold()
enabledAfterArm := sqliteAttributionEnabled.Load()
_, abortErr := AbortSQLiteHold(receipt)
// Then
if armErr != nil || abortErr != nil {
t.Fatalf("arm=%v abort=%v", armErr, abortErr)
}
if !enabledAfterArm {
t.Fatal("successful hold arm did not enable SQLite attribution")
}
if sqliteAttributionEnabled.Load() {
t.Fatal("successful hold abort left SQLite attribution enabled")
}
}
func TestSQLiteAttributionHoldFacadeDisablesAfterRelease(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
t.Cleanup(resetSQLiteAttributionForTest)
receipt, err := ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
transaction := sqliteHoldTestTransaction(204)
tracker := sqliteAttributionTracker.Load()
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "api_tokens", Journal: sqliteHoldTestJournal}); err != nil {
t.Fatal(err)
}
if _, err := tracker.BeginSQLiteFinalization(transaction); err != nil {
t.Fatal(err)
}
// When
result, releaseErr := ReleaseSQLiteHold(receipt)
// Then
if releaseErr != nil || result.State != SQLiteHoldControlStateReleased {
t.Fatalf("release=%+v err=%v", result, releaseErr)
}
if sqliteAttributionEnabled.Load() {
t.Fatal("successful hold release left SQLite attribution enabled")
}
}
func TestSQLiteAttributionHoldFacadeKeepsEnabledAfterCanceledWaitUntilAbort(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
t.Cleanup(resetSQLiteAttributionForTest)
receipt, err := ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
// When
_, waitErr := WaitSQLiteHoldSelected(ctx, receipt)
enabledAfterWait := sqliteAttributionEnabled.Load()
_, abortErr := AbortSQLiteHold(receipt)
// Then
if !errors.Is(waitErr, context.Canceled) || abortErr != nil {
t.Fatalf("wait=%v abort=%v", waitErr, abortErr)
}
if !enabledAfterWait {
t.Fatal("canceled wait disabled attribution before the active hold was aborted")
}
if sqliteAttributionEnabled.Load() {
t.Fatal("abort after canceled wait left SQLite attribution enabled")
}
}
func TestSQLiteAttributionHoldFacadeStaleReceiptCannotDisableNewHold(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
t.Cleanup(resetSQLiteAttributionForTest)
first, err := ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
if _, err := AbortSQLiteHold(first); err != nil {
t.Fatal(err)
}
second, err := ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
// When
_, staleErr := AbortSQLiteHold(first)
// Then
if !errors.Is(staleErr, ErrSQLiteHoldStaleSession) {
t.Fatalf("stale abort error = %v", staleErr)
}
if !sqliteAttributionEnabled.Load() {
t.Fatal("stale receipt disabled attribution owned by the current hold")
}
if _, err := AbortSQLiteHold(second); err != nil {
t.Fatal(err)
}
}
func TestSQLiteAttributionHoldFacadeSelectsGORMAPITokenUsageUpdate(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
t.Cleanup(resetSQLiteAttributionForTest)
database, err := gorm.Open(openSQLiteDialector(filepath.Join(t.TempDir(), "dashboard.sqlite")), &gorm.Config{})
if err != nil {
t.Fatal(err)
}
sqlDatabase, err := database.DB()
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := sqlDatabase.Close(); closeErr != nil {
t.Error(closeErr)
}
})
if err := database.AutoMigrate(&model.APIToken{}); err != nil {
t.Fatal(err)
}
token := model.APIToken{UserID: 1, Name: "usage-update", TokenHash: "hash"}
if err := database.Create(&token).Error; err != nil {
t.Fatal(err)
}
receipt, err := ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
ctx, cancel := context.WithCancel(t.Context())
defer cancel()
writerDone := make(chan error, 1)
usageTime := time.Unix(1_700_000_000, 0)
go func() {
writerDone <- database.Model(&model.APIToken{}).Where("id = ?", token.ID).Updates(map[string]any{
"last_used_at": usageTime,
"last_used_ip": "127.0.0.1",
}).Error
cancel()
}()
// When
selected, selectedErr := WaitSQLiteHoldSelected(ctx, receipt)
finalizing, finalizingErr := WaitSQLiteHoldFinalizing(ctx, selected)
_, releaseErr := ReleaseSQLiteHold(finalizing)
writerErr := <-writerDone
// Then
if selectedErr != nil || finalizingErr != nil || releaseErr != nil || writerErr != nil {
t.Fatalf("selected=%v finalizing=%v release=%v writer=%v", selectedErr, finalizingErr, releaseErr, writerErr)
}
var updated model.APIToken
if err := database.First(&updated, token.ID).Error; err != nil {
t.Fatal(err)
}
if updated.LastUsedAt == nil || !updated.LastUsedAt.Equal(usageTime) || updated.LastUsedIP != "127.0.0.1" {
t.Fatalf("persisted usage update = %+v", updated)
}
}
@@ -0,0 +1,42 @@
//go:build agentcompat && linux
package singleton
import (
"path/filepath"
"testing"
"golang.org/x/sys/unix"
)
func TestSQLiteAttributionDerivesJournalIdentityFromOpenedDescriptor(t *testing.T) {
// Given
journalPath := filepath.Join(t.TempDir(), "dashboard.sqlite-journal")
descriptor, err := unix.Open(journalPath, unix.O_CREAT|unix.O_WRONLY|unix.O_CLOEXEC, 0o600)
if err != nil {
t.Fatal(err)
}
if err := unix.Close(descriptor); err != nil {
t.Fatal(err)
}
descriptor, err = unix.Open(journalPath, unix.O_PATH|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := unix.Close(descriptor); closeErr != nil {
t.Error(closeErr)
}
})
// When
identity, err := sqliteAttributionJournalIdentityFromDescriptor(descriptor)
// Then
if err != nil {
t.Fatal(err)
}
if identity.MountID == 0 || identity.Inode == 0 || identity.BirthSeconds == 0 {
t.Fatal("descriptor-derived journal identity is incomplete")
}
}
@@ -0,0 +1,211 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"database/sql/driver"
"errors"
"testing"
"golang.org/x/sys/unix"
)
func TestSQLiteAttributionReturningCompletesExactlyOneRow(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
enableSQLiteAttribution()
transaction, err := database.BeginTx(context.Background(), nil)
if err != nil {
t.Fatal(err)
}
// When
rows, err := transaction.QueryContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id", "one-row")
if err != nil {
t.Fatal(err)
}
if !rows.Next() {
t.Fatal("RETURNING did not yield its inserted row")
}
var identifier int64
if err := rows.Scan(&identifier); err != nil {
t.Fatal(err)
}
if rows.Next() {
t.Fatal("RETURNING yielded more than one row")
}
rowsErr := rows.Err()
closeErr := rows.Close()
commitErr := transaction.Commit()
count := sqliteAttributionSettingsCount(t, database)
// Then
if rowsErr != nil || closeErr != nil || commitErr != nil {
t.Fatal("completed RETURNING did not finish successfully")
}
if count != 1 {
t.Fatalf("completed RETURNING persisted %d rows", count)
}
}
func TestSQLiteAttributionEarlyReturningCloseRollsBackTransaction(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
enableSQLiteAttribution()
transaction, err := database.BeginTx(context.Background(), nil)
if err != nil {
t.Fatal(err)
}
// When
rows, err := transaction.QueryContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id", "early-close")
if err != nil {
t.Fatal(err)
}
if !rows.Next() {
t.Fatal("RETURNING did not yield its inserted row")
}
if err := rows.Close(); err != nil {
t.Fatal(err)
}
commitErr := transaction.Commit()
count := sqliteAttributionSettingsCount(t, database)
// Then
if commitErr == nil {
t.Fatal("early RETURNING Close allowed Commit")
}
if count != 0 {
t.Fatalf("early RETURNING Close persisted %d rows", count)
}
}
func TestSQLiteAttributionZeroRowUpdateDoesNotPublishEvidence(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
enableSQLiteAttribution()
transaction, err := database.BeginTx(context.Background(), nil)
if err != nil {
t.Fatal(err)
}
// When
if _, err := transaction.Exec("UPDATE settings SET value = ? WHERE id = ?", "zero", -1); err != nil {
t.Fatal(err)
}
commitErr := transaction.Commit()
evidence := sqliteAttributionTrackerWriteEvidence()
// Then
if commitErr != nil {
t.Fatal(commitErr)
}
if evidence.hasWrite {
t.Fatal("zero-row UPDATE published write evidence")
}
}
func TestSQLiteAttributionDirectReadQueryClosesOwnedStatement(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
// When
rows, err := database.QueryContext(context.Background(), "SELECT id FROM settings")
if err != nil {
t.Fatal(err)
}
next := rows.Next()
rowsErr := rows.Err()
closeErr := rows.Close()
// Then
if next {
t.Fatal("empty SELECT returned a row")
}
if rowsErr != nil || closeErr != nil {
t.Fatal("direct readonly Query did not complete cleanly")
}
}
func TestSQLiteAttributionCommitClosesRetainedJournalDescriptor(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
rawConnection, err := sqliteAttributionDriver{}.Open(sqliteAttributionTestDatabasePath(t))
if err != nil {
t.Fatal(err)
}
connection := rawConnection.(*sqliteAttributionConnection)
t.Cleanup(func() {
if closeErr := connection.Close(); closeErr != nil {
t.Error(closeErr)
}
})
create, err := connection.Prepare("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)")
if err != nil {
t.Fatal(err)
}
if _, err := create.Exec(nil); err != nil {
t.Fatal(err)
}
if err := create.Close(); err != nil {
t.Fatal(err)
}
enableSQLiteAttribution()
transaction, err := connection.Begin()
if err != nil {
t.Fatal(err)
}
insert, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)")
if err != nil {
t.Fatal(err)
}
if _, err := insert.Exec([]driver.Value{"retained-descriptor"}); err != nil {
t.Fatal(err)
}
if err := insert.Close(); err != nil {
t.Fatal(err)
}
_, descriptor, active := sqliteAttributionTransactionState(t, connection)
if !active {
t.Fatal("active transaction is missing before Commit")
}
// When
commitErr := transaction.Commit()
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
// Then
if commitErr != nil {
t.Fatal(commitErr)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Fatalf("retained journal descriptor remains open: %v", descriptorErr)
}
}
func TestSQLiteAttributionQueryRowReturningPoisonsTransaction(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
enableSQLiteAttribution()
transaction, err := database.BeginTx(context.Background(), nil)
if err != nil {
t.Fatal(err)
}
// When
var identifier int64
scanErr := transaction.QueryRowContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id", "query-row").Scan(&identifier)
commitErr := transaction.Commit()
// Then
if scanErr != nil {
t.Fatal(scanErr)
}
if commitErr == nil {
t.Fatal("QueryRowContext RETURNING committed without reaching EOF")
}
if sqliteAttributionTrackerHasWrite() {
t.Fatal("QueryRowContext RETURNING published evidence without EOF")
}
}
@@ -0,0 +1,145 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"database/sql/driver"
"errors"
"testing"
)
type sqliteAttributionLegacyQueryStmt struct {
closeCount int
}
func (statement *sqliteAttributionLegacyQueryStmt) Close() error { statement.closeCount++; return nil }
func (sqliteAttributionLegacyQueryStmt) NumInput() int { return -1 }
func (sqliteAttributionLegacyQueryStmt) Exec([]driver.Value) (driver.Result, error) {
return nil, driver.ErrSkip
}
func (sqliteAttributionLegacyQueryStmt) Query([]driver.Value) (driver.Rows, error) {
return nil, driver.ErrSkip
}
type sqliteAttributionCountingStmt struct {
driver.Stmt
closeCount int
}
func (statement *sqliteAttributionCountingStmt) Close() error {
statement.closeCount++
return statement.Stmt.Close()
}
type sqliteAttributionQueryProbeStmt struct {
closeCount int
queryCount int
}
func (statement *sqliteAttributionQueryProbeStmt) Close() error { statement.closeCount++; return nil }
func (sqliteAttributionQueryProbeStmt) NumInput() int { return -1 }
func (sqliteAttributionQueryProbeStmt) Exec([]driver.Value) (driver.Result, error) {
return nil, driver.ErrSkip
}
func (statement *sqliteAttributionQueryProbeStmt) Query([]driver.Value) (driver.Rows, error) {
statement.queryCount++
return nil, driver.ErrSkip
}
func TestSQLiteAttributionDirectQueryClosesOwnedStatementWhenPreStepRejects(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
rawConnection, err := sqliteAttributionDriver{}.Open(sqliteAttributionTestDatabasePath(t))
if err != nil {
t.Fatal(err)
}
connection := rawConnection.(*sqliteAttributionConnection)
t.Cleanup(func() {
if closeErr := connection.Close(); closeErr != nil {
t.Error(closeErr)
}
})
if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil {
t.Fatal(err)
}
prepared, err := connection.Prepare("INSERT INTO settings (value) VALUES (?) RETURNING id")
if err != nil {
t.Fatal(err)
}
statement := prepared.(*sqliteAttributionStatement)
owner := &sqliteAttributionCountingStmt{Stmt: prepared}
enableSQLiteAttribution()
// When
rows, queryErr := statement.queryOwned(func() (driver.Rows, error) {
return statement.statement.Query([]driver.Value{"rejected"})
}, owner)
// Then
if rows != nil {
t.Fatal("rejected direct Query returned rows")
}
if !errors.Is(queryErr, ErrSQLiteAttributionUnboundWrite) {
t.Fatalf("rejected direct Query error = %v", queryErr)
}
if owner.closeCount != 1 {
t.Fatalf("owned statement Close count = %d, want 1", owner.closeCount)
}
}
func TestSQLiteAttributionDirectQueryRejectsUnboundWriteThroughPublicConnection(t *testing.T) {
resetSQLiteAttributionForTest()
probe := &sqliteAttributionQueryProbeStmt{}
var connection *sqliteAttributionConnection
connection = &sqliteAttributionConnection{
prepareStatement: func(context.Context, string) (driver.Stmt, error) {
return &sqliteAttributionStatement{
connection: connection,
statement: probe,
classification: sqliteAttributionClassification{
hasRowDML: true,
operation: SQLiteOperationInsert,
table: "settings",
},
}, nil
},
}
enableSQLiteAttribution()
rows, queryErr := connection.Query("INSERT INTO settings (value) VALUES (?) RETURNING id", []driver.Value{"public-rejected"})
if rows != nil {
t.Fatal("public direct Query returned rows after pre-step rejection")
}
if !errors.Is(queryErr, ErrSQLiteAttributionUnboundWrite) {
t.Fatalf("public direct Query error = %v", queryErr)
}
if probe.closeCount != 1 {
t.Fatalf("public direct Query statement Close count = %d, want 1", probe.closeCount)
}
if probe.queryCount != 0 {
t.Fatalf("public direct Query invoked underlying Query %d times, want 0", probe.queryCount)
}
}
func TestSQLiteAttributionQueryContextClosesUnsupportedStatementBeforeErrSkip(t *testing.T) {
// Given
statement := &sqliteAttributionLegacyQueryStmt{}
connection := &sqliteAttributionConnection{}
wrapped := &sqliteAttributionStatement{connection: connection, statement: statement}
// When
rows, queryErr := connection.queryContextStatement(context.Background(), wrapped, nil, statement)
// Then
if rows != nil {
t.Fatal("unsupported QueryContext returned rows")
}
if !errors.Is(queryErr, driver.ErrSkip) {
t.Fatalf("unsupported QueryContext error = %v", queryErr)
}
if statement.closeCount != 1 {
t.Fatalf("unsupported QueryContext Close count = %d, want 1", statement.closeCount)
}
}
@@ -0,0 +1,227 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"database/sql"
"database/sql/driver"
"errors"
"testing"
"golang.org/x/sys/unix"
)
func sqliteAttributionReleasedCommitFixture(t *testing.T, value string) (*sqliteAttributionConnection, *sqliteAttributionTx, SQLiteHoldSession, string, SQLiteTransaction, int) {
t.Helper()
resetSQLiteAttributionForTest()
databasePath := sqliteAttributionTestDatabasePath(t)
rawConnection, err := sqliteAttributionDriver{}.Open(databasePath)
if err != nil {
t.Fatal(err)
}
connection := rawConnection.(*sqliteAttributionConnection)
if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil {
t.Fatal(err)
}
enableSQLiteAttribution()
transaction, err := connection.BeginTx(context.Background(), driver.TxOptions{})
if err != nil {
t.Fatal(err)
}
statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)")
if err != nil {
t.Fatal(err)
}
if _, err := statement.Exec([]driver.Value{value}); err != nil {
t.Fatal(err)
}
if err := statement.Close(); err != nil {
t.Fatal(err)
}
session, err := sqliteAttributionTracker.Load().ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
wrapped := transaction.(*sqliteAttributionTx)
identity, descriptor, active := sqliteAttributionTransactionState(t, connection)
if !active {
t.Fatal("selected transaction is not active")
}
return connection, wrapped, session, databasePath, identity, descriptor
}
func sqliteAttributionReleasedTerminalSnapshot(t *testing.T, session SQLiteHoldSession) SQLiteHoldSnapshot {
t.Helper()
terminal, err := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing)
if err != nil {
t.Fatal(err)
}
return terminal
}
func TestSQLiteAttributionReleasedCommitOwnsTerminalAgainstConcurrentRollback(t *testing.T) {
// Given
connection, transaction, session, databasePath, identity, descriptor := sqliteAttributionReleasedCommitFixture(t, "released-rollback-owner")
t.Cleanup(func() {
if err := connection.Close(); err != nil && !errors.Is(err, driver.ErrBadConn) {
t.Error(err)
}
})
releasedBoundary := make(chan struct{})
allowCommitClaim := make(chan struct{})
loserWaiting := make(chan struct{})
transaction.state.releasedCommitBoundary = func() {
close(releasedBoundary)
<-allowCommitClaim
}
transaction.state.terminalLoserWaitBoundary = func() { close(loserWaiting) }
finalizing := make(chan error, 1)
commitResult := make(chan error, 1)
go func() {
_, err := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing)
finalizing <- err
}()
go func() { commitResult <- transaction.Commit() }()
if err := <-finalizing; err != nil {
t.Fatal(err)
}
// When
if err := sqliteAttributionTracker.Load().ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
<-releasedBoundary
rollbackResult := make(chan error, 1)
go func() { rollbackResult <- transaction.Rollback() }()
loserReturned := false
var rollbackErr error
select {
case <-loserWaiting:
case rollbackErr = <-rollbackResult:
loserReturned = true
}
close(allowCommitClaim)
commitErr := <-commitResult
if !loserReturned {
rollbackErr = <-rollbackResult
}
terminal := sqliteAttributionReleasedTerminalSnapshot(t, session)
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
count := sqliteAttributionPersistedCount(t, databasePath)
// Then
if loserReturned {
t.Errorf("released Commit lost terminal ownership: Rollback returned %v before waiting", rollbackErr)
}
if commitErr != nil {
t.Errorf("released Commit error = %v, want nil", commitErr)
}
if !errors.Is(rollbackErr, driver.ErrBadConn) {
t.Errorf("Rollback after successful Release error = %v, want driver.ErrBadConn", rollbackErr)
}
if !terminal.Selected || !terminal.Finalizing || !terminal.Released {
t.Errorf("released terminal snapshot = %+v", terminal)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Errorf("journal descriptor after released Commit = %v, want EBADF", descriptorErr)
}
if count != 1 {
t.Errorf("released Commit persisted %d rows, want 1", count)
}
if _, _, active := sqliteAttributionTransactionState(t, connection); active {
t.Error("released Commit left the connection transaction active")
}
if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) {
t.Error("released Commit left the tracker transaction active")
}
}
func TestSQLiteAttributionReleasedCommitOwnsTerminalAgainstConcurrentConnectionClose(t *testing.T) {
// Given
connection, transaction, session, databasePath, identity, descriptor := sqliteAttributionReleasedCommitFixture(t, "released-close-owner")
connectionClosed := false
t.Cleanup(func() {
if !connectionClosed {
if err := connection.Close(); err != nil && !errors.Is(err, driver.ErrBadConn) {
t.Error(err)
}
}
})
releasedBoundary := make(chan struct{})
allowCommitClaim := make(chan struct{})
loserWaiting := make(chan struct{})
transaction.state.releasedCommitBoundary = func() {
close(releasedBoundary)
<-allowCommitClaim
}
transaction.state.terminalLoserWaitBoundary = func() { close(loserWaiting) }
finalizing := make(chan error, 1)
commitResult := make(chan error, 1)
go func() {
_, err := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing)
finalizing <- err
}()
go func() { commitResult <- transaction.Commit() }()
if err := <-finalizing; err != nil {
t.Fatal(err)
}
// When
if err := sqliteAttributionTracker.Load().ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
<-releasedBoundary
closeResult := make(chan error, 1)
go func() { closeResult <- connection.Close() }()
loserReturned := false
var closeErr error
select {
case <-loserWaiting:
case closeErr = <-closeResult:
loserReturned = true
}
close(allowCommitClaim)
commitErr := <-commitResult
if !loserReturned {
closeErr = <-closeResult
}
connectionClosed = true
terminal := sqliteAttributionReleasedTerminalSnapshot(t, session)
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
database, err := sql.Open("sqlite3", databasePath)
if err != nil {
t.Fatal(err)
}
defer database.Close()
var count int
if err := database.QueryRow("SELECT COUNT(*) FROM settings").Scan(&count); err != nil {
t.Fatal(err)
}
// Then
if loserReturned {
t.Errorf("released Commit lost terminal ownership: connection Close returned %v before waiting", closeErr)
}
if commitErr != nil {
t.Errorf("released Commit error = %v, want nil", commitErr)
}
if closeErr != nil && !errors.Is(closeErr, driver.ErrBadConn) {
t.Errorf("connection Close after successful Release error = %v, want nil or driver.ErrBadConn", closeErr)
}
if !terminal.Selected || !terminal.Finalizing || !terminal.Released {
t.Errorf("released terminal snapshot = %+v", terminal)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Errorf("journal descriptor after released Commit = %v, want EBADF", descriptorErr)
}
if count != 1 {
t.Errorf("released Commit persisted %d rows, want 1", count)
}
if _, _, active := sqliteAttributionTransactionState(t, connection); active {
t.Error("released Commit left the connection transaction active")
}
if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) {
t.Error("released Commit left the tracker transaction active")
}
}
@@ -0,0 +1,121 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"testing"
"github.com/nezhahq/nezha/model"
"gorm.io/gorm"
)
func TestSQLiteAttributionRecordsDirectExplicitReturningEvidence(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
enableSQLiteAttribution()
transaction, err := database.BeginTx(context.Background(), nil)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if rollbackErr := transaction.Rollback(); rollbackErr != nil {
t.Error(rollbackErr)
}
})
// When
rows, queryErr := transaction.QueryContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id", "direct-explicit-returning")
rowsErr := sqliteAttributionConsumeRows(rows)
evidence := sqliteAttributionTrackerWriteEvidence()
// Then
if queryErr != nil || rowsErr != nil {
t.Fatalf("direct explicit RETURNING failed")
}
if !evidence.hasWrite {
t.Fatal("direct explicit RETURNING did not record atomic write evidence")
}
if evidence.write.Origin.StackHash == 0 {
t.Fatal("direct explicit RETURNING recorded a zero stack hash")
}
if evidence.write.Origin.FirstNezhaFrame == "" {
t.Fatal("direct explicit RETURNING recorded an empty Nezha frame")
}
}
func TestSQLiteAttributionRecordsPreparedExplicitReturningEvidence(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
enableSQLiteAttribution()
transaction, err := database.BeginTx(context.Background(), nil)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if rollbackErr := transaction.Rollback(); rollbackErr != nil {
t.Error(rollbackErr)
}
})
statement, err := transaction.PrepareContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := statement.Close(); closeErr != nil {
t.Error(closeErr)
}
})
// When
rows, queryErr := statement.QueryContext(context.Background(), "prepared-explicit-returning")
rowsErr := sqliteAttributionConsumeRows(rows)
evidence := sqliteAttributionTrackerWriteEvidence()
// Then
if queryErr != nil || rowsErr != nil {
t.Fatalf("prepared explicit RETURNING failed")
}
if !evidence.hasWrite {
t.Fatal("prepared explicit RETURNING did not record atomic write evidence")
}
if evidence.write.Origin.StackHash == 0 {
t.Fatal("prepared explicit RETURNING recorded a zero stack hash")
}
if evidence.write.Origin.FirstNezhaFrame == "" {
t.Fatal("prepared explicit RETURNING recorded an empty Nezha frame")
}
}
func TestSQLiteAttributionAllowsGORMReturningWhenDisabled(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
database, err := gorm.Open(openSQLiteDialector(sqliteAttributionTestDatabasePath(t)), &gorm.Config{})
if err != nil {
t.Fatal(err)
}
sqlDatabase, err := database.DB()
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := sqlDatabase.Close(); closeErr != nil {
t.Error(closeErr)
}
})
if err := database.AutoMigrate(&model.User{}); err != nil {
t.Fatal(err)
}
user := model.User{Username: "admin", Password: "hashed-password"}
// When
createErr := database.Create(&user).Error
// Then
if createErr != nil {
t.Fatalf("disabled attribution rejected GORM RETURNING: %v", createErr)
}
if user.ID == 0 {
t.Fatal("disabled attribution did not return the generated user ID")
}
}
@@ -0,0 +1,231 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"database/sql"
"database/sql/driver"
"errors"
"reflect"
"sync"
"testing"
)
func TestSQLiteAttributionLifecycleFieldsUseOneConnectionMutex(t *testing.T) {
// Given
mutexType := reflect.TypeFor[sync.Mutex]()
connectionType := reflect.TypeFor[sqliteAttributionConnection]()
transactionType := reflect.TypeFor[sqliteAttributionTransaction]()
mutexCount := 0
for _, typeUnderTest := range []reflect.Type{connectionType, transactionType} {
for fieldIndex := 0; fieldIndex < typeUnderTest.NumField(); fieldIndex++ {
if typeUnderTest.Field(fieldIndex).Type == mutexType {
mutexCount++
}
}
}
// When
field, found := connectionType.FieldByName("lifecycleMu")
// Then
if mutexCount != 1 {
t.Fatalf("direct lifecycle mutex count = %d, want 1 on sqliteAttributionConnection", mutexCount)
}
if !found || field.Type != mutexType {
t.Fatal("sqliteAttributionConnection.lifecycleMu is not the lifecycle mutex")
}
}
func TestSQLiteAttributionBeginRollsBackRawTransactionWhenTrackerRejects(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
rawConnection, err := sqliteAttributionDriver{}.Open(sqliteAttributionTestDatabasePath(t))
if err != nil {
t.Fatal(err)
}
connection := rawConnection.(*sqliteAttributionConnection)
t.Cleanup(func() {
_, _ = connection.connection.Exec("ROLLBACK", nil)
if closeErr := connection.Close(); closeErr != nil {
t.Error(closeErr)
}
})
identity := SQLiteTransaction{Connection: connection.identity, Identity: SQLiteTransactionIdentity(sqliteAttributionTransactionID.Load() + 1)}
if err := sqliteAttributionTracker.Load().BeginSQLiteTransaction(identity); err != nil {
t.Fatal(err)
}
// When
_, beginErr := connection.Begin()
followup, followupErr := connection.Begin()
if followup != nil {
_ = followup.Rollback()
}
// Then
if !errors.Is(beginErr, ErrSQLiteHoldTransactionActive) {
t.Fatalf("tracker rejection = %v, want ErrSQLiteHoldTransactionActive", beginErr)
}
if followupErr != nil {
t.Fatalf("tracker rejection left raw transaction active: %v", followupErr)
}
}
func TestSQLiteAttributionDirectQueryPreservesSQLiteColumnTypes(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
databasePath := sqliteAttributionTestDatabasePath(t)
attributed, err := openSQLiteAttributionTestDB(databasePath)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := attributed.Close(); closeErr != nil {
t.Error(closeErr)
}
})
standard, err := sql.Open("sqlite3", databasePath)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := standard.Close(); closeErr != nil {
t.Error(closeErr)
}
})
if _, err := attributed.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)"); err != nil {
t.Fatal(err)
}
if _, err := attributed.Exec("INSERT INTO settings (value) VALUES (?)", "typed"); err != nil {
t.Fatal(err)
}
// When
attributedRows, err := attributed.QueryContext(context.Background(), "SELECT id, value FROM settings")
if err != nil {
t.Fatal(err)
}
attributedTypes, attributedErr := attributedRows.ColumnTypes()
attributedCloseErr := attributedRows.Close()
standardRows, err := standard.QueryContext(context.Background(), "SELECT id, value FROM settings")
if err != nil {
t.Fatal(err)
}
standardTypes, standardErr := standardRows.ColumnTypes()
standardCloseErr := standardRows.Close()
// Then
if attributedErr != nil || attributedCloseErr != nil || standardErr != nil || standardCloseErr != nil {
t.Fatal("ColumnTypes did not complete cleanly")
}
if len(attributedTypes) != len(standardTypes) {
t.Fatalf("attributed ColumnTypes length = %d, standard = %d", len(attributedTypes), len(standardTypes))
}
for index := range standardTypes {
if attributedTypes[index].DatabaseTypeName() != standardTypes[index].DatabaseTypeName() || attributedTypes[index].ScanType() != standardTypes[index].ScanType() || !reflect.DeepEqual(attributedTypes[index], standardTypes[index]) {
t.Fatalf("column %d metadata differs from go-sqlite3", index)
}
}
}
func TestSQLiteAttributionFailsClosedWhenSQLiteRepreparesAfterSchemaChange(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
databasePath := sqliteAttributionTestDatabasePath(t)
rawConnection, err := sqliteAttributionDriver{}.Open(databasePath)
if err != nil {
t.Fatal(err)
}
connection := rawConnection.(*sqliteAttributionConnection)
t.Cleanup(func() {
if closeErr := connection.Close(); closeErr != nil {
t.Error(closeErr)
}
})
if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil {
t.Fatal(err)
}
statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := statement.Close(); closeErr != nil {
t.Error(closeErr)
}
})
other, err := sql.Open("sqlite3", databasePath)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := other.Close(); closeErr != nil {
t.Error(closeErr)
}
})
if _, err := other.Exec("CREATE TABLE audit (value TEXT)"); err != nil {
t.Fatal(err)
}
if _, err := other.Exec("CREATE TRIGGER settings_audit AFTER INSERT ON settings BEGIN INSERT INTO audit (value) VALUES (NEW.value); END"); err != nil {
t.Fatal(err)
}
enableSQLiteAttribution()
transaction, err := connection.Begin()
if err != nil {
t.Fatal(err)
}
// When
_, executionErr := statement.Exec([]driver.Value{"reprepared"})
commitErr := transaction.Commit()
var settingsCount, auditCount int
if err := other.QueryRow("SELECT COUNT(*) FROM settings").Scan(&settingsCount); err != nil {
t.Fatal(err)
}
if err := other.QueryRow("SELECT COUNT(*) FROM audit").Scan(&auditCount); err != nil {
t.Fatal(err)
}
// Then
if !errors.Is(errors.Join(executionErr, commitErr), ErrSQLiteAttributionUnsupportedWrite) {
t.Fatalf("reprepared write errors = %v / %v", executionErr, commitErr)
}
if settingsCount != 0 || auditCount != 0 {
t.Fatalf("reprepared write persisted settings=%d audit=%d", settingsCount, auditCount)
}
}
func TestSQLiteAttributionUpdateHookRejectsAuxiliaryDatabase(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
rawConnection, err := sqliteAttributionDriver{}.Open(sqliteAttributionTestDatabasePath(t))
if err != nil {
t.Fatal(err)
}
connection := rawConnection.(*sqliteAttributionConnection)
t.Cleanup(func() {
if closeErr := connection.Close(); closeErr != nil {
t.Error(closeErr)
}
})
connection.execution = &sqliteAttributionExecution{classification: sqliteAttributionClassification{operation: SQLiteOperationInsert, table: "settings"}}
if _, err := connection.connection.Exec("ATTACH DATABASE ':memory:' AS auxiliary", nil); err != nil {
t.Fatal(err)
}
if _, err := connection.connection.Exec("CREATE TABLE auxiliary.settings (value TEXT)", nil); err != nil {
t.Fatal(err)
}
// When
_, writeErr := connection.connection.Exec("INSERT INTO auxiliary.settings (value) VALUES ('auxiliary')", nil)
// Then
if writeErr != nil {
t.Fatal(writeErr)
}
if !connection.execution.hook.mismatch || connection.execution.hook.seen {
t.Fatal("auxiliary database update hook matched main attribution execution")
}
}
@@ -0,0 +1,186 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"database/sql/driver"
"errors"
"io"
"reflect"
)
type sqliteAttributionClassification struct {
readonly bool
hasRowDML bool
ambiguous bool
operation SQLiteOperation
table string
}
func (classification sqliteAttributionClassification) requiresAttribution() bool {
return !classification.readonly && classification.hasRowDML
}
func (classification sqliteAttributionClassification) valid() bool {
return !classification.ambiguous && classification.hasRowDML && classification.operation != "" && classification.table != ""
}
type sqliteAttributionStatement struct {
connection *sqliteAttributionConnection
statement driver.Stmt
classification sqliteAttributionClassification
}
var (
_ driver.Stmt = (*sqliteAttributionStatement)(nil)
_ driver.StmtExecContext = (*sqliteAttributionStatement)(nil)
_ driver.StmtQueryContext = (*sqliteAttributionStatement)(nil)
)
func (statement *sqliteAttributionStatement) Close() error { return statement.statement.Close() }
func (statement *sqliteAttributionStatement) NumInput() int { return statement.statement.NumInput() }
func (statement *sqliteAttributionStatement) Exec(values []driver.Value) (driver.Result, error) {
return statement.execute(func() (driver.Result, error) { return statement.statement.Exec(values) })
}
func (statement *sqliteAttributionStatement) ExecContext(ctx context.Context, values []driver.NamedValue) (driver.Result, error) {
contextStatement, ok := statement.statement.(driver.StmtExecContext)
if !ok {
return nil, driver.ErrSkip
}
return statement.execute(func() (driver.Result, error) { return contextStatement.ExecContext(ctx, values) })
}
func (statement *sqliteAttributionStatement) Query(values []driver.Value) (driver.Rows, error) {
return statement.queryOwned(func() (driver.Rows, error) { return statement.statement.Query(values) }, nil)
}
func (statement *sqliteAttributionStatement) QueryContext(ctx context.Context, values []driver.NamedValue) (driver.Rows, error) {
contextStatement, ok := statement.statement.(driver.StmtQueryContext)
if !ok {
return nil, driver.ErrSkip
}
return statement.queryOwned(func() (driver.Rows, error) { return contextStatement.QueryContext(ctx, values) }, nil)
}
func (statement *sqliteAttributionStatement) execute(run func() (driver.Result, error)) (driver.Result, error) {
if err := statement.connection.beforeWrite(statement.classification); err != nil {
return nil, err
}
result, err := run()
if err != nil {
statement.connection.discardExecution()
return nil, err
}
if rows, rowsErr := result.RowsAffected(); rowsErr != nil || rows > 0 {
if err := statement.connection.publishExecution(); err != nil {
return nil, err
}
}
statement.connection.discardExecution()
return result, nil
}
func (statement *sqliteAttributionStatement) queryOwned(run func() (driver.Rows, error), owner driver.Stmt) (driver.Rows, error) {
if err := statement.connection.beforeWrite(statement.classification); err != nil {
if owner != nil {
return nil, errors.Join(err, owner.Close())
}
return nil, err
}
rows, err := run()
if err != nil {
statement.connection.discardExecution()
if owner != nil {
return nil, errors.Join(err, owner.Close())
}
return nil, err
}
// Disabled attribution must preserve the stock driver rows lifecycle used during Dashboard bootstrap.
if !sqliteAttributionEnabled.Load() || !statement.classification.requiresAttribution() {
if owner == nil {
return rows, nil
}
return &sqliteAttributionRows{rows: rows, owner: owner}, nil
}
return &sqliteAttributionRows{rows: rows, owner: owner, connection: statement.connection}, nil
}
type sqliteAttributionRows struct {
rows driver.Rows
owner driver.Stmt
connection *sqliteAttributionConnection
finished bool
closed bool
}
func (rows *sqliteAttributionRows) Columns() []string { return rows.rows.Columns() }
func (rows *sqliteAttributionRows) ColumnTypeDatabaseTypeName(index int) string {
typed, ok := rows.rows.(driver.RowsColumnTypeDatabaseTypeName)
if !ok {
return ""
}
return typed.ColumnTypeDatabaseTypeName(index)
}
func (rows *sqliteAttributionRows) ColumnTypeNullable(index int) (bool, bool) {
typed, ok := rows.rows.(driver.RowsColumnTypeNullable)
if !ok {
return false, false
}
return typed.ColumnTypeNullable(index)
}
func (rows *sqliteAttributionRows) ColumnTypeScanType(index int) reflect.Type {
typed, ok := rows.rows.(driver.RowsColumnTypeScanType)
if !ok {
return nil
}
return typed.ColumnTypeScanType(index)
}
func (rows *sqliteAttributionRows) Close() error {
if rows.closed {
return nil
}
rows.closed = true
closeErr := rows.rows.Close()
if !rows.finished && rows.connection != nil {
rows.connection.poison(&SQLiteAttributionError{Cause: ErrSQLiteAttributionUnsupportedWrite})
rows.connection.discardExecution()
}
if rows.owner != nil {
return errors.Join(closeErr, rows.owner.Close())
}
return closeErr
}
func (rows *sqliteAttributionRows) Next(destination []driver.Value) error {
err := rows.rows.Next(destination)
if errors.Is(err, driver.ErrBadConn) {
// Read-only direct Query rows have no attribution connection; ErrBadConn must still propagate.
if rows.connection != nil {
rows.connection.discardExecution()
}
return err
}
if err == nil {
return nil
}
rows.finished = true
if err != io.EOF {
if rows.connection != nil {
rows.connection.discardExecution()
}
return err
}
if rows.connection != nil {
if publishErr := rows.connection.publishExecution(); publishErr != nil {
return publishErr
}
}
return io.EOF
}
@@ -0,0 +1,126 @@
//go:build agentcompat && linux
package singleton
import (
"database/sql/driver"
"errors"
)
type SQLiteHoldError struct{ Cause error }
func (err *SQLiteHoldError) Error() string { return "sqlite hold finalization failed" }
func (err *SQLiteHoldError) Unwrap() error { return err.Cause }
type sqliteAttributionTx struct {
connection *sqliteAttributionConnection
state *sqliteAttributionTransaction
}
var _ driver.Tx = (*sqliteAttributionTx)(nil)
func (transaction *sqliteAttributionTx) Commit() error {
state := transaction.state
transaction.connection.lifecycleMu.Lock()
if state.terminalPhase != sqliteAttributionTerminalOpen {
transaction.connection.lifecycleMu.Unlock()
return state.waitForTerminal()
}
poison := state.poison
if poison != nil {
journalFD := state.claimTerminalLocked(sqliteAttributionTerminalRollbackOwned)
transaction.connection.lifecycleMu.Unlock()
if errors.Is(poison, ErrSQLiteHoldAmbiguousCandidate) {
return errors.Join(&SQLiteHoldError{Cause: poison}, state.finish(transaction.connection, journalFD, state.raw.Rollback))
}
return errors.Join(poison, state.finish(transaction.connection, journalFD, state.raw.Rollback))
}
state.terminalPhase = sqliteAttributionTerminalCommitWaiting
transaction.connection.lifecycleMu.Unlock()
finalization, err := state.tracker.BeginSQLiteCommitFinalization(state.transaction)
if errors.Is(err, ErrSQLiteHoldNotSelected) {
return state.commit(transaction.connection)
}
if err != nil {
return errors.Join(&SQLiteHoldError{Cause: err}, state.rollback(transaction.connection))
}
// Wait before raw Commit so the rollback journal stays observable for deterministic drain and prevents the baseline sample-1 FD regression.
if err := state.tracker.WaitSQLiteCommitFinalization(state.context, finalization); err != nil {
return errors.Join(&SQLiteHoldError{Cause: errors.Join(ErrSQLiteHoldAborted, err)}, state.rollback(transaction.connection))
}
if state.releasedCommitBoundary != nil {
state.releasedCommitBoundary()
}
return state.commit(transaction.connection)
}
func (transaction *sqliteAttributionTx) Rollback() error {
return transaction.state.rollback(transaction.connection)
}
func (state *sqliteAttributionTransaction) commit(connection *sqliteAttributionConnection) error {
connection.lifecycleMu.Lock()
if state.terminalPhase != sqliteAttributionTerminalCommitWaiting {
connection.lifecycleMu.Unlock()
return state.waitForTerminal()
}
journalFD := state.claimTerminalLocked(sqliteAttributionTerminalCommitOwned)
connection.lifecycleMu.Unlock()
return state.finish(connection, journalFD, state.raw.Commit)
}
func (state *sqliteAttributionTransaction) rollback(connection *sqliteAttributionConnection) error {
connection.lifecycleMu.Lock()
switch state.terminalPhase {
case sqliteAttributionTerminalOpen:
journalFD := state.claimTerminalLocked(sqliteAttributionTerminalRollbackOwned)
connection.lifecycleMu.Unlock()
return state.finish(connection, journalFD, state.raw.Rollback)
case sqliteAttributionTerminalCommitWaiting:
connection.lifecycleMu.Unlock()
arbitration := state.tracker.ArbitrateSQLiteCommitWaitingTerminal(state.transaction)
if arbitration == sqliteAttributionTerminalReleaseWon || arbitration == sqliteAttributionTerminalRollbackReserved {
return state.waitForTerminal()
}
connection.lifecycleMu.Lock()
if state.terminalPhase != sqliteAttributionTerminalCommitWaiting {
connection.lifecycleMu.Unlock()
return state.waitForTerminal()
}
journalFD := state.claimTerminalLocked(sqliteAttributionTerminalRollbackOwned)
connection.lifecycleMu.Unlock()
return state.finish(connection, journalFD, state.raw.Rollback)
default:
connection.lifecycleMu.Unlock()
return state.waitForTerminal()
}
}
func (state *sqliteAttributionTransaction) claimTerminalLocked(phase sqliteAttributionTerminalPhase) int {
state.terminalPhase = phase
journalFD := state.journalFD
state.journalFD = -1
return journalFD
}
func (state *sqliteAttributionTransaction) waitForTerminal() error {
if state.terminalLoserWaitBoundary != nil {
state.terminalLoserWaitBoundary()
}
<-state.done
return driver.ErrBadConn
}
func (state *sqliteAttributionTransaction) finish(connection *sqliteAttributionConnection, journalFD int, completeRaw func() error) error {
defer close(state.done)
rawErr := completeRaw()
connection.discardExecution()
connection.lifecycleMu.Lock()
if connection.transaction == state {
connection.transaction = nil
}
connection.lifecycleMu.Unlock()
closeErr := sqliteAttributionCloseJournalDescriptor(journalFD)
trackerErr := state.tracker.FinishSQLiteTransaction(state.transaction)
return errors.Join(rawErr, closeErr, trackerErr)
}
@@ -0,0 +1,236 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"database/sql"
"database/sql/driver"
"errors"
"testing"
"golang.org/x/sys/unix"
)
func sqliteAttributionHeldTransaction(t *testing.T, value string, transactionContext context.Context) (*sqliteAttributionConnection, driver.Tx, SQLiteHoldSession, string) {
t.Helper()
resetSQLiteAttributionForTest()
databasePath := sqliteAttributionTestDatabasePath(t)
rawConnection, err := sqliteAttributionDriver{}.Open(databasePath)
if err != nil {
t.Fatal(err)
}
connection := rawConnection.(*sqliteAttributionConnection)
t.Cleanup(func() {
if err := connection.Close(); err != nil {
t.Error(err)
}
})
if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil {
t.Fatal(err)
}
enableSQLiteAttribution()
transaction, err := connection.BeginTx(transactionContext, driver.TxOptions{})
if err != nil {
t.Fatal(err)
}
statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)")
if err != nil {
t.Fatal(err)
}
if _, err := statement.Exec([]driver.Value{value}); err != nil {
t.Fatal(err)
}
if err := statement.Close(); err != nil {
t.Fatal(err)
}
session, err := sqliteAttributionTracker.Load().ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
return connection, transaction, session, databasePath
}
func sqliteAttributionStartHeldCommit(t *testing.T, transaction driver.Tx, session SQLiteHoldSession) <-chan error {
t.Helper()
finalizing := make(chan error, 1)
commit := make(chan error, 1)
go func() {
_, err := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing)
finalizing <- err
}()
go func() { commit <- transaction.Commit() }()
select {
case err := <-finalizing:
if err != nil {
t.Fatalf("Commit finalization wait error = %v", err)
}
case err := <-commit:
t.Fatalf("Commit completed before selected hold release: %v", err)
}
return commit
}
func sqliteAttributionPersistedCount(t *testing.T, databasePath string) int {
t.Helper()
database, err := sql.Open("sqlite3", databasePath)
if err != nil {
t.Fatal(err)
}
defer database.Close()
var count int
if err := database.QueryRow("SELECT COUNT(*) FROM settings").Scan(&count); err != nil {
t.Fatal(err)
}
return count
}
func sqliteAttributionTransactionState(t *testing.T, connection *sqliteAttributionConnection) (SQLiteTransaction, int, bool) {
t.Helper()
connection.lifecycleMu.Lock()
state := connection.transaction
if state == nil {
connection.lifecycleMu.Unlock()
return SQLiteTransaction{}, -1, false
}
identity, descriptor := state.transaction, state.journalFD
connection.lifecycleMu.Unlock()
return identity, descriptor, true
}
func sqliteAttributionTrackerTransactionActive(tracker *SQLiteHoldTracker, identity SQLiteTransaction) bool {
tracker.mu.Lock()
defer tracker.mu.Unlock()
_, active := tracker.transactions[identity]
return active
}
func TestSQLiteAttributionCommitWaitsForSelectedHoldRelease(t *testing.T) {
// Given
connection, transaction, session, databasePath := sqliteAttributionHeldTransaction(t, "held-commit", context.Background())
identity, descriptor, active := sqliteAttributionTransactionState(t, connection)
if !active {
t.Fatal("held transaction is not active")
}
// When
commit := sqliteAttributionStartHeldCommit(t, transaction, session)
select {
case err := <-commit:
t.Fatalf("Commit completed while selected hold remains unreleased: %v", err)
default:
}
if err := sqliteAttributionTracker.Load().ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
commitErr := <-commit
// Then
if commitErr != nil {
t.Fatal(commitErr)
}
if _, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0); !errors.Is(descriptorErr, unix.EBADF) {
t.Fatalf("journal descriptor after released Commit = %v, want EBADF", descriptorErr)
}
if _, _, active := sqliteAttributionTransactionState(t, connection); active {
t.Fatal("released Commit left the connection transaction active")
}
if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) {
t.Fatal("released Commit left the tracker transaction active")
}
if count := sqliteAttributionPersistedCount(t, databasePath); count != 1 {
t.Fatalf("released held Commit persisted %d rows", count)
}
}
func TestSQLiteAttributionCommitRollsBackWhenSelectedHoldAborts(t *testing.T) {
// Given
connection, transaction, session, databasePath := sqliteAttributionHeldTransaction(t, "aborted-commit", context.Background())
identity, descriptor, active := sqliteAttributionTransactionState(t, connection)
if !active {
t.Fatal("held transaction is not active")
}
// When
commit := sqliteAttributionStartHeldCommit(t, transaction, session)
if err := sqliteAttributionTracker.Load().AbortSQLiteHold(session); err != nil {
t.Fatal(err)
}
commitErr := <-commit
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
// Then
var holdErr *SQLiteHoldError
if !errors.As(commitErr, &holdErr) || !errors.Is(commitErr, ErrSQLiteHoldAborted) {
t.Fatalf("aborted held Commit error = %v", commitErr)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Fatalf("journal descriptor after aborted Commit = %v, want EBADF", descriptorErr)
}
if _, _, active := sqliteAttributionTransactionState(t, connection); active {
t.Fatal("aborted Commit left the connection transaction active")
}
if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) {
t.Fatal("aborted Commit left the tracker transaction active")
}
if count := sqliteAttributionPersistedCount(t, databasePath); count != 0 {
t.Fatalf("aborted held Commit persisted %d rows", count)
}
}
func TestSQLiteAttributionCommitRollsBackWhenBeginTxContextCancels(t *testing.T) {
// Given
transactionContext, cancel := context.WithCancel(context.Background())
defer cancel()
connection, transaction, session, databasePath := sqliteAttributionHeldTransaction(t, "cancelled-commit", transactionContext)
identity, descriptor, active := sqliteAttributionTransactionState(t, connection)
if !active {
t.Fatal("cancelled transaction is not active")
}
// When
commit := sqliteAttributionStartHeldCommit(t, transaction, session)
cancel()
commitErr := <-commit
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
// Then
var holdErr *SQLiteHoldError
if !errors.As(commitErr, &holdErr) || !errors.Is(commitErr, ErrSQLiteHoldAborted) || !errors.Is(commitErr, context.Canceled) {
t.Fatalf("cancelled held Commit error = %v", commitErr)
}
if count := sqliteAttributionPersistedCount(t, databasePath); count != 0 {
t.Fatalf("cancelled held Commit persisted %d rows", count)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Fatalf("journal descriptor after cancelled Commit = %v, want EBADF", descriptorErr)
}
if _, _, active := sqliteAttributionTransactionState(t, connection); active {
t.Fatal("cancelled Commit left the connection transaction active")
}
if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) {
t.Fatal("cancelled Commit left the tracker transaction active")
}
}
func TestSQLiteAttributionRollbackWakesSelectedCommit(t *testing.T) {
// Given
_, transaction, session, databasePath := sqliteAttributionHeldTransaction(t, "rollback-wakes-commit", context.Background())
// When
commit := sqliteAttributionStartHeldCommit(t, transaction, session)
rollbackErr := transaction.Rollback()
commitErr := <-commit
// Then
if rollbackErr != nil {
t.Fatal(rollbackErr)
}
var holdErr *SQLiteHoldError
if !errors.As(commitErr, &holdErr) || !errors.Is(commitErr, ErrSQLiteHoldAborted) {
t.Fatalf("Commit error after Rollback = %v", commitErr)
}
if count := sqliteAttributionPersistedCount(t, databasePath); count != 0 {
t.Fatalf("Rollback while Commit waited persisted %d rows", count)
}
}
@@ -0,0 +1,158 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"database/sql/driver"
"errors"
"testing"
"golang.org/x/sys/unix"
)
func sqliteAttributionTransactionWithWrite(t *testing.T, databasePath, value string) (*sqliteAttributionConnection, driver.Tx) {
t.Helper()
rawConnection, err := sqliteAttributionDriver{}.Open(databasePath)
if err != nil {
t.Fatal(err)
}
connection := rawConnection.(*sqliteAttributionConnection)
t.Cleanup(func() {
if err := connection.Close(); err != nil {
t.Error(err)
}
})
if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil {
t.Fatal(err)
}
transaction, err := connection.BeginTx(context.Background(), driver.TxOptions{})
if err != nil {
t.Fatal(err)
}
statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)")
if err != nil {
t.Fatal(err)
}
if _, err := statement.Exec([]driver.Value{value}); err != nil {
t.Fatal(err)
}
if err := statement.Close(); err != nil {
t.Fatal(err)
}
return connection, transaction
}
func sqliteAttributionTransactionWithAmbiguousWrite(t *testing.T, databasePath, value string) (*sqliteAttributionConnection, driver.Tx, error) {
t.Helper()
rawConnection, err := sqliteAttributionDriver{}.Open(databasePath)
if err != nil {
t.Fatal(err)
}
connection := rawConnection.(*sqliteAttributionConnection)
t.Cleanup(func() {
if err := connection.Close(); err != nil {
t.Error(err)
}
})
if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil {
t.Fatal(err)
}
transaction, err := connection.BeginTx(context.Background(), driver.TxOptions{})
if err != nil {
t.Fatal(err)
}
statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)")
if err != nil {
t.Fatal(err)
}
_, writeErr := statement.Exec([]driver.Value{value})
if err := statement.Close(); err != nil {
t.Fatal(err)
}
return connection, transaction, writeErr
}
func TestSQLiteAttributionUnselectedCommitPersistsImmediately(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
enableSQLiteAttribution()
databasePath := sqliteAttributionTestDatabasePath(t)
session, err := sqliteAttributionTracker.Load().ArmSQLiteHold(SQLiteJournalIdentity{})
if err != nil {
t.Fatal(err)
}
connection, transaction := sqliteAttributionTransactionWithWrite(t, databasePath, "unselected-commit")
identity, descriptor, active := sqliteAttributionTransactionState(t, connection)
if !active {
t.Fatal("unselected transaction is not active")
}
// When
commitErr := transaction.Commit()
abortErr := sqliteAttributionTracker.Load().AbortSQLiteHold(session)
_, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0)
// Then
if commitErr != nil {
t.Fatal(commitErr)
}
if abortErr != nil {
t.Fatal(abortErr)
}
if count := sqliteAttributionPersistedCount(t, databasePath); count != 1 {
t.Fatalf("unselected Commit persisted %d rows", count)
}
if !errors.Is(descriptorErr, unix.EBADF) {
t.Fatalf("journal descriptor after unselected Commit = %v, want EBADF", descriptorErr)
}
if _, _, active := sqliteAttributionTransactionState(t, connection); active {
t.Fatal("unselected Commit left the connection transaction active")
}
if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) {
t.Fatal("unselected Commit left the tracker transaction active")
}
}
func TestSQLiteAttributionFutureAmbiguityRollsBackEveryParticipatingCommit(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
enableSQLiteAttribution()
firstPath := sqliteAttributionTestDatabasePath(t)
secondPath := sqliteAttributionTestDatabasePath(t)
tracker := sqliteAttributionTracker.Load()
if _, err := tracker.ArmNextSQLiteHold(); err != nil {
t.Fatal(err)
}
firstConnection, firstTransaction := sqliteAttributionTransactionWithWrite(t, firstPath, "first-ambiguous")
secondConnection, secondTransaction, secondWriteErr := sqliteAttributionTransactionWithAmbiguousWrite(t, secondPath, "second-ambiguous")
if !errors.Is(secondWriteErr, ErrSQLiteHoldAmbiguousCandidate) {
t.Fatalf("second ambiguity write error = %v", secondWriteErr)
}
firstIdentity, _, firstActive := sqliteAttributionTransactionState(t, firstConnection)
secondIdentity, _, secondActive := sqliteAttributionTransactionState(t, secondConnection)
if !firstActive || !secondActive {
t.Fatal("ambiguity participants are not active")
}
// When
firstCommitErr := firstTransaction.Commit()
secondCommitErr := secondTransaction.Commit()
// Then
for _, commitErr := range []error{firstCommitErr, secondCommitErr} {
var holdErr *SQLiteHoldError
if !errors.As(commitErr, &holdErr) || !errors.Is(commitErr, ErrSQLiteHoldAmbiguousCandidate) {
t.Fatalf("ambiguous Commit error = %v", commitErr)
}
}
if count := sqliteAttributionPersistedCount(t, firstPath); count != 0 {
t.Fatalf("first ambiguous Commit persisted %d rows", count)
}
if count := sqliteAttributionPersistedCount(t, secondPath); count != 0 {
t.Fatalf("second ambiguous Commit persisted %d rows", count)
}
if sqliteAttributionTrackerTransactionActive(tracker, firstIdentity) || sqliteAttributionTrackerTransactionActive(tracker, secondIdentity) {
t.Fatal("ambiguous Commit left a tracker transaction active")
}
}
@@ -0,0 +1,217 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"database/sql/driver"
"errors"
"testing"
)
func TestSQLiteAttributionRejectsDirectUnboundExecBeforeInsert(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
enableSQLiteAttribution()
// When
_, err := database.Exec("INSERT INTO settings (value) VALUES (?)", "direct-unbound")
count := sqliteAttributionSettingsCount(t, database)
// Then
// An UpdateHook error after SQLite executes is too late for autocommit: the write may already be committed.
if !errors.Is(err, ErrSQLiteAttributionUnboundWrite) {
t.Error("direct unbound Exec error does not wrap ErrSQLiteAttributionUnboundWrite")
}
if count != 0 {
t.Errorf("direct unbound Exec persisted %d rows", count)
}
}
func TestSQLiteAttributionRejectsLegacyDriverQueryBeforeInsert(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
connection, err := sqliteAttributionDriver{}.Open(sqliteAttributionTestDatabasePath(t))
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := connection.Close(); closeErr != nil {
t.Error(closeErr)
}
})
if _, err := connection.(driver.Execer).Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil {
t.Fatal(err)
}
enableSQLiteAttribution()
// When
rows, queryErr := connection.(driver.Queryer).Query("INSERT INTO settings (value) VALUES (?) RETURNING id", []driver.Value{"legacy-direct"})
rowsErr := sqliteAttributionConsumeDriverRows(rows)
// Then
if !errors.Is(errors.Join(queryErr, rowsErr), ErrSQLiteAttributionUnboundWrite) {
t.Error("legacy driver Query does not reject unbound row DML")
}
}
func TestSQLiteAttributionRejectsPreparedLegacyDriverQueryBeforeInsert(t *testing.T) {
// Given
resetSQLiteAttributionForTest()
connection, err := sqliteAttributionDriver{}.Open(sqliteAttributionTestDatabasePath(t))
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := connection.Close(); closeErr != nil {
t.Error(closeErr)
}
})
if _, err := connection.(driver.Execer).Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil {
t.Fatal(err)
}
enableSQLiteAttribution()
statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?) RETURNING id")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := statement.Close(); closeErr != nil {
t.Error(closeErr)
}
})
// When
rows, queryErr := statement.Query([]driver.Value{"legacy-prepared"})
rowsErr := sqliteAttributionConsumeDriverRows(rows)
// Then
if !errors.Is(errors.Join(queryErr, rowsErr), ErrSQLiteAttributionUnboundWrite) {
t.Error("prepared legacy driver Query does not reject unbound row DML")
}
}
func TestSQLiteAttributionRejectsPreparedUnboundExecBeforeInsert(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
enableSQLiteAttribution()
statement, err := database.Prepare("INSERT INTO settings (value) VALUES (?)")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := statement.Close(); closeErr != nil {
t.Error(closeErr)
}
})
// When
_, err = statement.Exec("prepared-unbound")
count := sqliteAttributionSettingsCount(t, database)
// Then
if !errors.Is(err, ErrSQLiteAttributionUnboundWrite) {
t.Error("prepared unbound Exec error does not wrap ErrSQLiteAttributionUnboundWrite")
}
if count != 0 {
t.Errorf("prepared unbound Exec persisted %d rows", count)
}
}
func TestSQLiteAttributionRejectsDirectUnboundReturningBeforeInsert(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
enableSQLiteAttribution()
// When
rows, queryErr := database.QueryContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id", "direct-returning")
rowsErr := sqliteAttributionConsumeRows(rows)
count := sqliteAttributionSettingsCount(t, database)
// Then
if !errors.Is(errors.Join(queryErr, rowsErr), ErrSQLiteAttributionUnboundWrite) {
t.Error("direct unbound RETURNING does not propagate ErrSQLiteAttributionUnboundWrite through rows")
}
if count != 0 {
t.Errorf("direct unbound RETURNING persisted %d rows", count)
}
}
func TestSQLiteAttributionRejectsPreparedUnboundReturningBeforeInsert(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
enableSQLiteAttribution()
statement, err := database.PrepareContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := statement.Close(); closeErr != nil {
t.Error(closeErr)
}
})
// When
rows, queryErr := statement.QueryContext(context.Background(), "prepared-returning")
rowsErr := sqliteAttributionConsumeRows(rows)
count := sqliteAttributionSettingsCount(t, database)
// Then
if !errors.Is(errors.Join(queryErr, rowsErr), ErrSQLiteAttributionUnboundWrite) {
t.Error("prepared unbound RETURNING does not propagate ErrSQLiteAttributionUnboundWrite through rows")
}
if count != 0 {
t.Errorf("prepared unbound RETURNING persisted %d rows", count)
}
}
func TestSQLiteAttributionRejectsDirectUnboundUpdateBeforeSideEffect(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
if _, err := database.Exec("INSERT INTO settings (value) VALUES (?)", "original"); err != nil {
t.Fatal(err)
}
enableSQLiteAttribution()
// When
_, err := database.Exec("UPDATE settings SET value = ?", "changed")
value := sqliteAttributionSettingValue(t, database)
// Then
if !errors.Is(err, ErrSQLiteAttributionUnboundWrite) {
t.Error("direct unbound UPDATE error does not wrap ErrSQLiteAttributionUnboundWrite")
}
if value != "original" {
t.Error("direct unbound UPDATE changed the persisted value")
}
}
func TestSQLiteAttributionRejectsPreparedUnboundDeleteBeforeSideEffect(t *testing.T) {
// Given
database := openSQLiteAttributionTestDatabase(t)
if _, err := database.Exec("INSERT INTO settings (value) VALUES (?)", "preserved"); err != nil {
t.Fatal(err)
}
enableSQLiteAttribution()
statement, err := database.Prepare("DELETE FROM settings")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := statement.Close(); closeErr != nil {
t.Error(closeErr)
}
})
// When
_, err = statement.Exec()
count := sqliteAttributionSettingsCount(t, database)
// Then
if !errors.Is(err, ErrSQLiteAttributionUnboundWrite) {
t.Error("prepared unbound DELETE error does not wrap ErrSQLiteAttributionUnboundWrite")
}
if count != 1 {
t.Errorf("prepared unbound DELETE left %d rows", count)
}
}
@@ -0,0 +1,13 @@
//go:build agentcompat && linux
package singleton
import (
"gorm.io/driver/sqlite"
"gorm.io/gorm"
)
func openSQLiteDialector(path string) gorm.Dialector {
registerSQLiteAttributionDriver()
return sqlite.New(sqlite.Config{DriverName: sqliteAttributionDriverName, DSN: path})
}
@@ -0,0 +1,12 @@
//go:build !agentcompat || !linux
package singleton
import (
"gorm.io/driver/sqlite"
"gorm.io/gorm"
)
func openSQLiteDialector(path string) gorm.Dialector {
return sqlite.Open(path)
}
@@ -0,0 +1,250 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"database/sql"
"database/sql/driver"
"errors"
"fmt"
"strings"
"sync"
"sync/atomic"
"github.com/mattn/go-sqlite3"
)
const sqliteAttributionDriverName = "nezha_sqlite_attribution"
var (
ErrSQLiteAttributionUnsupportedDSN = errors.New("sqlite attribution requires a filesystem database")
ErrSQLiteAttributionUnboundWrite = errors.New("sqlite attribution observed a write without an explicit transaction")
ErrSQLiteAttributionJournalIdentity = errors.New("sqlite attribution could not identify the rollback journal")
ErrSQLiteAttributionUnsupportedWrite = errors.New("sqlite attribution cannot classify this write")
sqliteAttributionDriverRegistration sync.Once
sqliteAttributionEnabled atomic.Bool
sqliteAttributionConnectionID atomic.Uint64
sqliteAttributionTransactionID atomic.Uint64
sqliteAttributionTracker atomic.Pointer[SQLiteHoldTracker]
sqliteAttributionHoldControl *sqliteHoldControl
)
type SQLiteAttributionError struct{ Cause error }
func (err *SQLiteAttributionError) Error() string { return err.Cause.Error() }
func (err *SQLiteAttributionError) Unwrap() error { return err.Cause }
func init() { resetSQLiteAttributionControl() }
func registerSQLiteAttributionDriver() {
sqliteAttributionDriverRegistration.Do(func() { sql.Register(sqliteAttributionDriverName, sqliteAttributionDriver{}) })
}
func enableSQLiteAttribution() { sqliteAttributionEnabled.Store(true) }
func resetSQLiteAttributionForTest() {
sqliteAttributionEnabled.Store(false)
resetSQLiteAttributionControl()
}
func resetSQLiteAttributionControl() {
tracker := newSQLiteAttributionHoldTracker(&sqliteAttributionEnabled)
sqliteAttributionTracker.Store(tracker)
sqliteAttributionHoldControl = newProductionSQLiteHoldControl(tracker)
}
type sqliteAttributionDriver struct{}
func (sqliteAttributionDriver) Open(dataSourceName string) (driver.Conn, error) {
rawConnection, err := (&sqlite3.SQLiteDriver{}).Open(dataSourceName)
if err != nil {
return nil, err
}
connection, ok := rawConnection.(*sqlite3.SQLiteConn)
if !ok {
return nil, errors.New("sqlite attribution received an unsupported sqlite connection")
}
databasePath := connection.GetFilename("")
if databasePath == "" {
if closeErr := connection.Close(); closeErr != nil {
return nil, fmt.Errorf("close unsupported sqlite database: %w", closeErr)
}
return nil, &SQLiteAttributionError{Cause: ErrSQLiteAttributionUnsupportedDSN}
}
wrapped := &sqliteAttributionConnection{connection: connection, identity: SQLiteConnectionIdentity(sqliteAttributionConnectionID.Add(1)), journal: databasePath + "-journal", closeRawConnection: connection.Close}
wrapped.prepareStatement = wrapped.prepareSQLiteStatement
connection.RegisterAuthorizer(wrapped.authorize)
connection.RegisterUpdateHook(wrapped.recordUpdate)
return wrapped, nil
}
type sqliteAttributionConnection struct {
connection *sqlite3.SQLiteConn
identity SQLiteConnectionIdentity
journal string
// lifecycleMu owns the active transaction and its poison, journal descriptor, and completion state; never hold it across raw SQLite, tracker, wait, or close operations.
lifecycleMu sync.Mutex
transaction *sqliteAttributionTransaction
capturing bool
preparing sqliteAttributionClassification
execution *sqliteAttributionExecution
prepareStatement sqliteAttributionStatementPreparer
closeRawConnection func() error
}
type sqliteAttributionStatementPreparer func(context.Context, string) (driver.Stmt, error)
type sqliteAttributionTransaction struct {
transaction SQLiteTransaction
raw driver.Tx
tracker *SQLiteHoldTracker
context context.Context
journalFD int
poison error
terminalPhase sqliteAttributionTerminalPhase
done chan struct{}
releasedCommitBoundary func()
terminalLoserWaitBoundary func()
}
type sqliteAttributionTerminalPhase uint8
const (
sqliteAttributionTerminalOpen sqliteAttributionTerminalPhase = iota
sqliteAttributionTerminalCommitWaiting
sqliteAttributionTerminalCommitOwned
sqliteAttributionTerminalRollbackOwned
)
type sqliteAttributionExecution struct {
classification sqliteAttributionClassification
origin SQLiteExecutionOrigin
hook sqliteAttributionHook
}
type sqliteAttributionHook struct {
seen bool
mismatch bool
}
var (
_ driver.Driver = sqliteAttributionDriver{}
_ driver.Conn = (*sqliteAttributionConnection)(nil)
_ driver.Pinger = (*sqliteAttributionConnection)(nil)
_ driver.ConnPrepareContext = (*sqliteAttributionConnection)(nil)
_ driver.ConnBeginTx = (*sqliteAttributionConnection)(nil)
_ driver.Execer = (*sqliteAttributionConnection)(nil)
_ driver.ExecerContext = (*sqliteAttributionConnection)(nil)
_ driver.Queryer = (*sqliteAttributionConnection)(nil)
_ driver.QueryerContext = (*sqliteAttributionConnection)(nil)
)
func (connection *sqliteAttributionConnection) authorize(operation int, table, _, schema string) int {
if !connection.capturing {
return sqlite3.SQLITE_OK
}
classification := &connection.preparing
if operation != sqlite3.SQLITE_INSERT && operation != sqlite3.SQLITE_UPDATE && operation != sqlite3.SQLITE_DELETE {
return sqlite3.SQLITE_OK
}
rowOperation, ok := sqliteAttributionOperation(operation)
if !ok || schema != "main" || strings.HasPrefix(table, "sqlite_") {
classification.ambiguous = true
return sqlite3.SQLITE_OK
}
// SQLite authorizes UPDATE once per written column; only a different row-DML target is ambiguous.
if classification.hasRowDML {
if classification.operation != rowOperation || classification.table != table {
classification.ambiguous = true
}
return sqlite3.SQLITE_OK
}
classification.hasRowDML = true
classification.operation = rowOperation
classification.table = table
return sqlite3.SQLITE_OK
}
func (connection *sqliteAttributionConnection) Prepare(query string) (driver.Stmt, error) {
return connection.prepareStatement(context.Background(), query)
}
func (connection *sqliteAttributionConnection) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
return connection.prepareStatement(ctx, query)
}
func (connection *sqliteAttributionConnection) prepareSQLiteStatement(ctx context.Context, query string) (driver.Stmt, error) {
connection.preparing = sqliteAttributionClassification{}
connection.capturing = true
statement, err := connection.connection.PrepareContext(ctx, query)
connection.capturing = false
classification := connection.preparing
connection.preparing = sqliteAttributionClassification{}
if err != nil {
return nil, err
}
rawStatement, ok := statement.(*sqlite3.SQLiteStmt)
if !ok {
return nil, errors.New("sqlite attribution received an unsupported sqlite statement")
}
classification.readonly = rawStatement.Readonly()
return &sqliteAttributionStatement{connection: connection, statement: statement, classification: classification}, nil
}
func (connection *sqliteAttributionConnection) Ping(ctx context.Context) error {
return connection.connection.Ping(ctx)
}
func (connection *sqliteAttributionConnection) Begin() (driver.Tx, error) {
return connection.begin(context.Background(), driver.TxOptions{})
}
func (connection *sqliteAttributionConnection) BeginTx(ctx context.Context, options driver.TxOptions) (driver.Tx, error) {
return connection.begin(ctx, options)
}
func (connection *sqliteAttributionConnection) begin(ctx context.Context, options driver.TxOptions) (driver.Tx, error) {
transaction, err := connection.connection.BeginTx(ctx, options)
if err != nil {
return nil, err
}
identity := SQLiteTransaction{Connection: connection.identity, Identity: SQLiteTransactionIdentity(sqliteAttributionTransactionID.Add(1))}
tracker := sqliteAttributionTracker.Load()
if err := tracker.BeginSQLiteTransaction(identity); err != nil {
// Tracker rejection happens after BEGIN; leave no raw transaction to wedge this connection.
return nil, errors.Join(err, transaction.Rollback())
}
state := &sqliteAttributionTransaction{transaction: identity, raw: transaction, tracker: tracker, context: ctx, journalFD: -1, done: make(chan struct{})}
connection.lifecycleMu.Lock()
connection.transaction = state
connection.lifecycleMu.Unlock()
return &sqliteAttributionTx{connection: connection, state: state}, nil
}
func (connection *sqliteAttributionConnection) recordUpdate(operation int, database, table string, _ int64) {
execution := connection.execution
if execution == nil {
return
}
rowOperation, ok := sqliteAttributionOperation(operation)
if !ok || database != "main" || rowOperation != execution.classification.operation || table != execution.classification.table {
execution.hook.mismatch = true
return
}
execution.hook.seen = true
}
func sqliteAttributionOperation(operation int) (SQLiteOperation, bool) {
switch operation {
case sqlite3.SQLITE_INSERT:
return SQLiteOperationInsert, true
case sqlite3.SQLITE_UPDATE:
return SQLiteOperationUpdate, true
case sqlite3.SQLITE_DELETE:
return SQLiteOperationDelete, true
default:
return "", false
}
}
@@ -0,0 +1,204 @@
//go:build agentcompat && linux
package singleton
import (
"database/sql"
"database/sql/driver"
"errors"
"path/filepath"
"testing"
)
func TestSQLiteDriverAdapterRecordsExplicitInsert_when_AttributionEnabled(t *testing.T) {
resetSQLiteAttributionForTest()
databasePath := filepath.Join(t.TempDir(), "dashboard.sqlite")
database, err := openSQLiteAttributionTestDB(databasePath)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := database.Close(); closeErr != nil {
t.Error(closeErr)
}
})
if _, err := database.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)"); err != nil {
t.Fatal(err)
}
enableSQLiteAttribution()
transaction, err := database.Begin()
if err != nil {
t.Fatal(err)
}
if _, err := transaction.Exec("INSERT INTO settings (value) VALUES (?)", "opaque"); err != nil {
t.Fatal(err)
}
if !sqliteAttributionTrackerHasWrite() {
t.Fatal("explicit insert did not record an atomic sqlite write")
}
if err := transaction.Commit(); err != nil {
t.Fatal(err)
}
}
func TestSQLiteDriverAdapterRecordsPreparedExplicitInsert_when_AttributionEnabled(t *testing.T) {
resetSQLiteAttributionForTest()
databasePath := filepath.Join(t.TempDir(), "dashboard.sqlite")
database, err := openSQLiteAttributionTestDB(databasePath)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := database.Close(); closeErr != nil {
t.Error(closeErr)
}
})
if _, err := database.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)"); err != nil {
t.Fatal(err)
}
enableSQLiteAttribution()
transaction, err := database.Begin()
if err != nil {
t.Fatal(err)
}
statement, err := transaction.Prepare("INSERT INTO settings (value) VALUES (?)")
if err != nil {
t.Fatal(err)
}
if _, err := statement.Exec("prepared"); err != nil {
t.Fatal(err)
}
if err := statement.Close(); err != nil {
t.Fatal(err)
}
if !sqliteAttributionTrackerHasWrite() {
t.Fatal("prepared explicit insert did not record an atomic sqlite write")
}
if err := transaction.Commit(); err != nil {
t.Fatal(err)
}
}
func TestSQLiteDriverAdapterRejectsUnboundWrite_when_AttributionEnabled(t *testing.T) {
database := openSQLiteAttributionTestDatabase(t)
enableSQLiteAttribution()
_, err := database.Exec("INSERT INTO settings (value) VALUES (?)", "unbound")
if !errors.Is(err, ErrSQLiteAttributionUnboundWrite) {
t.Fatalf("expected unbound write instrumentation error, got %v", err)
}
}
func TestSQLiteDriverAdapterRejectsUnsupportedDSN(t *testing.T) {
database, err := openSQLiteAttributionTestDB(":memory:")
if err == nil {
err = database.Ping()
}
if database != nil {
t.Cleanup(func() {
if closeErr := database.Close(); closeErr != nil {
t.Error(closeErr)
}
})
}
if !errors.Is(err, ErrSQLiteAttributionUnsupportedDSN) {
t.Fatalf("expected unsupported DSN error, got %v", err)
}
}
func openSQLiteAttributionTestDB(path string) (*sql.DB, error) {
registerSQLiteAttributionDriver()
return sql.Open(sqliteAttributionDriverName, path)
}
func openSQLiteAttributionTestDatabase(t *testing.T) *sql.DB {
t.Helper()
resetSQLiteAttributionForTest()
database, err := openSQLiteAttributionTestDB(sqliteAttributionTestDatabasePath(t))
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if closeErr := database.Close(); closeErr != nil {
t.Error(closeErr)
}
})
if _, err := database.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)"); err != nil {
t.Fatal(err)
}
return database
}
func sqliteAttributionTestDatabasePath(t *testing.T) string {
t.Helper()
return filepath.Join(t.TempDir(), "dashboard.sqlite")
}
func sqliteAttributionSettingsCount(t *testing.T, database *sql.DB) int {
t.Helper()
var count int
if err := database.QueryRow("SELECT COUNT(*) FROM settings").Scan(&count); err != nil {
t.Fatal(err)
}
return count
}
func sqliteAttributionSettingValue(t *testing.T, database *sql.DB) string {
t.Helper()
var value string
if err := database.QueryRow("SELECT value FROM settings LIMIT 1").Scan(&value); err != nil {
t.Fatal(err)
}
return value
}
func sqliteAttributionConsumeRows(rows *sql.Rows) error {
if rows == nil {
return nil
}
var identifier int64
for rows.Next() {
if err := rows.Scan(&identifier); err != nil {
closeErr := rows.Close()
return errors.Join(err, closeErr)
}
}
err := rows.Err()
closeErr := rows.Close()
return errors.Join(err, closeErr)
}
func sqliteAttributionConsumeDriverRows(rows driver.Rows) error {
if rows == nil {
return nil
}
values := make([]driver.Value, len(rows.Columns()))
for {
err := rows.Next(values)
if err != nil {
closeErr := rows.Close()
return errors.Join(err, closeErr)
}
}
}
func sqliteAttributionTrackerHasWrite() bool {
return sqliteAttributionTrackerWriteEvidence().hasWrite
}
type sqliteAttributionWriteEvidence struct {
hasWrite bool
write SQLiteWriteObservation
}
func sqliteAttributionTrackerWriteEvidence() sqliteAttributionWriteEvidence {
tracker := sqliteAttributionTracker.Load()
tracker.mu.Lock()
defer tracker.mu.Unlock()
for _, held := range tracker.transactions {
if held.hasWrite {
return sqliteAttributionWriteEvidence{hasWrite: true, write: held.write}
}
}
return sqliteAttributionWriteEvidence{}
}
@@ -0,0 +1,176 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"crypto/rand"
"encoding/base64"
"errors"
"io"
"sync"
)
var ErrSQLiteHoldUnexpectedSelection = errors.New("sqlite hold selected an unexpected write")
type SQLiteHoldControlState string
const (
SQLiteHoldControlStateArmed SQLiteHoldControlState = "armed"
SQLiteHoldControlStateSelected SQLiteHoldControlState = "selected"
SQLiteHoldControlStateFinalizing SQLiteHoldControlState = "finalizing"
SQLiteHoldControlStateReleased SQLiteHoldControlState = "released"
SQLiteHoldControlStateAborted SQLiteHoldControlState = "aborted"
)
type SQLiteHoldReceipt struct {
ID string `json:"id"`
State SQLiteHoldControlState `json:"state"`
}
type sqliteHoldControlRecord struct {
receipt SQLiteHoldReceipt
session SQLiteHoldSession
}
type sqliteHoldControl struct {
mu sync.Mutex
tracker *SQLiteHoldTracker
random io.Reader
record *sqliteHoldControlRecord
}
func newSQLiteHoldControl(tracker *SQLiteHoldTracker, randomSource io.Reader) *sqliteHoldControl {
return &sqliteHoldControl{tracker: tracker, random: randomSource}
}
func newProductionSQLiteHoldControl(tracker *SQLiteHoldTracker) *sqliteHoldControl {
return newSQLiteHoldControl(tracker, rand.Reader)
}
func (control *sqliteHoldControl) ArmNextSQLiteHold() (SQLiteHoldReceipt, error) {
control.mu.Lock()
defer control.mu.Unlock()
identifier, err := control.newReceiptID()
if err != nil {
return SQLiteHoldReceipt{}, err
}
session, err := control.tracker.ArmNextSQLiteHold()
if err != nil {
return SQLiteHoldReceipt{}, err
}
receipt := SQLiteHoldReceipt{ID: identifier, State: SQLiteHoldControlStateArmed}
control.record = &sqliteHoldControlRecord{receipt: receipt, session: session}
return receipt, nil
}
func (control *sqliteHoldControl) WaitSelected(ctx context.Context, receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) {
return control.wait(ctx, receipt, SQLiteHoldWaitSelected, SQLiteHoldControlStateSelected)
}
func (control *sqliteHoldControl) WaitFinalizing(ctx context.Context, receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) {
return control.wait(ctx, receipt, SQLiteHoldWaitFinalizing, SQLiteHoldControlStateFinalizing)
}
func (control *sqliteHoldControl) wait(ctx context.Context, receipt SQLiteHoldReceipt, target SQLiteHoldWaitTarget, state SQLiteHoldControlState) (SQLiteHoldReceipt, error) {
record, err := control.activeRecord(receipt)
if err != nil {
return SQLiteHoldReceipt{}, err
}
snapshot, err := control.tracker.WaitSQLiteHold(ctx, record.session, target)
if err != nil {
return control.terminalReceipt(record, err)
}
if snapshot.Operation != SQLiteOperationUpdate || snapshot.Table != "api_tokens" {
_ = control.tracker.AbortSQLiteHold(record.session)
return control.terminalReceipt(record, ErrSQLiteHoldUnexpectedSelection)
}
return control.updateReceipt(record, state), nil
}
func (control *sqliteHoldControl) Release(receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) {
record, err := control.activeRecord(receipt)
if err != nil {
return SQLiteHoldReceipt{}, err
}
if err := control.tracker.ReleaseSQLiteHold(record.session); err != nil {
return control.terminalReceipt(record, err)
}
return control.updateReceipt(record, SQLiteHoldControlStateReleased), nil
}
func (control *sqliteHoldControl) Abort(receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) {
record, err := control.activeRecord(receipt)
if err != nil {
return SQLiteHoldReceipt{}, err
}
if err := control.tracker.AbortSQLiteHold(record.session); err != nil {
return control.terminalReceipt(record, err)
}
return control.updateReceipt(record, SQLiteHoldControlStateAborted), nil
}
func (control *sqliteHoldControl) Snapshot(receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) {
record, err := control.activeRecord(receipt)
if err != nil {
return SQLiteHoldReceipt{}, err
}
current := control.currentReceipt(record)
if current.State == SQLiteHoldControlStateReleased || current.State == SQLiteHoldControlStateAborted {
return current, nil
}
snapshot, snapshotErr := control.tracker.SQLiteHoldSnapshot(record.session)
if snapshotErr != nil {
if errors.Is(snapshotErr, ErrSQLiteHoldStaleSession) {
return control.updateReceipt(record, SQLiteHoldControlStateAborted), nil
}
return control.terminalReceipt(record, snapshotErr)
}
state := SQLiteHoldControlStateArmed
if snapshot.Finalizing {
state = SQLiteHoldControlStateFinalizing
} else if snapshot.Selected {
state = SQLiteHoldControlStateSelected
}
return control.updateReceipt(record, state), nil
}
func (control *sqliteHoldControl) currentReceipt(record *sqliteHoldControlRecord) SQLiteHoldReceipt {
control.mu.Lock()
defer control.mu.Unlock()
return record.receipt
}
func (control *sqliteHoldControl) newReceiptID() (string, error) {
bytes := make([]byte, 32)
if _, err := io.ReadFull(control.random, bytes); err != nil {
return "", err
}
return base64.RawURLEncoding.EncodeToString(bytes), nil
}
func (control *sqliteHoldControl) activeRecord(receipt SQLiteHoldReceipt) (*sqliteHoldControlRecord, error) {
control.mu.Lock()
defer control.mu.Unlock()
if control.record == nil || control.record.receipt.ID != receipt.ID {
return nil, ErrSQLiteHoldStaleSession
}
return control.record, nil
}
func (control *sqliteHoldControl) updateReceipt(record *sqliteHoldControlRecord, state SQLiteHoldControlState) SQLiteHoldReceipt {
control.mu.Lock()
defer control.mu.Unlock()
if control.record == record {
control.record.receipt.State = state
}
return record.receipt
}
func (control *sqliteHoldControl) terminalReceipt(record *sqliteHoldControlRecord, cause error) (SQLiteHoldReceipt, error) {
state := SQLiteHoldControlStateAborted
if cause == nil {
state = SQLiteHoldControlStateReleased
}
return control.updateReceipt(record, state), cause
}
@@ -0,0 +1,111 @@
//go:build agentcompat && linux
package singleton
import (
"bytes"
"context"
"encoding/json"
"errors"
"strings"
"testing"
)
func TestSQLiteHoldControlIssuesOpaqueReceiptAndStalesPriorReceipt(t *testing.T) {
// Given
control := newSQLiteHoldControl(NewSQLiteHoldTracker(), bytes.NewReader(append(bytes.Repeat([]byte{7}, 32), bytes.Repeat([]byte{8}, 32)...)))
// When
first, firstErr := control.ArmNextSQLiteHold()
if _, abortErr := control.Abort(first); abortErr != nil {
t.Fatal(abortErr)
}
aborted, abortSnapshotErr := control.Snapshot(first)
second, secondErr := control.ArmNextSQLiteHold()
_, staleErr := control.Snapshot(first)
// Then
if firstErr != nil || secondErr != nil {
t.Fatalf("first=%v second=%v", firstErr, secondErr)
}
if len(first.ID) != 43 || strings.Contains(first.ID, "=") || first.ID == second.ID {
t.Fatalf("opaque receipt IDs = %q, %q", first.ID, second.ID)
}
if !errors.Is(staleErr, ErrSQLiteHoldStaleSession) {
t.Fatalf("old receipt error = %v", staleErr)
}
if abortSnapshotErr != nil || aborted.State != SQLiteHoldControlStateAborted {
t.Fatalf("aborted receipt=%+v err=%v", aborted, abortSnapshotErr)
}
}
func TestSQLiteHoldControlReadsTerminalStatesAndRejectsUnexpectedSelection(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
control := newSQLiteHoldControl(tracker, bytes.NewReader(bytes.Repeat([]byte{9}, 64)))
receipt, err := control.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
transaction := sqliteHoldTestTransaction(66)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{Operation: SQLiteOperationInsert, Table: "settings", Journal: sqliteHoldTestJournal}); err != nil {
t.Fatal(err)
}
// When
_, waitErr := control.WaitSelected(context.Background(), receipt)
aborted, snapshotErr := control.Snapshot(receipt)
wire, marshalErr := json.Marshal(aborted)
// Then
if !errors.Is(waitErr, ErrSQLiteHoldUnexpectedSelection) {
t.Fatalf("unexpected selection error = %v", waitErr)
}
if snapshotErr != nil || aborted.State != SQLiteHoldControlStateAborted {
t.Fatalf("terminal receipt=%+v err=%v", aborted, snapshotErr)
}
if marshalErr != nil {
t.Fatal(marshalErr)
}
for _, forbidden := range []string{"transaction", "connection", "journal", "operation", "table", "path", "origin", "session_id"} {
if strings.Contains(strings.ToLower(string(wire)), forbidden) {
t.Fatalf("opaque receipt leaked %q: %s", forbidden, wire)
}
}
}
func TestSQLiteHoldControlReadsReleasedReceipt(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
control := newSQLiteHoldControl(tracker, bytes.NewReader(bytes.Repeat([]byte{11}, 32)))
receipt, err := control.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
transaction := sqliteHoldTestTransaction(67)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "api_tokens", Journal: sqliteHoldTestJournal}); err != nil {
t.Fatal(err)
}
if _, err := tracker.BeginSQLiteFinalization(transaction); err != nil {
t.Fatal(err)
}
// When
_, releaseErr := control.Release(receipt)
released, snapshotErr := control.Snapshot(receipt)
_, abortErr := control.Abort(receipt)
// Then
if releaseErr != nil || snapshotErr != nil || released.State != SQLiteHoldControlStateReleased {
t.Fatalf("release=%v receipt=%+v snapshot=%v", releaseErr, released, snapshotErr)
}
if !errors.Is(abortErr, ErrSQLiteHoldStaleSession) {
t.Fatalf("abort after release = %v", abortErr)
}
}
@@ -0,0 +1,24 @@
//go:build agentcompat && linux
package singleton
import "context"
func ArmNextSQLiteHold() (SQLiteHoldReceipt, error) {
return sqliteAttributionHoldControl.ArmNextSQLiteHold()
}
func WaitSQLiteHoldSelected(ctx context.Context, receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) {
return sqliteAttributionHoldControl.WaitSelected(ctx, receipt)
}
func WaitSQLiteHoldFinalizing(ctx context.Context, receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) {
return sqliteAttributionHoldControl.WaitFinalizing(ctx, receipt)
}
func SnapshotSQLiteHold(receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) {
return sqliteAttributionHoldControl.Snapshot(receipt)
}
func ReleaseSQLiteHold(receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) {
return sqliteAttributionHoldControl.Release(receipt)
}
func AbortSQLiteHold(receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) {
return sqliteAttributionHoldControl.Abort(receipt)
}
@@ -0,0 +1,289 @@
//go:build agentcompat && linux
package singleton
import (
"errors"
"testing"
)
var sqliteHoldTestJournal = SQLiteJournalIdentity{
MountID: 2, DeviceMajor: 8, DeviceMinor: 1, Inode: 13, BirthSeconds: 10, BirthNanoseconds: 20,
}
func sqliteHoldTestTransaction(id SQLiteTransactionIdentity) SQLiteTransaction {
return SQLiteTransaction{Connection: SQLiteConnectionIdentity(3), Identity: id}
}
func recordSQLiteHoldTestUpdate(t *testing.T, tracker *SQLiteHoldTracker, transaction SQLiteTransaction) {
t.Helper()
if err := tracker.RecordSQLiteExecution(transaction, SQLiteExecutionOrigin{
Operation: SQLiteOperationUpdate, Table: "settings", StackHash: 21, FirstNezhaFrame: "singleton.test",
}); err != nil {
t.Fatal(err)
}
if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{
Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal,
}); err != nil {
t.Fatal(err)
}
}
func TestSQLiteHoldTrackerSelectsFutureCandidateAfterZeroCandidateArm(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
transaction := sqliteHoldTestTransaction(1)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
// When
recordSQLiteHoldTestUpdate(t, tracker, transaction)
finalization, err := tracker.BeginSQLiteFinalization(transaction)
// Then
if err != nil {
t.Fatal(err)
}
if err := tracker.ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
if err := finalization.Wait(); err != nil {
t.Fatal(err)
}
}
func TestSQLiteHoldTrackerSelectsSingleActiveCandidate(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(2)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
// When
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
finalization, finalizationErr := tracker.BeginSQLiteFinalization(transaction)
// Then
if err != nil || finalizationErr != nil {
t.Fatalf("arm=%v finalization=%v", err, finalizationErr)
}
if err := tracker.ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
if err := finalization.Wait(); err != nil {
t.Fatal(err)
}
}
func TestSQLiteHoldTrackerAbortsDuplicateCandidates(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
first := sqliteHoldTestTransaction(3)
second := sqliteHoldTestTransaction(4)
for _, transaction := range []SQLiteTransaction{first, second} {
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
}
// When
_, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
// Then
if !errors.Is(err, ErrSQLiteHoldAmbiguousCandidate) {
t.Fatalf("expected ambiguous candidate error, got %v", err)
}
if _, err := tracker.BeginSQLiteFinalization(first); !errors.Is(err, ErrSQLiteHoldNotSelected) {
t.Fatalf("expected no selected transaction, got %v", err)
}
}
func TestSQLiteHoldTrackerStaleReleaseCannotReleaseNewSession(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
first := sqliteHoldTestTransaction(5)
if err := tracker.BeginSQLiteTransaction(first); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, first)
staleSession, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
if err := tracker.AbortSQLiteHold(staleSession); err != nil {
t.Fatal(err)
}
if err := tracker.FinishSQLiteTransaction(first); err != nil {
t.Fatal(err)
}
second := sqliteHoldTestTransaction(6)
if err := tracker.BeginSQLiteTransaction(second); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, second)
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
finalization, err := tracker.BeginSQLiteFinalization(second)
if err != nil {
t.Fatal(err)
}
// When
err = tracker.ReleaseSQLiteHold(staleSession)
// Then
if !errors.Is(err, ErrSQLiteHoldStaleSession) {
t.Fatalf("expected stale session error, got %v", err)
}
if finalization.Released() {
t.Fatal("stale release released the new session")
}
if err := tracker.ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
if err := finalization.Wait(); err != nil {
t.Fatal(err)
}
}
func TestSQLiteHoldTrackerLinearizesArmBeforeFinalization(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(7)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
// When
finalization, err := tracker.BeginSQLiteFinalization(transaction)
// Then
if err != nil {
t.Fatal(err)
}
if finalization.Released() {
t.Fatal("finalization released before its session release")
}
if err := tracker.ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
if err := finalization.Wait(); err != nil {
t.Fatal(err)
}
}
func TestSQLiteHoldTrackerRollbackNeverWaits(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(8)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
finalization, err := tracker.BeginSQLiteFinalization(transaction)
if err != nil {
t.Fatal(err)
}
// When
err = tracker.FinishSQLiteTransaction(transaction)
// Then
if err != nil {
t.Fatal(err)
}
if err := finalization.Wait(); !errors.Is(err, ErrSQLiteHoldAborted) {
t.Fatalf("expected aborted finalization, got %v", err)
}
if err := tracker.ReleaseSQLiteHold(session); !errors.Is(err, ErrSQLiteHoldStaleSession) {
t.Fatalf("expected stale release, got %v", err)
}
}
func TestSQLiteHoldTrackerAbortUnblocksSelectedFinalization(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(9)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
finalization, err := tracker.BeginSQLiteFinalization(transaction)
if err != nil {
t.Fatal(err)
}
// When
err = tracker.AbortSQLiteHold(session)
// Then
if err != nil {
t.Fatal(err)
}
if err := finalization.Wait(); !errors.Is(err, ErrSQLiteHoldAborted) {
t.Fatalf("expected aborted finalization, got %v", err)
}
}
func TestSQLiteHoldTrackerCleansUpAfterReleasedFinalization(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
first := sqliteHoldTestTransaction(10)
if err := tracker.BeginSQLiteTransaction(first); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, first)
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
finalization, err := tracker.BeginSQLiteFinalization(first)
if err != nil {
t.Fatal(err)
}
if err := tracker.ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
if err := finalization.Wait(); err != nil {
t.Fatal(err)
}
if err := tracker.FinishSQLiteTransaction(first); err != nil {
t.Fatal(err)
}
second := sqliteHoldTestTransaction(11)
if err := tracker.BeginSQLiteTransaction(second); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, second)
// When
_, err = tracker.ArmSQLiteHold(sqliteHoldTestJournal)
// Then
if err != nil {
t.Fatal(err)
}
}
@@ -0,0 +1,239 @@
//go:build agentcompat && linux
package singleton
import "context"
func (tracker *SQLiteHoldTracker) ArmSQLiteHold(journal SQLiteJournalIdentity) (SQLiteHoldSession, error) {
return tracker.armSQLiteHold(SQLiteHoldSelectionModeKnownJournal, journal)
}
func (tracker *SQLiteHoldTracker) ArmNextSQLiteHold() (SQLiteHoldSession, error) {
return tracker.armSQLiteHold(SQLiteHoldSelectionModeNextWriter, SQLiteJournalIdentity{})
}
func (tracker *SQLiteHoldTracker) armSQLiteHold(mode SQLiteHoldSelectionMode, journal SQLiteJournalIdentity) (SQLiteHoldSession, error) {
tracker.mu.Lock()
defer tracker.mu.Unlock()
if tracker.session != nil {
return SQLiteHoldSession{}, ErrSQLiteHoldSessionActive
}
tracker.nextSession++
tracker.session = &sqliteHoldSessionState{
id: SQLiteHoldSession{identity: tracker.nextSession}, mode: mode, journal: journal, notify: make(chan struct{}),
}
var candidates []SQLiteTransaction
for transaction, held := range tracker.transactions {
if tracker.isEligibleLocked(held) {
candidates = append(candidates, transaction)
}
}
if len(candidates) > 1 {
tracker.abortSQLiteHoldWithCandidatesLocked(ErrSQLiteHoldAmbiguousCandidate, candidates)
return SQLiteHoldSession{}, ErrSQLiteHoldAmbiguousCandidate
}
if len(candidates) == 1 {
tracker.selectSQLiteHoldLocked(candidates[0])
}
tracker.setAttributionEnabledLocked(true)
return tracker.session.id, nil
}
func (tracker *SQLiteHoldTracker) WaitSQLiteHold(ctx context.Context, session SQLiteHoldSession, target SQLiteHoldWaitTarget) (SQLiteHoldSnapshot, error) {
if !target.valid() {
return SQLiteHoldSnapshot{}, ErrSQLiteHoldInvalidWaitTarget
}
for {
tracker.mu.Lock()
if tracker.session == nil || tracker.session.id != session {
if tracker.terminal != nil && tracker.terminal.id == session {
if tracker.terminal.released {
snapshot := SQLiteHoldSnapshot{SessionID: session.ID(), Selected: tracker.terminal.hasSelected, Transaction: tracker.terminal.selected, Finalizing: true, Released: true}
tracker.mu.Unlock()
return snapshot, nil
}
cause := tracker.terminal.cause
tracker.mu.Unlock()
return SQLiteHoldSnapshot{}, cause
}
tracker.mu.Unlock()
return SQLiteHoldSnapshot{}, ErrSQLiteHoldStaleSession
}
if tracker.waitTargetReachedLocked(target) {
snapshot, err := tracker.snapshotLocked(session)
tracker.mu.Unlock()
return snapshot, err
}
notify := tracker.session.notify
tracker.mu.Unlock()
select {
case <-notify:
case <-ctx.Done():
return SQLiteHoldSnapshot{}, ctx.Err()
}
}
}
func (tracker *SQLiteHoldTracker) BeginSQLiteFinalization(transaction SQLiteTransaction) (*SQLiteHoldFinalization, error) {
tracker.mu.Lock()
defer tracker.mu.Unlock()
if tracker.session == nil || !tracker.session.hasSelected || tracker.session.selected != transaction {
return nil, ErrSQLiteHoldNotSelected
}
return tracker.beginSQLiteFinalizationLocked(transaction)
}
func (tracker *SQLiteHoldTracker) BeginSQLiteCommitFinalization(transaction SQLiteTransaction) (*SQLiteHoldFinalization, error) {
tracker.mu.Lock()
defer tracker.mu.Unlock()
if tracker.session != nil && tracker.session.hasSelected && tracker.session.selected == transaction {
return tracker.beginSQLiteFinalizationLocked(transaction)
}
if cause, ok := tracker.causes[transaction]; ok {
return nil, cause
}
return nil, ErrSQLiteHoldNotSelected
}
func (tracker *SQLiteHoldTracker) ReleaseSQLiteHold(session SQLiteHoldSession) error {
tracker.mu.Lock()
defer tracker.mu.Unlock()
if !tracker.matchesSessionLocked(session) || tracker.session.released {
return ErrSQLiteHoldStaleSession
}
if !tracker.session.hasSelected || tracker.session.finalization == nil {
return ErrSQLiteHoldFinalizationNotStarted
}
tracker.session.released = true
tracker.setAttributionEnabledLocked(false)
close(tracker.session.finalization.done)
tracker.notifySQLiteHoldLocked()
return nil
}
func (tracker *SQLiteHoldTracker) AbortSQLiteHold(session SQLiteHoldSession) error {
tracker.mu.Lock()
defer tracker.mu.Unlock()
if !tracker.matchesSessionLocked(session) || tracker.session.released {
return ErrSQLiteHoldStaleSession
}
tracker.abortSQLiteHoldLocked(ErrSQLiteHoldAborted)
return nil
}
func (tracker *SQLiteHoldTracker) SQLiteHoldSnapshot(session SQLiteHoldSession) (SQLiteHoldSnapshot, error) {
tracker.mu.Lock()
defer tracker.mu.Unlock()
if !tracker.matchesSessionLocked(session) {
return SQLiteHoldSnapshot{}, ErrSQLiteHoldStaleSession
}
return tracker.snapshotLocked(session)
}
func (tracker *SQLiteHoldTracker) snapshotLocked(session SQLiteHoldSession) (SQLiteHoldSnapshot, error) {
snapshot := SQLiteHoldSnapshot{SessionID: session.ID(), Mode: tracker.session.mode, Selected: tracker.session.hasSelected, Journal: tracker.session.journal, Released: tracker.session.released}
if !tracker.session.hasSelected {
return snapshot, nil
}
held := tracker.transactions[tracker.session.selected]
if held == nil {
return SQLiteHoldSnapshot{}, ErrSQLiteHoldNotSelected
}
snapshot.Transaction = held.transaction
if held.hasWrite {
snapshot.Operation, snapshot.Table, snapshot.Journal = held.write.Update.Operation, held.write.Update.Table, held.write.Update.Journal
snapshot.StackHash, snapshot.FirstNezhaFrame = held.write.Origin.StackHash, held.write.Origin.FirstNezhaFrame
} else {
snapshot.Operation, snapshot.Table, snapshot.Journal = held.update.Operation, held.update.Table, held.update.Journal
snapshot.StackHash, snapshot.FirstNezhaFrame = held.origin.StackHash, held.origin.FirstNezhaFrame
}
snapshot.Finalizing = held.finalizing
return snapshot, nil
}
func (tracker *SQLiteHoldTracker) selectSQLiteHoldLocked(transaction SQLiteTransaction) error {
if !tracker.session.hasSelected {
held := tracker.transactions[transaction]
tracker.session.selected, tracker.session.hasSelected = transaction, true
if tracker.session.mode == SQLiteHoldSelectionModeNextWriter {
tracker.session.journal = held.update.Journal
}
tracker.notifySQLiteHoldLocked()
return nil
}
if tracker.session.selected == transaction {
return nil
}
tracker.abortSQLiteHoldWithCandidatesLocked(ErrSQLiteHoldAmbiguousCandidate, []SQLiteTransaction{tracker.session.selected, transaction})
return ErrSQLiteHoldAmbiguousCandidate
}
func (tracker *SQLiteHoldTracker) matchesSessionLocked(session SQLiteHoldSession) bool {
return tracker.session != nil && tracker.session.id == session
}
func (tracker *SQLiteHoldTracker) activeSessionHasSelectedTransactionLocked(transaction SQLiteTransaction) bool {
return tracker.session != nil && tracker.session.hasSelected && tracker.session.selected == transaction
}
func (tracker *SQLiteHoldTracker) beginSQLiteFinalizationLocked(transaction SQLiteTransaction) (*SQLiteHoldFinalization, error) {
held, ok := tracker.transactionLocked(transaction)
if !ok || held.finalizing {
return nil, ErrSQLiteHoldFinalizationStarted
}
held.finalizing = true
finalization := &SQLiteHoldFinalization{done: make(chan struct{})}
tracker.session.finalization = finalization
tracker.notifySQLiteHoldLocked()
return finalization, nil
}
func (tracker *SQLiteHoldTracker) waitTargetReachedLocked(target SQLiteHoldWaitTarget) bool {
switch target {
case SQLiteHoldWaitSelected:
return tracker.session.hasSelected
case SQLiteHoldWaitFinalizing:
return tracker.session.finalization != nil
default:
return false
}
}
func (target SQLiteHoldWaitTarget) valid() bool {
return target == SQLiteHoldWaitSelected || target == SQLiteHoldWaitFinalizing
}
// Closing and replacing the channel under mu makes a snapshot-to-wait handoff race-free.
func (tracker *SQLiteHoldTracker) notifySQLiteHoldLocked() {
close(tracker.session.notify)
tracker.session.notify = make(chan struct{})
}
func (tracker *SQLiteHoldTracker) abortSQLiteHoldLocked(cause error) {
tracker.abortSQLiteHoldWithCandidatesLocked(cause, nil)
}
func (tracker *SQLiteHoldTracker) abortSQLiteHoldWithCandidatesLocked(cause error, candidates []SQLiteTransaction) {
tracker.setAttributionEnabledLocked(false)
if tracker.session.finalization != nil {
tracker.session.finalization.err = cause
close(tracker.session.finalization.done)
}
if tracker.session.hasSelected {
tracker.causes[tracker.session.selected] = cause
}
for _, transaction := range candidates {
tracker.causes[transaction] = cause
}
tracker.terminal = &sqliteHoldTerminalState{id: tracker.session.id, selected: tracker.session.selected, hasSelected: tracker.session.hasSelected, cause: cause}
tracker.notifySQLiteHoldLocked()
tracker.session = nil
}
func (tracker *SQLiteHoldTracker) setAttributionEnabledLocked(enabled bool) {
if tracker.attributionEnabled == nil {
return
}
// Attribution is hold-scoped because ordinary GORM RETURNING closes rows before EOF.
tracker.attributionEnabled.Store(enabled)
}
@@ -0,0 +1,75 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"errors"
"testing"
)
func TestSQLiteHoldTrackerCommitFinalizationCancellationAbortsUnreleasedHold(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(202)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
finalization, err := tracker.BeginSQLiteCommitFinalization(transaction)
if err != nil {
t.Fatal(err)
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
// When
waitErr := tracker.WaitSQLiteCommitFinalization(ctx, finalization)
releaseErr := tracker.ReleaseSQLiteHold(session)
// Then
if !errors.Is(waitErr, context.Canceled) {
t.Fatalf("cancelled finalization wait error = %v", waitErr)
}
if !errors.Is(finalization.Wait(), ErrSQLiteHoldAborted) {
t.Fatalf("finalization after cancellation = %v", finalization.Wait())
}
if !errors.Is(releaseErr, ErrSQLiteHoldStaleSession) {
t.Fatalf("release after cancellation-owned abort = %v", releaseErr)
}
}
func TestSQLiteHoldTrackerCommitFinalizationReleaseWinsOverAlreadyCancelledContext(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(203)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
finalization, err := tracker.BeginSQLiteCommitFinalization(transaction)
if err != nil {
t.Fatal(err)
}
if err := tracker.ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
// When
waitErr := tracker.WaitSQLiteCommitFinalization(ctx, finalization)
// Then
if waitErr != nil {
t.Fatalf("released finalization lost to already-cancelled context: %v", waitErr)
}
}
@@ -0,0 +1,33 @@
//go:build agentcompat && linux
package singleton
import "context"
// WaitSQLiteCommitFinalization linearizes cancellation against release under the tracker lock.
func (tracker *SQLiteHoldTracker) WaitSQLiteCommitFinalization(ctx context.Context, finalization *SQLiteHoldFinalization) error {
for {
tracker.mu.Lock()
select {
case <-finalization.done:
err := finalization.err
tracker.mu.Unlock()
return err
default:
}
if err := ctx.Err(); err != nil {
if tracker.session != nil && tracker.session.finalization == finalization && !tracker.session.released {
tracker.abortSQLiteHoldLocked(ErrSQLiteHoldAborted)
}
tracker.mu.Unlock()
return err
}
done := finalization.done
tracker.mu.Unlock()
select {
case <-done:
return finalization.err
case <-ctx.Done():
}
}
}
@@ -0,0 +1,113 @@
//go:build agentcompat && linux
package singleton
func (tracker *SQLiteHoldTracker) BeginSQLiteTransaction(transaction SQLiteTransaction) error {
tracker.mu.Lock()
defer tracker.mu.Unlock()
if _, active := tracker.transactions[transaction]; active {
return ErrSQLiteHoldTransactionActive
}
tracker.transactions[transaction] = &sqliteHeldTransaction{transaction: transaction}
return nil
}
func (tracker *SQLiteHoldTracker) RecordSQLiteExecution(transaction SQLiteTransaction, origin SQLiteExecutionOrigin) error {
tracker.mu.Lock()
defer tracker.mu.Unlock()
held, ok := tracker.transactionLocked(transaction)
if !ok {
return ErrSQLiteHoldNotSelected
}
if held.finalizing {
return ErrSQLiteHoldFinalizationStarted
}
if held.hasWrite {
return ErrSQLiteHoldAtomicWriteRecorded
}
if tracker.activeSessionHasSelectedTransactionLocked(transaction) {
return ErrSQLiteHoldEvidenceFrozen
}
held.origin = origin
return nil
}
func (tracker *SQLiteHoldTracker) RecordSQLiteUpdate(transaction SQLiteTransaction, update SQLiteUpdateObservation) error {
tracker.mu.Lock()
defer tracker.mu.Unlock()
held, ok := tracker.transactionLocked(transaction)
if !ok {
return ErrSQLiteHoldNotSelected
}
if held.finalizing {
return ErrSQLiteHoldFinalizationStarted
}
if held.hasWrite {
return ErrSQLiteHoldAtomicWriteRecorded
}
if tracker.activeSessionHasSelectedTransactionLocked(transaction) {
return ErrSQLiteHoldEvidenceFrozen
}
held.update, held.hasUpdate = update, true
if tracker.session != nil && !tracker.session.released && tracker.isEligibleLocked(held) {
return tracker.selectSQLiteHoldLocked(transaction)
}
return nil
}
func (tracker *SQLiteHoldTracker) RecordSQLiteWrite(transaction SQLiteTransaction, write SQLiteWriteObservation) error {
tracker.mu.Lock()
defer tracker.mu.Unlock()
held, ok := tracker.transactionLocked(transaction)
if !ok {
return ErrSQLiteHoldNotSelected
}
if held.finalizing {
return ErrSQLiteHoldFinalizationStarted
}
if held.hasWrite {
return nil
}
if tracker.activeSessionHasSelectedTransactionLocked(transaction) {
return ErrSQLiteHoldEvidenceFrozen
}
held.write, held.hasWrite = write, true
held.origin, held.update, held.hasUpdate = write.Origin, write.Update, true
if tracker.session != nil && !tracker.session.released && tracker.isEligibleLocked(held) {
return tracker.selectSQLiteHoldLocked(transaction)
}
return nil
}
func (tracker *SQLiteHoldTracker) FinishSQLiteTransaction(transaction SQLiteTransaction) error {
tracker.mu.Lock()
defer tracker.mu.Unlock()
if _, ok := tracker.transactionLocked(transaction); !ok {
return ErrSQLiteHoldNotSelected
}
delete(tracker.transactions, transaction)
defer delete(tracker.causes, transaction)
defer delete(tracker.terminalArbitrations, transaction)
if tracker.session != nil && tracker.session.hasSelected && tracker.session.selected == transaction {
if tracker.session.released {
tracker.terminal = &sqliteHoldTerminalState{id: tracker.session.id, selected: transaction, hasSelected: true, released: true}
tracker.notifySQLiteHoldLocked()
tracker.session = nil
} else {
tracker.abortSQLiteHoldLocked(ErrSQLiteHoldAborted)
}
}
return nil
}
func (tracker *SQLiteHoldTracker) transactionLocked(transaction SQLiteTransaction) (*sqliteHeldTransaction, bool) {
held, ok := tracker.transactions[transaction]
return held, ok
}
func (tracker *SQLiteHoldTracker) isEligibleLocked(held *sqliteHeldTransaction) bool {
if !held.hasUpdate || held.finalizing {
return false
}
return tracker.session.mode == SQLiteHoldSelectionModeNextWriter || held.update.Journal == tracker.session.journal
}
@@ -0,0 +1,228 @@
//go:build agentcompat && linux
package singleton
import (
"errors"
"testing"
)
func sqliteHoldTestTransactionOnConnection(connection SQLiteConnectionIdentity, id SQLiteTransactionIdentity) SQLiteTransaction {
return SQLiteTransaction{Connection: connection, Identity: id}
}
func TestSQLiteHoldTrackerTracksSameTransactionIdentityAcrossConnections(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
first := sqliteHoldTestTransactionOnConnection(3, 31)
second := sqliteHoldTestTransactionOnConnection(4, 31)
for _, transaction := range []SQLiteTransaction{first, second} {
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
}
// When
_, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
// Then
if !errors.Is(err, ErrSQLiteHoldAmbiguousCandidate) {
t.Fatalf("expected both connections to be tracked, got %v", err)
}
}
func TestSQLiteHoldTrackerRejectsExactDuplicateTransaction(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransactionOnConnection(3, 32)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
// When
err := tracker.BeginSQLiteTransaction(transaction)
// Then
if !errors.Is(err, ErrSQLiteHoldTransactionActive) {
t.Fatalf("expected duplicate composite rejection, got %v", err)
}
}
func TestSQLiteHoldTrackerFreezesSelectedMetadataAfterFinalization(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(33)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
if err := tracker.RecordSQLiteExecution(transaction, SQLiteExecutionOrigin{
Operation: SQLiteOperationInsert, Table: "initial", StackHash: 44, FirstNezhaFrame: "singleton.initial",
}); err != nil {
t.Fatal(err)
}
if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{
Operation: SQLiteOperationInsert, Table: "initial", Journal: sqliteHoldTestJournal,
}); err != nil {
t.Fatal(err)
}
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
if _, err := tracker.BeginSQLiteFinalization(transaction); err != nil {
t.Fatal(err)
}
// When
executionErr := tracker.RecordSQLiteExecution(transaction, SQLiteExecutionOrigin{
Operation: SQLiteOperationDelete, Table: "changed", StackHash: 45, FirstNezhaFrame: "singleton.changed",
})
updateErr := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{
Operation: SQLiteOperationDelete, Table: "changed", Journal: SQLiteJournalIdentity{Inode: 99},
})
snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session)
// Then
if !errors.Is(executionErr, ErrSQLiteHoldFinalizationStarted) || !errors.Is(updateErr, ErrSQLiteHoldFinalizationStarted) {
t.Fatalf("execution=%v update=%v", executionErr, updateErr)
}
if snapshotErr != nil {
t.Fatal(snapshotErr)
}
if snapshot.Operation != SQLiteOperationInsert || snapshot.Table != "initial" || snapshot.StackHash != 44 ||
snapshot.FirstNezhaFrame != "singleton.initial" || snapshot.Journal != sqliteHoldTestJournal {
t.Fatalf("finalizing metadata changed: %+v", snapshot)
}
}
func TestSQLiteHoldTrackerRejectsSelectedReleaseBeforeFinalizationThenCompletes(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(34)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
// When
releaseErr := tracker.ReleaseSQLiteHold(session)
finalization, finalizationErr := tracker.BeginSQLiteFinalization(transaction)
// Then
if !errors.Is(releaseErr, ErrSQLiteHoldFinalizationNotStarted) || finalizationErr != nil {
t.Fatalf("release=%v finalization=%v", releaseErr, finalizationErr)
}
if err := tracker.ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
if err := finalization.Wait(); err != nil {
t.Fatal(err)
}
}
func TestSQLiteHoldTrackerCompletesSecondLifecycleAfterReleasedCleanup(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
first := sqliteHoldTestTransaction(35)
if err := tracker.BeginSQLiteTransaction(first); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, first)
firstSession, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
firstFinalization, err := tracker.BeginSQLiteFinalization(first)
if err != nil {
t.Fatal(err)
}
if err := tracker.ReleaseSQLiteHold(firstSession); err != nil {
t.Fatal(err)
}
if err := firstFinalization.Wait(); err != nil {
t.Fatal(err)
}
if err := tracker.FinishSQLiteTransaction(first); err != nil {
t.Fatal(err)
}
second := sqliteHoldTestTransaction(36)
if err := tracker.BeginSQLiteTransaction(second); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, second)
// When
secondSession, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
secondFinalization, finalizationErr := tracker.BeginSQLiteFinalization(second)
// Then
if err != nil || finalizationErr != nil {
t.Fatalf("arm=%v finalization=%v", err, finalizationErr)
}
if err := tracker.ReleaseSQLiteHold(secondSession); err != nil {
t.Fatal(err)
}
if err := secondFinalization.Wait(); err != nil {
t.Fatal(err)
}
if err := tracker.FinishSQLiteTransaction(second); err != nil {
t.Fatal(err)
}
}
func TestSQLiteHoldTrackerFreezesLegacyEvidenceWhenSelected(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(37)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
originalOrigin := SQLiteExecutionOrigin{
Operation: SQLiteOperationInsert, Table: "original", StackHash: 46, FirstNezhaFrame: "singleton.original",
}
if err := tracker.RecordSQLiteExecution(transaction, originalOrigin); err != nil {
t.Fatal(err)
}
if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{
Operation: SQLiteOperationInsert, Table: "original", Journal: sqliteHoldTestJournal,
}); err != nil {
t.Fatal(err)
}
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
// When
executionErr := tracker.RecordSQLiteExecution(transaction, SQLiteExecutionOrigin{
Operation: SQLiteOperationDelete, Table: "changed", StackHash: 47, FirstNezhaFrame: "singleton.changed",
})
updateErr := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{
Operation: SQLiteOperationDelete, Table: "changed", Journal: SQLiteJournalIdentity{Inode: 94},
})
snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session)
finalization, finalizationErr := tracker.BeginSQLiteFinalization(transaction)
// Then
if !errors.Is(executionErr, ErrSQLiteHoldEvidenceFrozen) || !errors.Is(updateErr, ErrSQLiteHoldEvidenceFrozen) {
t.Fatalf("execution=%v update=%v", executionErr, updateErr)
}
if snapshotErr != nil || finalizationErr != nil {
t.Fatalf("snapshot=%v finalization=%v", snapshotErr, finalizationErr)
}
if snapshot.Operation != SQLiteOperationInsert || snapshot.Table != "original" || snapshot.StackHash != 46 ||
snapshot.FirstNezhaFrame != "singleton.original" || snapshot.Journal != sqliteHoldTestJournal {
t.Fatalf("selected legacy evidence changed: %+v", snapshot)
}
if err := tracker.ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
if err := finalization.Wait(); err != nil {
t.Fatal(err)
}
}
@@ -0,0 +1,141 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"errors"
"os"
"os/exec"
"testing"
"time"
)
func TestMain(m *testing.M) {
if os.Getenv("NEZHA_SQLITE_HOLD_INVALID_TARGET_HELPER") == "1" {
tracker := NewSQLiteHoldTracker()
session, err := tracker.ArmNextSQLiteHold()
if err != nil {
os.Exit(2)
}
_, waitErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitTarget(99))
if !errors.Is(waitErr, ErrSQLiteHoldInvalidWaitTarget) {
os.Exit(3)
}
os.Exit(0)
}
os.Exit(m.Run())
}
func TestSQLiteHoldTrackerWaitRejectsInvalidTargetImmediately(t *testing.T) {
// Given
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
command := exec.CommandContext(ctx, os.Args[0], "-test.run=^TestSQLiteHoldTrackerWaitRejectsInvalidTargetImmediately$")
command.Env = append(os.Environ(), "NEZHA_SQLITE_HOLD_INVALID_TARGET_HELPER=1")
// When
err := command.Run()
// Then
if err != nil {
t.Fatalf("invalid target helper did not return before watchdog deadline: %v", err)
}
}
func TestSQLiteHoldTrackerRecordsAmbiguityForEveryFutureConflictCandidate(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
_, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
first, second, mismatch := sqliteHoldTestTransaction(81), sqliteHoldTestTransaction(82), sqliteHoldTestTransaction(83)
for _, transaction := range []SQLiteTransaction{first, second, mismatch} {
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
}
if err := tracker.RecordSQLiteUpdate(first, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal}); err != nil {
t.Fatal(err)
}
if err := tracker.RecordSQLiteUpdate(mismatch, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: SQLiteJournalIdentity{Inode: 84}}); err != nil {
t.Fatal(err)
}
// When
duplicateErr := tracker.RecordSQLiteUpdate(second, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal})
_, firstErr := tracker.BeginSQLiteCommitFinalization(first)
_, secondErr := tracker.BeginSQLiteCommitFinalization(second)
_, mismatchErr := tracker.BeginSQLiteCommitFinalization(mismatch)
// Then
if !errors.Is(duplicateErr, ErrSQLiteHoldAmbiguousCandidate) || !errors.Is(firstErr, ErrSQLiteHoldAmbiguousCandidate) || !errors.Is(secondErr, ErrSQLiteHoldAmbiguousCandidate) {
t.Fatalf("duplicate=%v first=%v second=%v", duplicateErr, firstErr, secondErr)
}
if !errors.Is(mismatchErr, ErrSQLiteHoldNotSelected) {
t.Fatalf("mismatched transaction cause = %v", mismatchErr)
}
}
func TestSQLiteHoldTrackerRecordsAmbiguityForEveryArmTimeConflictCandidate(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
first, second, mismatch := sqliteHoldTestTransaction(85), sqliteHoldTestTransaction(86), sqliteHoldTestTransaction(87)
for _, transaction := range []SQLiteTransaction{first, second, mismatch} {
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
}
for _, transaction := range []SQLiteTransaction{first, second} {
if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal}); err != nil {
t.Fatal(err)
}
}
if err := tracker.RecordSQLiteUpdate(mismatch, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: SQLiteJournalIdentity{Inode: 88}}); err != nil {
t.Fatal(err)
}
// When
_, armErr := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
_, firstErr := tracker.BeginSQLiteCommitFinalization(first)
_, secondErr := tracker.BeginSQLiteCommitFinalization(second)
_, mismatchErr := tracker.BeginSQLiteCommitFinalization(mismatch)
// Then
if !errors.Is(armErr, ErrSQLiteHoldAmbiguousCandidate) || !errors.Is(firstErr, ErrSQLiteHoldAmbiguousCandidate) || !errors.Is(secondErr, ErrSQLiteHoldAmbiguousCandidate) {
t.Fatalf("arm=%v first=%v second=%v", armErr, firstErr, secondErr)
}
if !errors.Is(mismatchErr, ErrSQLiteHoldNotSelected) {
t.Fatalf("mismatched transaction cause = %v", mismatchErr)
}
}
func TestSQLiteHoldTrackerFinishDeletesAmbiguityCause(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
first, second := sqliteHoldTestTransaction(89), sqliteHoldTestTransaction(90)
for _, transaction := range []SQLiteTransaction{first, second} {
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal}); err != nil {
t.Fatal(err)
}
}
if _, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal); !errors.Is(err, ErrSQLiteHoldAmbiguousCandidate) {
t.Fatal(err)
}
// When
if err := tracker.FinishSQLiteTransaction(first); err != nil {
t.Fatal(err)
}
_, lookupErr := tracker.BeginSQLiteCommitFinalization(first)
// Then
if !errors.Is(lookupErr, ErrSQLiteHoldNotSelected) {
t.Fatalf("finished transaction ambiguity cause = %v", lookupErr)
}
}
@@ -0,0 +1,265 @@
//go:build agentcompat && linux
package singleton
import (
"errors"
"testing"
)
func TestSQLiteHoldTrackerAbortsSessionWhenFutureDuplicateArrives(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
first := sqliteHoldTestTransaction(12)
if err := tracker.BeginSQLiteTransaction(first); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, first)
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
second := sqliteHoldTestTransaction(13)
if err := tracker.BeginSQLiteTransaction(second); err != nil {
t.Fatal(err)
}
// When
err = tracker.RecordSQLiteUpdate(second, SQLiteUpdateObservation{
Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal,
})
// Then
if !errors.Is(err, ErrSQLiteHoldAmbiguousCandidate) {
t.Fatalf("expected ambiguous candidate error, got %v", err)
}
if err := tracker.ReleaseSQLiteHold(session); !errors.Is(err, ErrSQLiteHoldStaleSession) {
t.Fatalf("expected stale session after duplicate candidate, got %v", err)
}
if _, err := tracker.BeginSQLiteFinalization(first); !errors.Is(err, ErrSQLiteHoldNotSelected) {
t.Fatalf("expected aborted session, got %v", err)
}
}
func TestSQLiteHoldTrackerRejectsRepeatedRelease(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(21)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
if _, err := tracker.BeginSQLiteFinalization(transaction); err != nil {
t.Fatal(err)
}
if err := tracker.ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
// When
err = tracker.ReleaseSQLiteHold(session)
// Then
if !errors.Is(err, ErrSQLiteHoldStaleSession) {
t.Fatalf("expected stale repeated release, got %v", err)
}
}
func TestSQLiteHoldTrackerRejectsReleaseBeforeFinalization(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
// When
err = tracker.ReleaseSQLiteHold(session)
// Then
if !errors.Is(err, ErrSQLiteHoldFinalizationNotStarted) {
t.Fatalf("expected finalization-not-started error, got %v", err)
}
}
func TestSQLiteHoldTrackerRejectsAbortAfterRelease(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(14)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
finalization, err := tracker.BeginSQLiteFinalization(transaction)
if err != nil {
t.Fatal(err)
}
if err := tracker.ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
// When
err = tracker.AbortSQLiteHold(session)
// Then
if !errors.Is(err, ErrSQLiteHoldStaleSession) {
t.Fatalf("expected stale session error, got %v", err)
}
if err := finalization.Wait(); err != nil {
t.Fatal(err)
}
}
func TestSQLiteHoldTrackerIgnoresFutureCandidateAfterRelease(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
first := sqliteHoldTestTransaction(15)
if err := tracker.BeginSQLiteTransaction(first); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, first)
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
finalization, err := tracker.BeginSQLiteFinalization(first)
if err != nil {
t.Fatal(err)
}
if err := tracker.ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
second := sqliteHoldTestTransaction(16)
if err := tracker.BeginSQLiteTransaction(second); err != nil {
t.Fatal(err)
}
// When
err = tracker.RecordSQLiteUpdate(second, SQLiteUpdateObservation{
Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal,
})
// Then
if err != nil {
t.Fatal(err)
}
if err := finalization.Wait(); err != nil {
t.Fatal(err)
}
}
func TestSQLiteHoldTrackerDoesNotMatchJournalWithDifferentMountOrBirthTime(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
journal := SQLiteJournalIdentity{MountID: 2, DeviceMajor: 8, DeviceMinor: 1, Inode: 13, BirthSeconds: 10, BirthNanoseconds: 20}
transaction := sqliteHoldTestTransaction(17)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
// When
if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{
Operation: SQLiteOperationInsert, Table: "settings", Journal: SQLiteJournalIdentity{
MountID: 3, DeviceMajor: 8, DeviceMinor: 1, Inode: 13, BirthSeconds: 10, BirthNanoseconds: 20,
},
}); err != nil {
t.Fatal(err)
}
session, err := tracker.ArmSQLiteHold(journal)
// Then
if err != nil {
t.Fatal(err)
}
if _, err := tracker.BeginSQLiteFinalization(transaction); !errors.Is(err, ErrSQLiteHoldNotSelected) {
t.Fatalf("expected mount mismatch not to select, got %v", err)
}
if err := tracker.AbortSQLiteHold(session); err != nil {
t.Fatal(err)
}
}
func TestSQLiteHoldTrackerDoesNotMatchJournalWithDifferentBirthTime(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
journal := SQLiteJournalIdentity{MountID: 2, DeviceMajor: 8, DeviceMinor: 1, Inode: 13, BirthSeconds: 10, BirthNanoseconds: 20}
transaction := sqliteHoldTestTransaction(18)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
// When
if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{
Operation: SQLiteOperationDelete, Table: "settings", Journal: SQLiteJournalIdentity{
MountID: 2, DeviceMajor: 8, DeviceMinor: 1, Inode: 13, BirthSeconds: 10, BirthNanoseconds: 21,
},
}); err != nil {
t.Fatal(err)
}
session, err := tracker.ArmSQLiteHold(journal)
// Then
if err != nil {
t.Fatal(err)
}
if _, err := tracker.BeginSQLiteFinalization(transaction); !errors.Is(err, ErrSQLiteHoldNotSelected) {
t.Fatalf("expected birth-time mismatch not to select, got %v", err)
}
if err := tracker.AbortSQLiteHold(session); err != nil {
t.Fatal(err)
}
}
func TestSQLiteHoldTrackerRejectsDuplicateTransactionIdentity(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(19)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
// When
err := tracker.BeginSQLiteTransaction(transaction)
// Then
if !errors.Is(err, ErrSQLiteHoldTransactionActive) {
t.Fatalf("expected transaction-active error, got %v", err)
}
}
func TestSQLiteHoldTrackerSelectsDeleteOperation(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(20)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{
Operation: SQLiteOperationDelete, Table: "settings", Journal: sqliteHoldTestJournal,
}); err != nil {
t.Fatal(err)
}
// When
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
finalization, finalizationErr := tracker.BeginSQLiteFinalization(transaction)
// Then
if err != nil || finalizationErr != nil {
t.Fatalf("arm=%v finalization=%v", err, finalizationErr)
}
if err := tracker.ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
if err := finalization.Wait(); err != nil {
t.Fatal(err)
}
}
@@ -0,0 +1,29 @@
//go:build agentcompat && linux
package singleton
type sqliteAttributionTerminalArbitration uint8
const (
sqliteAttributionTerminalNoReservation sqliteAttributionTerminalArbitration = iota
sqliteAttributionTerminalReleaseWon
sqliteAttributionTerminalRollbackGranted
sqliteAttributionTerminalRollbackReserved
)
func (tracker *SQLiteHoldTracker) ArbitrateSQLiteCommitWaitingTerminal(transaction SQLiteTransaction) sqliteAttributionTerminalArbitration {
tracker.mu.Lock()
defer tracker.mu.Unlock()
if !tracker.activeSessionHasSelectedTransactionLocked(transaction) {
if _, reserved := tracker.terminalArbitrations[transaction]; reserved {
return sqliteAttributionTerminalRollbackReserved
}
return sqliteAttributionTerminalNoReservation
}
if tracker.session.released {
return sqliteAttributionTerminalReleaseWon
}
tracker.abortSQLiteHoldLocked(ErrSQLiteHoldAborted)
tracker.terminalArbitrations[transaction] = struct{}{}
return sqliteAttributionTerminalRollbackGranted
}
@@ -0,0 +1,169 @@
//go:build agentcompat && linux
package singleton
import (
"errors"
"sync"
"sync/atomic"
)
var (
ErrSQLiteHoldSessionActive = errors.New("sqlite hold session already active")
ErrSQLiteHoldAmbiguousCandidate = errors.New("sqlite hold has multiple matching candidates")
ErrSQLiteHoldNotSelected = errors.New("sqlite hold transaction is not selected")
ErrSQLiteHoldFinalizationStarted = errors.New("sqlite hold finalization already started")
ErrSQLiteHoldFinalizationNotStarted = errors.New("sqlite hold finalization not started")
ErrSQLiteHoldStaleSession = errors.New("sqlite hold session is stale")
ErrSQLiteHoldAborted = errors.New("sqlite hold session aborted")
ErrSQLiteHoldTransactionActive = errors.New("sqlite hold transaction already active")
ErrSQLiteHoldAtomicWriteRecorded = errors.New("sqlite hold atomic write already recorded")
ErrSQLiteHoldEvidenceFrozen = errors.New("sqlite hold selected evidence is frozen")
ErrSQLiteHoldInvalidWaitTarget = errors.New("sqlite hold wait target is invalid")
)
type SQLiteConnectionIdentity uint64
type SQLiteTransactionIdentity uint64
type SQLiteJournalIdentity struct {
MountID uint64
DeviceMajor uint32
DeviceMinor uint32
Inode uint64
BirthSeconds int64
BirthNanoseconds uint32
}
type SQLiteOperation string
const (
SQLiteOperationInsert SQLiteOperation = "insert"
SQLiteOperationUpdate SQLiteOperation = "update"
SQLiteOperationDelete SQLiteOperation = "delete"
)
type SQLiteTransaction struct {
Connection SQLiteConnectionIdentity
Identity SQLiteTransactionIdentity
}
type SQLiteExecutionOrigin struct {
Operation SQLiteOperation
Table string
StackHash uint64
FirstNezhaFrame string
}
type SQLiteUpdateObservation struct {
Operation SQLiteOperation
Table string
Journal SQLiteJournalIdentity
}
type SQLiteWriteObservation struct {
Origin SQLiteExecutionOrigin
Update SQLiteUpdateObservation
}
type SQLiteHoldSelectionMode uint8
const (
SQLiteHoldSelectionModeKnownJournal SQLiteHoldSelectionMode = iota + 1
SQLiteHoldSelectionModeNextWriter
)
type SQLiteHoldSession struct{ identity uint64 }
func (session SQLiteHoldSession) ID() uint64 { return session.identity }
type SQLiteHoldSnapshot struct {
SessionID uint64
Mode SQLiteHoldSelectionMode
Selected bool
Transaction SQLiteTransaction
Operation SQLiteOperation
Table string
StackHash uint64
FirstNezhaFrame string
Journal SQLiteJournalIdentity
Finalizing bool
Released bool
Aborted bool
}
type SQLiteHoldFinalization struct {
done chan struct{}
err error
}
type SQLiteHoldWaitTarget uint8
const (
SQLiteHoldWaitSelected SQLiteHoldWaitTarget = iota + 1
SQLiteHoldWaitFinalizing
)
func (finalization *SQLiteHoldFinalization) Wait() error {
<-finalization.done
return finalization.err
}
func (finalization *SQLiteHoldFinalization) Released() bool {
select {
case <-finalization.done:
return finalization.err == nil
default:
return false
}
}
type sqliteHeldTransaction struct {
transaction SQLiteTransaction
origin SQLiteExecutionOrigin
update SQLiteUpdateObservation
hasUpdate bool
write SQLiteWriteObservation
hasWrite bool
finalizing bool
}
type sqliteHoldSessionState struct {
id SQLiteHoldSession
mode SQLiteHoldSelectionMode
journal SQLiteJournalIdentity
selected SQLiteTransaction
hasSelected bool
released bool
finalization *SQLiteHoldFinalization
notify chan struct{}
}
type sqliteHoldTerminalState struct {
id SQLiteHoldSession
selected SQLiteTransaction
hasSelected bool
released bool
cause error
}
// SQLiteHoldTracker linearizes state transitions; finalization waits only after releasing mu.
type SQLiteHoldTracker struct {
mu sync.Mutex
attributionEnabled *atomic.Bool
nextSession uint64
transactions map[SQLiteTransaction]*sqliteHeldTransaction
session *sqliteHoldSessionState
terminal *sqliteHoldTerminalState
causes map[SQLiteTransaction]error
terminalArbitrations map[SQLiteTransaction]struct{}
}
func NewSQLiteHoldTracker() *SQLiteHoldTracker {
return &SQLiteHoldTracker{transactions: make(map[SQLiteTransaction]*sqliteHeldTransaction), causes: make(map[SQLiteTransaction]error), terminalArbitrations: make(map[SQLiteTransaction]struct{})}
}
func newSQLiteAttributionHoldTracker(enabled *atomic.Bool) *SQLiteHoldTracker {
tracker := NewSQLiteHoldTracker()
tracker.attributionEnabled = enabled
return tracker
}
@@ -0,0 +1,149 @@
//go:build agentcompat && linux
package singleton
import (
"errors"
"testing"
)
func TestSQLiteHoldTrackerNextWriterSelectsFutureUpdateAndPublishesSnapshot(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
session, err := tracker.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
before, err := tracker.SQLiteHoldSnapshot(session)
if err != nil {
t.Fatal(err)
}
transaction := sqliteHoldTestTransaction(22)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
// When
if err := tracker.RecordSQLiteExecution(transaction, SQLiteExecutionOrigin{
Operation: SQLiteOperationInsert, Table: "settings", StackHash: 34, FirstNezhaFrame: "singleton.next",
}); err != nil {
t.Fatal(err)
}
if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{
Operation: SQLiteOperationInsert, Table: "settings", Journal: sqliteHoldTestJournal,
}); err != nil {
t.Fatal(err)
}
after, err := tracker.SQLiteHoldSnapshot(session)
// Then
if err != nil {
t.Fatal(err)
}
if before.SessionID != session.ID() || before.Mode != SQLiteHoldSelectionModeNextWriter || before.Selected {
t.Fatalf("unexpected unselected snapshot: %+v", before)
}
if !after.Selected || after.Transaction != transaction || after.Operation != SQLiteOperationInsert || after.Table != "settings" ||
after.StackHash != 34 || after.FirstNezhaFrame != "singleton.next" || after.Journal != sqliteHoldTestJournal {
t.Fatalf("unexpected selected snapshot: %+v", after)
}
}
func TestSQLiteHoldTrackerNextWriterSelectsOneActiveUpdate(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(23)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
// When
session, err := tracker.ArmNextSQLiteHold()
snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session)
// Then
if err != nil || snapshotErr != nil {
t.Fatalf("arm=%v snapshot=%v", err, snapshotErr)
}
if !snapshot.Selected || snapshot.Journal != sqliteHoldTestJournal {
t.Fatalf("unexpected active selection snapshot: %+v", snapshot)
}
}
func TestSQLiteHoldTrackerNextWriterAbortsMultipleActiveUpdates(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
for _, transaction := range []SQLiteTransaction{sqliteHoldTestTransaction(24), sqliteHoldTestTransaction(25)} {
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
}
// When
_, err := tracker.ArmNextSQLiteHold()
// Then
if !errors.Is(err, ErrSQLiteHoldAmbiguousCandidate) {
t.Fatalf("expected ambiguous active updates, got %v", err)
}
}
func TestSQLiteHoldTrackerNextWriterAbortsDuplicateFutureUpdate(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
session, err := tracker.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
first := sqliteHoldTestTransaction(26)
second := sqliteHoldTestTransaction(27)
for _, transaction := range []SQLiteTransaction{first, second} {
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
}
if err := tracker.RecordSQLiteUpdate(first, SQLiteUpdateObservation{
Operation: SQLiteOperationDelete, Table: "settings", Journal: sqliteHoldTestJournal,
}); err != nil {
t.Fatal(err)
}
// When
err = tracker.RecordSQLiteUpdate(second, SQLiteUpdateObservation{
Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal,
})
// Then
if !errors.Is(err, ErrSQLiteHoldAmbiguousCandidate) {
t.Fatalf("expected ambiguous future updates, got %v", err)
}
if _, err := tracker.SQLiteHoldSnapshot(session); !errors.Is(err, ErrSQLiteHoldStaleSession) {
t.Fatalf("expected stale session after duplicate future update, got %v", err)
}
}
func TestSQLiteHoldTrackerNextWriterRequiresFinalizationBeforeReleaseAndStalesAfterAbort(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
session, err := tracker.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
// When
releaseErr := tracker.ReleaseSQLiteHold(session)
abortErr := tracker.AbortSQLiteHold(session)
// Then
if !errors.Is(releaseErr, ErrSQLiteHoldFinalizationNotStarted) {
t.Fatalf("expected release rejection, got %v", releaseErr)
}
if abortErr != nil {
t.Fatal(abortErr)
}
if _, err := tracker.SQLiteHoldSnapshot(session); !errors.Is(err, ErrSQLiteHoldStaleSession) {
t.Fatalf("expected stale aborted session, got %v", err)
}
}
@@ -0,0 +1,238 @@
//go:build agentcompat && linux
package singleton
import (
"context"
"errors"
"testing"
)
func TestSQLiteHoldTrackerWaitSelectedAndFinalizingObserveTransitionsWithoutLostWakeup(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
session, err := tracker.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
transaction := sqliteHoldTestTransaction(61)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
// When
selected, selectedErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitSelected)
if _, err := tracker.BeginSQLiteFinalization(transaction); err != nil {
t.Fatal(err)
}
finalizing, finalizingErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing)
// Then
if selectedErr != nil || !selected.Selected {
t.Fatalf("selected=%+v err=%v", selected, selectedErr)
}
if finalizingErr != nil || !finalizing.Finalizing {
t.Fatalf("finalizing=%+v err=%v", finalizing, finalizingErr)
}
}
func TestSQLiteHoldTrackerWaitCancellationLeavesSessionActive(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
session, err := tracker.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
// When
_, waitErr := tracker.WaitSQLiteHold(ctx, session, SQLiteHoldWaitSelected)
snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session)
// Then
if !errors.Is(waitErr, context.Canceled) {
t.Fatalf("wait error = %v, want context cancellation", waitErr)
}
if snapshotErr != nil || snapshot.Selected {
t.Fatalf("snapshot=%+v err=%v", snapshot, snapshotErr)
}
}
func TestSQLiteHoldTrackerWaitReturnsTerminalAbortAndAmbiguityCauses(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
session, err := tracker.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
abortResult := make(chan error, 1)
go func() {
_, waitErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitSelected)
abortResult <- waitErr
}()
// When
if err := tracker.AbortSQLiteHold(session); err != nil {
t.Fatal(err)
}
abortErr := <-abortResult
if _, err := tracker.ArmNextSQLiteHold(); err != nil {
t.Fatal(err)
}
first, second := sqliteHoldTestTransaction(62), sqliteHoldTestTransaction(63)
for _, transaction := range []SQLiteTransaction{first, second} {
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
}
if err := tracker.RecordSQLiteUpdate(first, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal}); err != nil {
t.Fatal(err)
}
ambiguityErr := tracker.RecordSQLiteUpdate(second, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal})
// Then
if !errors.Is(abortErr, ErrSQLiteHoldAborted) {
t.Fatalf("abort waiter error = %v", abortErr)
}
if !errors.Is(ambiguityErr, ErrSQLiteHoldAmbiguousCandidate) {
t.Fatalf("ambiguity error = %v", ambiguityErr)
}
if _, lookupErr := tracker.BeginSQLiteCommitFinalization(first); !errors.Is(lookupErr, ErrSQLiteHoldAmbiguousCandidate) {
t.Fatalf("ambiguity lookup error = %v", lookupErr)
}
}
func TestSQLiteHoldTrackerCommitFinalizationReturnsSelectedAbortCause(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(64)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
session, err := tracker.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
if err := tracker.AbortSQLiteHold(session); err != nil {
t.Fatal(err)
}
// When
_, selectedErr := tracker.BeginSQLiteCommitFinalization(transaction)
_, unselectedErr := tracker.BeginSQLiteCommitFinalization(sqliteHoldTestTransaction(65))
// Then
if !errors.Is(selectedErr, ErrSQLiteHoldAborted) {
t.Fatalf("selected aborted lookup error = %v", selectedErr)
}
if !errors.Is(unselectedErr, ErrSQLiteHoldNotSelected) {
t.Fatalf("unselected lookup error = %v", unselectedErr)
}
}
func TestSQLiteHoldTrackerCommitFinalizationPreservesCauseAcrossLaterLifecycle(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
first := sqliteHoldTestTransaction(70)
if err := tracker.BeginSQLiteTransaction(first); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, first)
firstSession, err := tracker.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
if err := tracker.AbortSQLiteHold(firstSession); err != nil {
t.Fatal(err)
}
second := sqliteHoldTestTransaction(71)
if err := tracker.BeginSQLiteTransaction(second); err != nil {
t.Fatal(err)
}
secondJournal := SQLiteJournalIdentity{Inode: 72}
if err := tracker.RecordSQLiteUpdate(second, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: secondJournal}); err != nil {
t.Fatal(err)
}
secondSession, err := tracker.ArmSQLiteHold(secondJournal)
if err != nil {
t.Fatal(err)
}
if err := tracker.AbortSQLiteHold(secondSession); err != nil {
t.Fatal(err)
}
// When
_, lookupErr := tracker.BeginSQLiteCommitFinalization(first)
// Then
if !errors.Is(lookupErr, ErrSQLiteHoldAborted) {
t.Fatalf("first lifecycle abort cause after second lifecycle = %v", lookupErr)
}
}
func TestSQLiteHoldTrackerWaitSelectedReturnsAbortWhenSelectedTransactionFinishes(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(68)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
session, err := tracker.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
result := make(chan error, 1)
go func() {
_, waitErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing)
result <- waitErr
}()
// When
if err := tracker.FinishSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
// Then
if waitErr := <-result; !errors.Is(waitErr, ErrSQLiteHoldAborted) {
t.Fatalf("finished selected transaction wait error = %v", waitErr)
}
}
func TestSQLiteHoldTrackerWaitTreatsFinishedReleasedSessionAsSelectedAndFinalizing(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(69)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
recordSQLiteHoldTestUpdate(t, tracker, transaction)
session, err := tracker.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
if _, err := tracker.BeginSQLiteFinalization(transaction); err != nil {
t.Fatal(err)
}
if err := tracker.ReleaseSQLiteHold(session); err != nil {
t.Fatal(err)
}
if err := tracker.FinishSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
// When
selected, selectedErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitSelected)
finalizing, finalizingErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing)
// Then
if selectedErr != nil || !selected.Released || !selected.Selected {
t.Fatalf("selected=%+v err=%v", selected, selectedErr)
}
if finalizingErr != nil || !finalizing.Released || !finalizing.Finalizing {
t.Fatalf("finalizing=%+v err=%v", finalizing, finalizingErr)
}
}
@@ -0,0 +1,215 @@
//go:build agentcompat && linux
package singleton
import (
"errors"
"testing"
)
func sqliteHoldTestWrite(operation SQLiteOperation, table string, stackHash uint64, journal SQLiteJournalIdentity) SQLiteWriteObservation {
return SQLiteWriteObservation{
Origin: SQLiteExecutionOrigin{Operation: operation, Table: table, StackHash: stackHash, FirstNezhaFrame: "singleton.write"},
Update: SQLiteUpdateObservation{Operation: operation, Table: table, Journal: journal},
}
}
func TestSQLiteHoldTrackerRecordsAtomicFirstWriteAttribution(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(41)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
session, err := tracker.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
// When
err = tracker.RecordSQLiteWrite(transaction, sqliteHoldTestWrite(SQLiteOperationInsert, "settings", 51, sqliteHoldTestJournal))
snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session)
// Then
if err != nil || snapshotErr != nil {
t.Fatalf("write=%v snapshot=%v", err, snapshotErr)
}
if !snapshot.Selected || snapshot.Transaction != transaction || snapshot.Operation != SQLiteOperationInsert ||
snapshot.Table != "settings" || snapshot.StackHash != 51 || snapshot.FirstNezhaFrame != "singleton.write" ||
snapshot.Journal != sqliteHoldTestJournal {
t.Fatalf("unexpected atomic snapshot: %+v", snapshot)
}
}
func TestSQLiteHoldTrackerPreservesFirstWriteForRepeatedSameTransaction(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(42)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
session, err := tracker.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
if err := tracker.RecordSQLiteWrite(transaction, sqliteHoldTestWrite(SQLiteOperationUpdate, "first", 52, sqliteHoldTestJournal)); err != nil {
t.Fatal(err)
}
// When
err = tracker.RecordSQLiteWrite(transaction, sqliteHoldTestWrite(SQLiteOperationDelete, "later", 53, SQLiteJournalIdentity{Inode: 98}))
snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session)
// Then
if err != nil || snapshotErr != nil {
t.Fatalf("repeat=%v snapshot=%v", err, snapshotErr)
}
if snapshot.Operation != SQLiteOperationUpdate || snapshot.Table != "first" || snapshot.StackHash != 52 || snapshot.Journal != sqliteHoldTestJournal {
t.Fatalf("repeated callback changed first write: %+v", snapshot)
}
}
func TestSQLiteHoldTrackerAbortsDifferentAtomicWriterBeforeRelease(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
session, err := tracker.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
first, second := sqliteHoldTestTransaction(43), sqliteHoldTestTransaction(44)
for _, transaction := range []SQLiteTransaction{first, second} {
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
}
if err := tracker.RecordSQLiteWrite(first, sqliteHoldTestWrite(SQLiteOperationInsert, "first", 54, sqliteHoldTestJournal)); err != nil {
t.Fatal(err)
}
// When
err = tracker.RecordSQLiteWrite(second, sqliteHoldTestWrite(SQLiteOperationDelete, "second", 55, sqliteHoldTestJournal))
// Then
if !errors.Is(err, ErrSQLiteHoldAmbiguousCandidate) {
t.Fatalf("expected ambiguity, got %v", err)
}
if _, err := tracker.SQLiteHoldSnapshot(session); !errors.Is(err, ErrSQLiteHoldStaleSession) {
t.Fatalf("expected stale aborted session, got %v", err)
}
}
func TestSQLiteHoldTrackerKnownJournalMakesFirstMismatchedWriteIneligible(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(45)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal)
if err != nil {
t.Fatal(err)
}
if err := tracker.RecordSQLiteWrite(transaction, sqliteHoldTestWrite(SQLiteOperationUpdate, "wrong", 56, SQLiteJournalIdentity{Inode: 97})); err != nil {
t.Fatal(err)
}
// When
err = tracker.RecordSQLiteWrite(transaction, sqliteHoldTestWrite(SQLiteOperationUpdate, "matching", 57, sqliteHoldTestJournal))
snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session)
// Then
if err != nil || snapshotErr != nil {
t.Fatalf("later write=%v snapshot=%v", err, snapshotErr)
}
if snapshot.Selected || snapshot.Journal != sqliteHoldTestJournal {
t.Fatalf("mismatched first write selected session: %+v", snapshot)
}
}
func TestSQLiteHoldTrackerRejectsAtomicWriteAfterFinalization(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(46)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
session, err := tracker.ArmNextSQLiteHold()
if err != nil {
t.Fatal(err)
}
first := sqliteHoldTestWrite(SQLiteOperationInsert, "first", 58, sqliteHoldTestJournal)
if err := tracker.RecordSQLiteWrite(transaction, first); err != nil {
t.Fatal(err)
}
if _, err := tracker.BeginSQLiteFinalization(transaction); err != nil {
t.Fatal(err)
}
// When
err = tracker.RecordSQLiteWrite(transaction, sqliteHoldTestWrite(SQLiteOperationDelete, "later", 59, SQLiteJournalIdentity{Inode: 96}))
snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session)
// Then
if !errors.Is(err, ErrSQLiteHoldFinalizationStarted) || snapshotErr != nil {
t.Fatalf("write=%v snapshot=%v", err, snapshotErr)
}
if snapshot.Operation != first.Update.Operation || snapshot.Table != first.Update.Table || snapshot.StackHash != first.Origin.StackHash || snapshot.Journal != first.Update.Journal {
t.Fatalf("finalizing atomic write changed snapshot: %+v", snapshot)
}
}
func TestSQLiteHoldTrackerRejectsLegacyWritesAfterAtomicFirstWrite(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(47)
if err := tracker.BeginSQLiteTransaction(transaction); err != nil {
t.Fatal(err)
}
journalA := sqliteHoldTestJournal
journalB := SQLiteJournalIdentity{Inode: 95}
session, err := tracker.ArmSQLiteHold(journalA)
if err != nil {
t.Fatal(err)
}
first := sqliteHoldTestWrite(SQLiteOperationInsert, "atomic", 60, journalB)
if err := tracker.RecordSQLiteWrite(transaction, first); err != nil {
t.Fatal(err)
}
// When
executionErr := tracker.RecordSQLiteExecution(transaction, SQLiteExecutionOrigin{
Operation: SQLiteOperationUpdate, Table: "legacy", StackHash: 61, FirstNezhaFrame: "singleton.legacy",
})
updateErr := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{
Operation: SQLiteOperationUpdate, Table: "legacy", Journal: journalA,
})
snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session)
// Then
if !errors.Is(executionErr, ErrSQLiteHoldAtomicWriteRecorded) || !errors.Is(updateErr, ErrSQLiteHoldAtomicWriteRecorded) {
t.Fatalf("execution=%v update=%v", executionErr, updateErr)
}
if snapshotErr != nil {
t.Fatal(snapshotErr)
}
if snapshot.Selected || snapshot.Operation != "" || snapshot.Journal != journalA {
t.Fatalf("legacy write selected mismatched atomic transaction: %+v", snapshot)
}
if _, err := tracker.BeginSQLiteFinalization(transaction); !errors.Is(err, ErrSQLiteHoldNotSelected) {
t.Fatalf("expected mismatched atomic transaction to remain unselected, got %v", err)
}
}
func TestSQLiteHoldTrackerRejectsAtomicWriteForUnregisteredTransaction(t *testing.T) {
// Given
tracker := NewSQLiteHoldTracker()
transaction := sqliteHoldTestTransaction(48)
// When
err := tracker.RecordSQLiteWrite(transaction, sqliteHoldTestWrite(SQLiteOperationInsert, "settings", 62, sqliteHoldTestJournal))
// Then
if !errors.Is(err, ErrSQLiteHoldNotSelected) {
t.Fatalf("expected unregistered atomic write rejection, got %v", err)
}
}
+70
View File
@@ -0,0 +1,70 @@
package singleton
import "github.com/nezhahq/nezha/model"
// NewEmptyServerClassForTest 构造一个不依赖 DB 的空 ServerClass,仅用于单测。
// 生产路径请用 NewServerClass。
func NewEmptyServerClassForTest() *ServerClass {
sc := &ServerClass{
class: class[uint64, *model.Server]{
list: make(map[uint64]*model.Server),
},
uuidToID: make(map[string]uint64),
}
model.OwnerServerIDsLookup = sc.ownerServerIDs
model.AllServerIDsLookup = sc.allServerIDs
model.OwnerIsAdminLookup = ownerIsAdmin
return sc
}
// InsertForTest 把一个 server 直接塞进内存表与排序快照,跳过 DB & InitServer 逻辑。
// 调用方需保证 server.ID 已经设置。
func (c *ServerClass) InsertForTest(s *model.Server) {
c.lockLifecycleWrite()
defer c.unlockLifecycleWrite()
c.listMu.Lock()
c.list[s.ID] = s
if s.UUID != "" {
c.uuidToID[s.UUID] = s.ID
}
c.listMu.Unlock()
c.sortList()
}
// NewEmptyDDNSClassForTest 构造一个不依赖 DB 的空 DDNSClass,仅用于单测。
func NewEmptyDDNSClassForTest() *DDNSClass {
return &DDNSClass{
class: class[uint64, *model.DDNSProfile]{
list: make(map[uint64]*model.DDNSProfile),
},
}
}
// InsertForTest 把一个 DDNS profile 直接塞进内存表,跳过 DB。
func (c *DDNSClass) InsertForTest(p *model.DDNSProfile) {
c.listMu.Lock()
c.list[p.ID] = p
c.listMu.Unlock()
c.sortList()
}
// NewEmptyNotificationClassForTest 构造空 NotificationClass。
func NewEmptyNotificationClassForTest() *NotificationClass {
return &NotificationClass{
class: class[uint64, *model.Notification]{
list: make(map[uint64]*model.Notification),
},
groupToIDList: make(map[uint64]map[uint64]*model.Notification),
idToGroupList: make(map[uint64]map[uint64]struct{}),
groupList: make(map[uint64]string),
}
}
// InsertForTest 把一个 Notification 直接塞进内存表。
func (c *NotificationClass) InsertForTest(n *model.Notification) {
c.listMu.Lock()
c.list[n.ID] = n
c.listMu.Unlock()
c.sortList()
}
+39 -2
View File
@@ -23,7 +23,10 @@ func initUser() {
var users []model.User
DB.Find(&users)
// for backward compatibility
// Backward compatibility for pre-user-scoped Agents. AgentSecretKey is a
// deployment-wide migration/master credential, so user 0 is intentionally
// not tenant-scoped. Do not remove this mapping until every legacy Agent has
// rotated to a per-user/per-Agent credential; doing so would disconnect them.
UserInfoMap[0] = model.UserInfo{
Role: model.RoleAdmin,
AgentSecret: Conf.AgentSecretKey,
@@ -40,10 +43,34 @@ func initUser() {
UserInfoMap[u.ID] = model.UserInfo{
Role: u.Role,
Username: u.Username,
AgentSecret: u.AgentSecret,
}
AgentSecretToUserId[u.AgentSecret] = u.ID
}
model.ServerOwnerLookup = lookupServerOwner
}
// lookupServerOwner resolves Server.UserID into a display-ready owner
// record for model.Server.MarshalJSON. uid=0 is the legacy global agent
// secret (a pseudo-owner with no User row) and intentionally returns
// ok=false with no username; the frontend renders it as "Global". Other
// uids return ok=false when the user has been deleted, so the JSON still
// carries the bare id and the frontend can render an "Unknown (#<uid>)"
// placeholder. RLock is required because OnUserUpdate / OnUserDelete may
// mutate UserInfoMap concurrently with serialization.
func lookupServerOwner(uid uint64) (model.ServerOwnerInfo, bool) {
if uid == 0 {
return model.ServerOwnerInfo{}, false
}
UserLock.RLock()
info, ok := UserInfoMap[uid]
UserLock.RUnlock()
if !ok {
return model.ServerOwnerInfo{}, false
}
return model.ServerOwnerInfo{ID: uid, Username: info.Username}, true
}
func OnUserUpdate(u *model.User) {
@@ -56,6 +83,7 @@ func OnUserUpdate(u *model.User) {
UserInfoMap[u.ID] = model.UserInfo{
Role: u.Role,
Username: u.Username,
AgentSecret: u.AgentSecret,
}
AgentSecretToUserId[u.AgentSecret] = u.ID
@@ -69,6 +97,10 @@ func OnUserDelete(id []uint64, errorFunc func(string, ...any) error) error {
return Localizer.ErrorT("user id not specified")
}
if ServerTransferShared != nil {
ServerTransferShared.OnUsersDeleted(id)
}
var (
cron, server bool
crons, servers []uint64
@@ -101,7 +133,7 @@ func OnUserDelete(id []uint64, errorFunc func(string, ...any) error) error {
return err
}
if err := tx.Where("id IN (?)", id).Delete(&model.User{}).Error; err != nil {
if err := tx.Where("id = ?", uid).Delete(&model.User{}).Error; err != nil {
return err
}
return nil
@@ -127,6 +159,11 @@ func OnUserDelete(id []uint64, errorFunc func(string, ...any) error) error {
}
}
AlertsLock.Unlock()
// Cancel pending transfers before ServerShared drops the
// in-memory entry: same ordering rationale as batchDeleteServer.
if ServerTransferShared != nil {
ServerTransferShared.OnServersDeleted(servers)
}
ServerShared.Delete(servers)
}
+90
View File
@@ -0,0 +1,90 @@
package singleton
import (
"testing"
"github.com/stretchr/testify/require"
"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/i18n"
)
func setupOnUserDeleteFixture(t *testing.T) (*ServerTransferClass, func()) {
t.Helper()
c, transferCleanup := setupTransferFixture(t)
require.NoError(t, DB.AutoMigrate(&model.Cron{}, &model.Transfer{}, &model.ServerGroupServer{}))
originalCronShared := CronShared
CronShared = &CronClass{
class: class[uint64, *model.Cron]{list: map[uint64]*model.Cron{}},
}
originalLocalizer := Localizer
Localizer = i18n.NewLocalizer("zh_CN", domain, "translations", i18n.Translations)
cleanup := func() {
Localizer = originalLocalizer
CronShared = originalCronShared
transferCleanup()
}
return c, cleanup
}
func TestOnUserDeleteCancelsPendingTransfersAwayFromDeletedUser(t *testing.T) {
c, cleanup := setupOnUserDeleteFixture(t)
defer cleanup()
const fromUser = uint64(100)
const toUser = uint64(200)
const serverID = uint64(1)
seedServerForTransfer(t, serverID, fromUser)
require.NoError(t, DB.AutoMigrate(&model.User{}))
require.NoError(t, DB.Create(&model.User{
Common: model.Common{ID: fromUser},
Username: "alice",
AgentSecret: "alice-secret",
}).Error)
require.NoError(t, DB.Create(&model.User{
Common: model.Common{ID: toUser},
Username: "bob",
AgentSecret: "bob-secret",
}).Error)
UserLock.Lock()
UserInfoMap[fromUser] = model.UserInfo{Role: model.RoleMember, AgentSecret: "alice-secret"}
UserInfoMap[toUser] = model.UserInfo{Role: model.RoleMember, AgentSecret: "bob-secret"}
UserLock.Unlock()
tr := initiateAndRegister(t, c, serverID, fromUser, toUser, fromUser)
require.True(t, c.HasPending(serverID), "precondition: pending transfer published")
srv, ok := ServerShared.Get(serverID)
require.True(t, ok)
require.Equal(t, toUser, srv.GetUserID(), "precondition: pending transfer flipped owner to ToUserID")
require.NoError(t, OnUserDelete([]uint64{fromUser}, func(format string, args ...any) error {
return nil
}))
if c.HasPending(serverID) {
t.Fatal("OnUserDelete on the transfer FromUserID must terminate the pending transfer so a later Cancel/Fail/Timeout cannot revert ownership to the deleted user")
}
if srv, ok := ServerShared.Get(serverID); ok {
require.NotEqual(t, fromUser, srv.GetUserID(),
"server owner must not be reverted to the deleted FromUserID; got owner=%d", srv.GetUserID())
}
if _, err := c.Cancel(tr.ID); err == nil {
var refreshed model.ServerTransfer
if err := DB.First(&refreshed, tr.ID).Error; err == nil {
require.NotEqual(t, model.ServerTransferStatusPending, refreshed.Status,
"after OnUserDelete a subsequent Cancel must not leave the transfer Pending")
if srv, ok := ServerShared.Get(serverID); ok {
require.NotEqual(t, fromUser, srv.GetUserID(),
"a late Cancel against the terminated transfer must not revert server.UserID to the deleted FromUserID; got owner=%d", srv.GetUserID())
}
}
}
}