feat: implement precise token logout

server/store/group_message_runtime.go

@@ -2,6 +2,8 @@ package store
 
 import (
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
 	"encoding/json"
 	"fmt"
 	"sort"
@@ -78,6 +80,7 @@ type groupMessageRuntime interface {
 	AllowRate(tenantID types.TenantID, scope string, limit int, window time.Duration) (bool, int, time.Duration, error)
 	ClaimShortState(namespace, id string, ttl time.Duration) (bool, error)
 	ClaimShortStateValue(namespace, id, value string, ttl time.Duration) (bool, bool, error)
+	ShortStateExists(namespace, id string) (bool, error)
 }
 
 var groupMsgRuntime groupMessageRuntime
@@ -146,6 +149,26 @@ func ClaimShortStateValue(namespace, id, value string, ttl time.Duration) (bool,
 	return claimed, sameValue, true, err
 }
 
+func RevokeLoginToken(tenantID types.TenantID, token []byte, ttl time.Duration) (bool, error) {
+	if groupMsgRuntime == nil || !tenantID.IsValid() || len(token) == 0 || ttl <= 0 {
+		return false, nil
+	}
+	digest := sha256.Sum256(token)
+	id := strconv.FormatInt(int64(tenantID), 10) + ":" + hex.EncodeToString(digest[:])
+	_, err := groupMsgRuntime.ClaimShortState("revoked_login_token", id, ttl)
+	return true, err
+}
+
+func IsLoginTokenRevoked(tenantID types.TenantID, token []byte) (bool, bool, error) {
+	if groupMsgRuntime == nil || !tenantID.IsValid() || len(token) == 0 {
+		return false, false, nil
+	}
+	digest := sha256.Sum256(token)
+	id := strconv.FormatInt(int64(tenantID), 10) + ":" + hex.EncodeToString(digest[:])
+	revoked, err := groupMsgRuntime.ShortStateExists("revoked_login_token", id)
+	return revoked, true, err
+}
+
 func RefreshUserPresence(tenantID types.TenantID, uid types.Uid, sid, nodeID string) error {
 	if groupMsgRuntime == nil || !tenantID.IsValid() || uid.IsZero() || sid == "" {
 		return nil
@@ -881,6 +904,11 @@ func (r *redisGroupMessageRuntime) ClaimShortStateValue(namespace, id, value str
 	return false, existing == value, nil
 }
 
+func (r *redisGroupMessageRuntime) ShortStateExists(namespace, id string) (bool, error) {
+	count, err := r.client.Exists(context.Background(), r.shortStateKey(namespace, id)).Result()
+	return count > 0, err
+}
+
 func (r *redisGroupMessageRuntime) hotSeqIDs(tenantID types.TenantID, topic string, opts *types.QueryOpt, limit int) ([]int64, bool, error) {
 	if opts != nil && len(opts.IdRanges) > 0 {
 		seqs := expandRanges(opts.IdRanges, limit)

server/store/group_message_runtime_test.go

@@ -11,13 +11,15 @@ import (
 type testGroupMessageRuntime struct {
 	claims          map[string]*groupMessageUniqueRecord
 	seqs            map[string]int64
+	shortStates     map[string]bool
 	deletedPrefixes []string
 }
 
 func newTestGroupMessageRuntime() *testGroupMessageRuntime {
 	return &testGroupMessageRuntime{
-		claims: make(map[string]*groupMessageUniqueRecord),
-		seqs:   make(map[string]int64),
+		claims:      make(map[string]*groupMessageUniqueRecord),
+		seqs:        make(map[string]int64),
+		shortStates: make(map[string]bool),
 	}
 }
 
@@ -138,7 +140,12 @@ func (r *testGroupMessageRuntime) AllowRate(types.TenantID, string, int, time.Du
 	return true, 0, 0, nil
 }
 
-func (r *testGroupMessageRuntime) ClaimShortState(string, string, time.Duration) (bool, error) {
+func (r *testGroupMessageRuntime) ClaimShortState(namespace, id string, _ time.Duration) (bool, error) {
+	key := namespace + ":" + id
+	if r.shortStates[key] {
+		return false, nil
+	}
+	r.shortStates[key] = true
 	return true, nil
 }
 
@@ -146,6 +153,30 @@ func (r *testGroupMessageRuntime) ClaimShortStateValue(string, string, string, t
 	return true, true, nil
 }
 
+func (r *testGroupMessageRuntime) ShortStateExists(namespace, id string) (bool, error) {
+	return r.shortStates[namespace+":"+id], nil
+}
+
+func TestLoginTokenRevocation(t *testing.T) {
+	oldRuntime := groupMsgRuntime
+	groupMsgRuntime = newTestGroupMessageRuntime()
+	t.Cleanup(func() { groupMsgRuntime = oldRuntime })
+
+	token := []byte("current-device-token")
+	if revoked, enabled, err := IsLoginTokenRevoked(7, token); err != nil || !enabled || revoked {
+		t.Fatalf("unexpected initial state: revoked=%v enabled=%v err=%v", revoked, enabled, err)
+	}
+	if enabled, err := RevokeLoginToken(7, token, time.Hour); err != nil || !enabled {
+		t.Fatalf("revoke token: enabled=%v err=%v", enabled, err)
+	}
+	if revoked, enabled, err := IsLoginTokenRevoked(7, token); err != nil || !enabled || !revoked {
+		t.Fatalf("revoked token was accepted: revoked=%v enabled=%v err=%v", revoked, enabled, err)
+	}
+	if revoked, _, err := IsLoginTokenRevoked(7, []byte("other-device-token")); err != nil || revoked {
+		t.Fatalf("other device token was revoked: revoked=%v err=%v", revoked, err)
+	}
+}
+
 func TestMessageRuntimeSeqAndClientMidIdempotency(t *testing.T) {
 	oldRuntime := groupMsgRuntime
 	runtime := newTestGroupMessageRuntime()