Fix token (#2653)

* fix: kick token * fix: kick token * fix: change config
openimsdk · Sep 25, 2024 · f6364a4 · f6364a4
1 parent 031c1cd
commit f6364a4
Show file tree

Hide file tree

Showing 6 changed files with 48 additions and 13 deletions.
diff --git a/config/openim-msggateway.yml b/config/openim-msggateway.yml
@@ -23,5 +23,4 @@ longConnSvr:
   # WebSocket connection handshake timeout in seconds
   websocketTimeout: 10
 
-# 1: For Android, iOS, Windows, Mac, and web platforms, only one instance can be online at a time
-multiLoginPolicy: 1
+
diff --git a/config/share.yml b/config/share.yml
@@ -12,3 +12,5 @@ rpcRegisterName:
 
 imAdminUserID: [ imAdmin ]
 
+# 1: For Android, iOS, Windows, Mac, and web platforms, only one instance can be online at a time
+multiLoginPolicy: 1
diff --git a/internal/msggateway/ws_server.go b/internal/msggateway/ws_server.go
@@ -321,7 +321,7 @@ func (ws *WsServer) KickUserConn(client *Client) error {
 }
 
 func (ws *WsServer) multiTerminalLoginChecker(clientOK bool, oldClients []*Client, newClient *Client) {
-	switch ws.msgGatewayConfig.MsgGateway.MultiLoginPolicy {
+	switch ws.msgGatewayConfig.Share.MultiLoginPolicy {
 	case constant.DefalutNotKick:
 	case constant.PCAndOther:
 		if constant.PlatformIDToClass(newClient.PlatformID) == constant.TerminalPC {

diff --git a/internal/rpc/auth/auth.go b/internal/rpc/auth/auth.go
@@ -64,6 +64,7 @@ func Start(ctx context.Context, config *Config, client discovery.SvcDiscoveryReg
 			redis2.NewTokenCacheModel(rdb, config.RpcConfig.TokenPolicy.Expire),
 			config.Share.Secret,
 			config.RpcConfig.TokenPolicy.Expire,
+			config.Share.MultiLoginPolicy,
 		),
 		config: config,
 	})

diff --git a/pkg/common/config/config.go b/pkg/common/config/config.go
@@ -185,7 +185,6 @@ type MsgGateway struct {
 		WebsocketMaxMsgLen  int   `mapstructure:"websocketMaxMsgLen"`
 		WebsocketTimeout    int   `mapstructure:"websocketTimeout"`
 	} `mapstructure:"longConnSvr"`
-	MultiLoginPolicy int `mapstructure:"multiLoginPolicy"`
 }
 
 type MsgTransfer struct {
@@ -358,9 +357,10 @@ type AfterConfig struct {
 }
 
 type Share struct {
-	Secret          string          `mapstructure:"secret"`
-	RpcRegisterName RpcRegisterName `mapstructure:"rpcRegisterName"`
-	IMAdminUserID   []string        `mapstructure:"imAdminUserID"`
+	Secret           string          `mapstructure:"secret"`
+	RpcRegisterName  RpcRegisterName `mapstructure:"rpcRegisterName"`
+	IMAdminUserID    []string        `mapstructure:"imAdminUserID"`
+	MultiLoginPolicy int             `mapstructure:"multiLoginPolicy"`
 }
 type RpcRegisterName struct {
 	User           string `mapstructure:"user"`

diff --git a/pkg/common/storage/controller/auth.go b/pkg/common/storage/controller/auth.go
@@ -35,13 +35,14 @@ type AuthDatabase interface {
 }
 
 type authDatabase struct {
-	cache        cache.TokenModel
-	accessSecret string
-	accessExpire int64
+	cache            cache.TokenModel
+	accessSecret     string
+	accessExpire     int64
+	multiLoginPolicy int
 }
 
-func NewAuthDatabase(cache cache.TokenModel, accessSecret string, accessExpire int64) AuthDatabase {
-	return &authDatabase{cache: cache, accessSecret: accessSecret, accessExpire: accessExpire}
+func NewAuthDatabase(cache cache.TokenModel, accessSecret string, accessExpire int64, policy int) AuthDatabase {
+	return &authDatabase{cache: cache, accessSecret: accessSecret, accessExpire: accessExpire, multiLoginPolicy: policy}
 }
 
 // If the result is empty.
@@ -55,15 +56,19 @@ func (a *authDatabase) SetTokenMapByUidPid(ctx context.Context, userID string, p
 
 // Create Token.
 func (a *authDatabase) CreateToken(ctx context.Context, userID string, platformID int) (string, error) {
+	// todo: get all platform token
 	tokens, err := a.cache.GetTokensWithoutError(ctx, userID, platformID)
 	if err != nil {
 		return "", err
 	}
 	var deleteTokenKey []string
+	var kickedTokenKey []string
 	for k, v := range tokens {
-		_, err = tokenverify.GetClaimFromToken(k, authverify.Secret(a.accessSecret))
+		t, err := tokenverify.GetClaimFromToken(k, authverify.Secret(a.accessSecret))
 		if err != nil || v != constant.NormalToken {
 			deleteTokenKey = append(deleteTokenKey, k)
+		} else if a.checkKickToken(ctx, platformID, t) {
+			kickedTokenKey = append(kickedTokenKey, k)
 		}
 	}
 	if len(deleteTokenKey) != 0 {
@@ -72,6 +77,14 @@ func (a *authDatabase) CreateToken(ctx context.Context, userID string, platformI
 			return "", err
 		}
 	}
+	if len(kickedTokenKey) != 0 {
+		for _, k := range kickedTokenKey {
+			err := a.cache.SetTokenFlagEx(ctx, userID, platformID, k, constant.KickedToken)
+			if err != nil {
+				return "", err
+			}
+		}
+	}
 
 	claims := tokenverify.BuildClaims(userID, platformID, a.accessExpire)
 	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
@@ -85,3 +98,23 @@ func (a *authDatabase) CreateToken(ctx context.Context, userID string, platformI
 	}
 	return tokenString, nil
 }
+
+func (a *authDatabase) checkKickToken(ctx context.Context, platformID int, token *tokenverify.Claims) bool {
+	switch a.multiLoginPolicy {
+	case constant.DefalutNotKick:
+		return false
+	case constant.PCAndOther:
+		if constant.PlatformIDToClass(platformID) == constant.TerminalPC ||
+			constant.PlatformIDToClass(token.PlatformID) == constant.TerminalPC {
+			return false
+		}
+		return true
+	case constant.AllLoginButSameTermKick:
+		if platformID == token.PlatformID {
+			return true
+		}
+		return false
+	default:
+		return false
+	}
+}