diff --git a/database/db.go b/database/db.go index e0c50551..62cf1f11 100644 --- a/database/db.go +++ b/database/db.go @@ -40,6 +40,8 @@ func initModels() error { &model.InboundClientIps{}, &xray.ClientTraffic{}, &model.HistoryOfSeeders{}, + &model.SharedState{}, + &model.NodeState{}, } for _, model := range models { if err := db.AutoMigrate(model); err != nil { @@ -47,6 +49,9 @@ func initModels() error { return err } } + if err := seedSharedAccountsVersion(db); err != nil { + return err + } return nil } diff --git a/database/db_test.go b/database/db_test.go index a51b5c64..a64cb501 100644 --- a/database/db_test.go +++ b/database/db_test.go @@ -256,3 +256,67 @@ func TestSettingKey_IsUnique(t *testing.T) { t.Fatal("expected duplicate setting key insert to fail") } } + +func TestInitDB_CreatesSharedMetadataTables(t *testing.T) { + setupTestDB(t) + + for _, table := range []string{"shared_states", "node_states"} { + var count int64 + if err := db.Table(table).Count(&count).Error; err != nil { + t.Fatalf("table %s should exist: %v", table, err) + } + } +} + +func TestBumpSharedAccountsVersion(t *testing.T) { + setupTestDB(t) + + version, err := GetSharedAccountsVersion(GetDB()) + if err != nil { + t.Fatalf("GetSharedAccountsVersion error: %v", err) + } + if version != 0 { + t.Fatalf("expected seeded version 0, got %d", version) + } + + tx := GetDB().Begin() + if err := BumpSharedAccountsVersion(tx); err != nil { + t.Fatalf("BumpSharedAccountsVersion error: %v", err) + } + if err := tx.Commit().Error; err != nil { + t.Fatalf("Commit error: %v", err) + } + + version, err = GetSharedAccountsVersion(GetDB()) + if err != nil { + t.Fatalf("GetSharedAccountsVersion error: %v", err) + } + if version != 1 { + t.Fatalf("expected bumped version 1, got %d", version) + } +} + +func TestUpsertNodeState(t *testing.T) { + setupTestDB(t) + + state := &model.NodeState{ + NodeID: "worker-1", + NodeRole: "worker", + LastSeenVersion: 7, + LastError: "dial tcp timeout", + } + if err := UpsertNodeState(GetDB(), state); err != nil { + t.Fatalf("UpsertNodeState error: %v", err) + } + + var stored model.NodeState + if err := GetDB().First(&stored, "node_id = ?", "worker-1").Error; err != nil { + t.Fatalf("lookup node state failed: %v", err) + } + if stored.LastSeenVersion != 7 { + t.Fatalf("expected last seen version 7, got %d", stored.LastSeenVersion) + } + if stored.LastError != "dial tcp timeout" { + t.Fatalf("expected last error to round-trip, got %q", stored.LastError) + } +} diff --git a/database/model/node_state.go b/database/model/node_state.go new file mode 100644 index 00000000..79476a2f --- /dev/null +++ b/database/model/node_state.go @@ -0,0 +1,11 @@ +package model + +type NodeState struct { + NodeID string `json:"nodeId" gorm:"primaryKey"` + NodeRole string `json:"nodeRole" gorm:"not null"` + LastSyncAt int64 `json:"lastSyncAt"` + LastHeartbeatAt int64 `json:"lastHeartbeatAt"` + LastSeenVersion int64 `json:"lastSeenVersion"` + LastError string `json:"lastError"` + UpdatedAt int64 `json:"updatedAt"` +} diff --git a/database/model/shared_state.go b/database/model/shared_state.go new file mode 100644 index 00000000..bb58fcc6 --- /dev/null +++ b/database/model/shared_state.go @@ -0,0 +1,7 @@ +package model + +type SharedState struct { + Key string `json:"key" gorm:"primaryKey"` + Version int64 `json:"version" gorm:"not null;default:0"` + UpdatedAt int64 `json:"updatedAt"` +} diff --git a/database/shared_state.go b/database/shared_state.go new file mode 100644 index 00000000..1b4d5859 --- /dev/null +++ b/database/shared_state.go @@ -0,0 +1,50 @@ +package database + +import ( + "time" + + "github.com/mhsanaei/3x-ui/v2/database/model" + "gorm.io/gorm" +) + +const SharedAccountsVersionKey = "shared_accounts_version" + +func txOrDB(tx *gorm.DB) *gorm.DB { + if tx != nil { + return tx + } + return GetDB() +} + +func seedSharedAccountsVersion(tx *gorm.DB) error { + return txOrDB(tx).FirstOrCreate( + &model.SharedState{}, + &model.SharedState{ + Key: SharedAccountsVersionKey, + Version: 0, + UpdatedAt: time.Now().Unix(), + }, + ).Error +} + +func GetSharedAccountsVersion(tx *gorm.DB) (int64, error) { + state := &model.SharedState{} + if err := txOrDB(tx).First(state, "key = ?", SharedAccountsVersionKey).Error; err != nil { + return 0, err + } + return state.Version, nil +} + +func BumpSharedAccountsVersion(tx *gorm.DB) error { + return txOrDB(tx).Model(&model.SharedState{}). + Where("key = ?", SharedAccountsVersionKey). + Updates(map[string]any{ + "version": gorm.Expr("version + 1"), + "updated_at": time.Now().Unix(), + }).Error +} + +func UpsertNodeState(tx *gorm.DB, state *model.NodeState) error { + state.UpdatedAt = time.Now().Unix() + return txOrDB(tx).Save(state).Error +}