diff --git a/server/internal/auth/locked_test.go b/server/internal/auth/locked_test.go new file mode 100644 index 0000000..498642b --- /dev/null +++ b/server/internal/auth/locked_test.go @@ -0,0 +1,34 @@ +package auth + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" +) + +func TestRefuseLocked(t *testing.T) { + gin.SetMode(gin.TestMode) + prev := instanceLocked + t.Cleanup(func() { instanceLocked = prev }) + instanceLocked = func(id string) bool { return id == "locked" } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + if !refuseLocked(c, "locked") { + t.Fatal("a locked instance must be refused") + } + if w.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want 401", w.Code) + } + + w = httptest.NewRecorder() + c, _ = gin.CreateTestContext(w) + if refuseLocked(c, "open") { + t.Fatal("an unlocked instance must pass") + } + if c.IsAborted() { + t.Fatal("an unlocked instance must not abort") + } +} diff --git a/server/internal/auth/middleware.go b/server/internal/auth/middleware.go index ecd4840..ebc4af8 100644 --- a/server/internal/auth/middleware.go +++ b/server/internal/auth/middleware.go @@ -12,6 +12,20 @@ import ( const ctxSessionKey = "km_session" +// instanceLocked is a variable so tests can stub the Mongo-backed check. +var instanceLocked = services.InstanceLocked + +// refuseLocked answers 401 for a session or token on an instance Vantage HQ has +// locked under a dispute. Same body as an expired session: a locked instance +// is not announced as locked to whoever holds a credential for it. +func refuseLocked(c *gin.Context, instanceID string) bool { + if !instanceLocked(instanceID) { + return false + } + c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "session expired"}) + return true +} + func GetSessionFromContext(c *gin.Context) *Session { v, _ := c.Get(ctxSessionKey) sess, _ := v.(*Session) @@ -81,6 +95,13 @@ func Middleware() gin.HandlerFunc { return } + // Explicit, because the host guard below cannot do this: the resolver + // hides a locked instance, so its host resolves to nothing and that + // guard is skipped rather than tripped. + if refuseLocked(c, sess.InstanceID) { + return + } + c.Set(ctxSessionKey, sess) // The host guard applies to both credential kinds. A token carries an diff --git a/server/internal/services/instancelock.go b/server/internal/services/instancelock.go new file mode 100644 index 0000000..4a77831 --- /dev/null +++ b/server/internal/services/instancelock.go @@ -0,0 +1,73 @@ +package services + +import ( + "context" + "log" + "sync" + "time" + + "gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/db" + "go.mongodb.org/mongo-driver/v2/bson" +) + +// instanceLockTTL matches the host resolver's cache, so a lock and an unlock +// both take effect within the same minute everywhere. +const instanceLockTTL = 60 * time.Second + +type lockEntry struct { + locked bool + at time.Time +} + +var ( + lockMu sync.Mutex + lockCache = map[string]lockEntry{} + + // lockLookup and lockNow are variables so tests can replace them. + lockLookup = func(ctx context.Context, instanceID string) (bool, error) { + n, err := db.Col("instances").CountDocuments(ctx, bson.M{ + "instance_id": instanceID, + "locked_at": bson.M{"$exists": true}, + }) + return n > 0, err + } + lockNow = time.Now +) + +// InstanceLocked reports whether Vantage HQ has locked this instance under an +// account dispute. HQ writes instances.locked_at through its cloudprov package; +// nothing on this side ever sets or clears it. +// +// A read error answers false and is not cached. A database that cannot answer +// this cannot serve the request that asked either, so failing open here costs +// nothing that failing closed would save. +func InstanceLocked(instanceID string) bool { + if instanceID == "" { + return false + } + lockMu.Lock() + if e, ok := lockCache[instanceID]; ok && lockNow().Sub(e.at) < instanceLockTTL { + lockMu.Unlock() + return e.locked + } + lockMu.Unlock() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + locked, err := lockLookup(ctx, instanceID) + if err != nil { + log.Printf("instance lock: read %s: %v", instanceID, err) + return false + } + lockMu.Lock() + lockCache[instanceID] = lockEntry{locked: locked, at: lockNow()} + lockMu.Unlock() + return locked +} + +// unlockedFilter narrows an instance query to instances HQ has not locked. +// HQ clears the field with $unset, so absence is the whole test. +func unlockedFilter(f bson.M) bson.M { + f["locked_at"] = bson.M{"$exists": false} + return f +} diff --git a/server/internal/services/instancelock_test.go b/server/internal/services/instancelock_test.go new file mode 100644 index 0000000..bd8bf56 --- /dev/null +++ b/server/internal/services/instancelock_test.go @@ -0,0 +1,86 @@ +package services + +import ( + "context" + "errors" + "testing" + "time" + + "go.mongodb.org/mongo-driver/v2/bson" +) + +// stubLock replaces the Mongo lookup and the clock, and empties the cache, for +// one test. +func stubLock(t *testing.T, lookup func(context.Context, string) (bool, error), at *time.Time) { + t.Helper() + prevLookup, prevNow := lockLookup, lockNow + lockLookup = lookup + lockNow = func() time.Time { return *at } + lockMu.Lock() + lockCache = map[string]lockEntry{} + lockMu.Unlock() + t.Cleanup(func() { + lockLookup, lockNow = prevLookup, prevNow + lockMu.Lock() + lockCache = map[string]lockEntry{} + lockMu.Unlock() + }) +} + +func TestInstanceLockedCachesForTTL(t *testing.T) { + calls := 0 + at := time.Date(2026, 9, 10, 10, 0, 0, 0, time.UTC) + stubLock(t, func(context.Context, string) (bool, error) { calls++; return true, nil }, &at) + + if !InstanceLocked("i1") { + t.Fatal("want locked") + } + InstanceLocked("i1") + if calls != 1 { + t.Fatalf("lookups = %d within the TTL, want 1", calls) + } + at = at.Add(instanceLockTTL + time.Second) + InstanceLocked("i1") + if calls != 2 { + t.Fatalf("lookups = %d after the TTL, want 2", calls) + } +} + +func TestInstanceLockedReadErrorIsNotCached(t *testing.T) { + calls := 0 + at := time.Date(2026, 9, 10, 10, 0, 0, 0, time.UTC) + stubLock(t, func(context.Context, string) (bool, error) { + calls++ + return false, errors.New("mongo down") + }, &at) + + if InstanceLocked("i1") { + t.Fatal("a read error must answer false") + } + InstanceLocked("i1") + if calls != 2 { + t.Fatalf("lookups = %d, want 2: an error must not be cached", calls) + } +} + +func TestInstanceLockedEmptyID(t *testing.T) { + at := time.Now() + stubLock(t, func(context.Context, string) (bool, error) { + t.Fatal("no lookup for an empty id") + return false, nil + }, &at) + if InstanceLocked("") { + t.Fatal("empty id is never locked") + } +} + +func TestUnlockedFilter(t *testing.T) { + got := unlockedFilter(bson.M{"slug": "acme"}) + want := bson.M{"slug": "acme", "locked_at": bson.M{"$exists": false}} + if len(got) != 2 || got["slug"] != "acme" { + t.Fatalf("got %v, want %v", got, want) + } + if cond, ok := got["locked_at"].(bson.M); !ok || cond["$exists"] != false { + t.Fatalf("got %v, want %v", got, want) + } +} diff --git a/server/internal/services/instances.go b/server/internal/services/instances.go index dc34119..93f9457 100644 --- a/server/internal/services/instances.go +++ b/server/internal/services/instances.go @@ -28,7 +28,7 @@ func GetInstanceBySlug(slug string) (*models.Instance, error) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() var o models.Instance - err := db.Col("instances").FindOne(ctx, bson.M{"slug": slug}).Decode(&o) + err := db.Col("instances").FindOne(ctx, unlockedFilter(bson.M{"slug": slug})).Decode(&o) if err != nil { return nil, err } diff --git a/server/internal/services/servers.go b/server/internal/services/servers.go index b71b976..05c6389 100644 --- a/server/internal/services/servers.go +++ b/server/internal/services/servers.go @@ -201,6 +201,12 @@ func ValidateAgentToken(serverID, agentToken string) (*models.Server, error) { if s.InstanceID == "" { return nil, fmt.Errorf("server %s has no org", serverID) } + + // An agent of a locked instance is refused exactly like a bad token, so it + // keeps retrying with backoff and reconnects on its own if HQ restores it. + if InstanceLocked(s.InstanceID) { + return nil, fmt.Errorf("invalid agent token") + } return &s, nil }