Skip to content

Commit

Permalink
fix: admin token limit
Browse files Browse the repository at this point in the history
  • Loading branch information
icey-yu committed Nov 22, 2024
1 parent 15648cd commit c00d6c5
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 33 deletions.
10 changes: 8 additions & 2 deletions internal/rpc/auth/auth.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ package auth

import (
"context"
"errors"

"github.com/openimsdk/open-im-server/v3/pkg/common/config"
redis2 "github.com/openimsdk/open-im-server/v3/pkg/common/storage/cache/redis"
Expand Down Expand Up @@ -66,6 +67,7 @@ func Start(ctx context.Context, config *Config, client discovery.SvcDiscoveryReg
config.Share.Secret,
config.RpcConfig.TokenPolicy.Expire,
config.Share.MultiLogin,
config.Share.IMAdminUserID,
),
config: config,
})
Expand Down Expand Up @@ -129,6 +131,10 @@ func (s *authServer) parseToken(ctx context.Context, tokensString string) (claim
if err != nil {
return nil, errs.Wrap(err)
}
isAdmin := authverify.IsManagerUserID(claims.UserID, s.config.Share.IMAdminUserID)
if isAdmin {
return claims, nil
}
m, err := s.authDatabase.GetTokensWithoutError(ctx, claims.UserID, claims.PlatformID)
if err != nil {
return nil, err
Expand Down Expand Up @@ -190,7 +196,7 @@ func (s *authServer) forceKickOff(ctx context.Context, userID string, platformID
}

m, err := s.authDatabase.GetTokensWithoutError(ctx, userID, int(platformID))
if err != nil && err != redis.Nil {
if err != nil && errors.Is(err, redis.Nil) {
return err
}
for k := range m {
Expand All @@ -208,7 +214,7 @@ func (s *authServer) forceKickOff(ctx context.Context, userID string, platformID

func (s *authServer) InvalidateToken(ctx context.Context, req *pbauth.InvalidateTokenReq) (*pbauth.InvalidateTokenResp, error) {
m, err := s.authDatabase.GetTokensWithoutError(ctx, req.UserID, int(req.PlatformID))
if err != nil && err != redis.Nil {
if err != nil && errors.Is(err, redis.Nil) {
return nil, err
}
if m == nil {
Expand Down
69 changes: 39 additions & 30 deletions pkg/common/storage/controller/auth.go
Original file line number Diff line number Diff line change
Expand Up @@ -35,9 +35,10 @@ type authDatabase struct {
accessSecret string
accessExpire int64
multiLogin multiLoginConfig
adminUserIDs []string
}

func NewAuthDatabase(cache cache.TokenModel, accessSecret string, accessExpire int64, multiLogin config.MultiLogin) AuthDatabase {
func NewAuthDatabase(cache cache.TokenModel, accessSecret string, accessExpire int64, multiLogin config.MultiLogin, adminUserIDs []string) AuthDatabase {
return &authDatabase{cache: cache, accessSecret: accessSecret, accessExpire: accessExpire, multiLogin: multiLoginConfig{
Policy: multiLogin.Policy,
MaxNumOneEnd: multiLogin.MaxNumOneEnd,
Expand All @@ -53,7 +54,8 @@ func NewAuthDatabase(cache cache.TokenModel, accessSecret string, accessExpire i
constant.IPadPlatformID: multiLogin.CustomizeLoginNum.IPad,
constant.AdminPlatformID: multiLogin.CustomizeLoginNum.Admin,
},
}}
}, adminUserIDs: adminUserIDs,
}
}

// If the result is empty.
Expand Down Expand Up @@ -90,27 +92,31 @@ func (a *authDatabase) BatchSetTokenMapByUidPid(ctx context.Context, tokens []st

// Create Token.
func (a *authDatabase) CreateToken(ctx context.Context, userID string, platformID int) (string, error) {
tokens, err := a.cache.GetAllTokensWithoutError(ctx, userID)
if err != nil {
return "", err
}
deleteTokenKey, kickedTokenKey, err := a.checkToken(ctx, tokens, platformID)
if err != nil {
return "", err
}
if len(deleteTokenKey) != 0 {
err = a.cache.DeleteTokenByUidPid(ctx, userID, platformID, deleteTokenKey)
isAdmin := authverify.IsManagerUserID(userID, a.adminUserIDs)
if !isAdmin {
tokens, err := a.cache.GetAllTokensWithoutError(ctx, userID)
if err != nil {
return "", err
}
}
if len(kickedTokenKey) != 0 {
for _, k := range kickedTokenKey {
err := a.cache.SetTokenFlagEx(ctx, userID, platformID, k, constant.KickedToken)

deleteTokenKey, kickedTokenKey, err := a.checkToken(ctx, tokens, platformID)
if err != nil {
return "", err
}
if len(deleteTokenKey) != 0 {
err = a.cache.DeleteTokenByUidPid(ctx, userID, platformID, deleteTokenKey)
if err != nil {
return "", err
}
log.ZDebug(ctx, "kicked token in create token", "token", k)
}
if len(kickedTokenKey) != 0 {
for _, k := range kickedTokenKey {
err := a.cache.SetTokenFlagEx(ctx, userID, platformID, k, constant.KickedToken)
if err != nil {
return "", err
}
log.ZDebug(ctx, "kicked token in create token", "token", k)
}
}
}

Expand All @@ -121,9 +127,12 @@ func (a *authDatabase) CreateToken(ctx context.Context, userID string, platformI
return "", errs.WrapMsg(err, "token.SignedString")
}

if err = a.cache.SetTokenFlagEx(ctx, userID, platformID, tokenString, constant.NormalToken); err != nil {
return "", err
if !isAdmin {
if err = a.cache.SetTokenFlagEx(ctx, userID, platformID, tokenString, constant.NormalToken); err != nil {
return "", err
}
}

return tokenString, nil
}

Expand Down Expand Up @@ -226,16 +235,16 @@ func (a *authDatabase) checkToken(ctx context.Context, tokens map[int]map[string
return nil, nil, errs.New("unknown multiLogin policy").Wrap()
}

var adminTokenMaxNum = a.multiLogin.MaxNumOneEnd
if a.multiLogin.Policy == constant.Customize {
adminTokenMaxNum = a.multiLogin.CustomizeLoginNum[constant.AdminPlatformID]
}
l := len(adminToken)
if platformID == constant.AdminPlatformID {
l++
}
if l > adminTokenMaxNum {
kickToken = append(kickToken, adminToken[:l-adminTokenMaxNum]...)
}
//var adminTokenMaxNum = a.multiLogin.MaxNumOneEnd
//if a.multiLogin.Policy == constant.Customize {
// adminTokenMaxNum = a.multiLogin.CustomizeLoginNum[constant.AdminPlatformID]
//}
//l := len(adminToken)
//if platformID == constant.AdminPlatformID {
// l++
//}
//if l > adminTokenMaxNum {
// kickToken = append(kickToken, adminToken[:l-adminTokenMaxNum]...)
//}
return deleteToken, kickToken, nil
}
2 changes: 1 addition & 1 deletion pkg/common/storage/controller/msg.go
Original file line number Diff line number Diff line change
Expand Up @@ -490,7 +490,7 @@ func (db *commonMsgDatabase) GetMsgBySeqs(ctx context.Context, userID string, co
}
successMsgs, failedSeqs, err := db.msg.GetMessagesBySeq(ctx, conversationID, newSeqs)
if err != nil {
if err != redis.Nil {
if errors.Is(err, redis.Nil) {
log.ZError(ctx, "get message from redis exception", err, "failedSeqs", failedSeqs, "conversationID", conversationID)
}
}
Expand Down

0 comments on commit c00d6c5

Please sign in to comment.