From 91296257fc857539c4ebb1eee01beefc6d862bc4 Mon Sep 17 00:00:00 2001
From: Athurg Gooth <athurg@gooth.org>
Date: Wed, 22 Nov 2023 23:20:45 +0800
Subject: [PATCH] chore: remove invalid access token from db (#2539)

Remove invalid access token from db
---
 api/v1/auth.go        | 33 ++++++++++-----------------------
 api/v1/jwt.go         | 12 ++++++++++--
 store/user_setting.go | 27 +++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 25 deletions(-)

diff --git a/api/v1/auth.go b/api/v1/auth.go
index d10b9b30..625873ea 100644
--- a/api/v1/auth.go
+++ b/api/v1/auth.go
@@ -254,33 +254,14 @@ func (s *APIV1Service) SignInSSO(c echo.Context) error {
 //	@Success	200	{boolean}	true	"Sign-out success"
 //	@Router		/api/v1/auth/signout [POST]
 func (s *APIV1Service) SignOut(c echo.Context) error {
-	ctx := c.Request().Context()
 	accessToken := findAccessToken(c)
 	userID, _ := getUserIDFromAccessToken(accessToken, s.Secret)
-	userAccessTokens, err := s.Store.GetUserAccessTokens(ctx, userID)
-	// Auto remove the current access token from the user access tokens.
-	if err == nil && len(userAccessTokens) != 0 {
-		accessTokens := []*storepb.AccessTokensUserSetting_AccessToken{}
-		for _, userAccessToken := range userAccessTokens {
-			if accessToken != userAccessToken.AccessToken {
-				accessTokens = append(accessTokens, userAccessToken)
-			}
-		}
 
-		if _, err := s.Store.UpsertUserSettingV1(ctx, &storepb.UserSetting{
-			UserId: userID,
-			Key:    storepb.UserSettingKey_USER_SETTING_ACCESS_TOKENS,
-			Value: &storepb.UserSetting_AccessTokens{
-				AccessTokens: &storepb.AccessTokensUserSetting{
-					AccessTokens: accessTokens,
-				},
-			},
-		}); err != nil {
-			return echo.NewHTTPError(http.StatusInternalServerError, fmt.Sprintf("failed to upsert user setting, err: %s", err)).SetInternal(err)
-		}
+	err := removeAccessTokenAndCookies(c, s.Store, userID, accessToken)
+	if err != nil {
+		return echo.NewHTTPError(http.StatusInternalServerError, fmt.Sprintf("failed to remove access token, err: %s", err)).SetInternal(err)
 	}
 
-	removeAccessTokenAndCookies(c)
 	return c.JSON(http.StatusOK, true)
 }
 
@@ -393,9 +374,15 @@ func (s *APIV1Service) UpsertAccessTokenToStore(ctx context.Context, user *store
 }
 
 // removeAccessTokenAndCookies removes the jwt token from the cookies.
-func removeAccessTokenAndCookies(c echo.Context) {
+func removeAccessTokenAndCookies(c echo.Context, s *store.Store, userID int32, token string) error {
+	err := s.RemoveUserAccessToken(c.Request().Context(), userID, token)
+	if err != nil {
+		return err
+	}
+
 	cookieExp := time.Now().Add(-1 * time.Hour)
 	setTokenCookie(c, auth.AccessTokenCookieName, "", cookieExp)
+	return nil
 }
 
 // setTokenCookie sets the token to the cookie.
diff --git a/api/v1/jwt.go b/api/v1/jwt.go
index 7aa940d3..a596f3c7 100644
--- a/api/v1/jwt.go
+++ b/api/v1/jwt.go
@@ -8,8 +8,10 @@ import (
 	"github.com/golang-jwt/jwt/v4"
 	"github.com/labstack/echo/v4"
 	"github.com/pkg/errors"
+	"go.uber.org/zap"
 
 	"github.com/usememos/memos/api/auth"
+	"github.com/usememos/memos/internal/log"
 	"github.com/usememos/memos/internal/util"
 	storepb "github.com/usememos/memos/proto/gen/store"
 	"github.com/usememos/memos/store"
@@ -79,7 +81,10 @@ func JWTMiddleware(server *APIV1Service, next echo.HandlerFunc, secret string) e
 
 		userID, err := getUserIDFromAccessToken(accessToken, secret)
 		if err != nil {
-			removeAccessTokenAndCookies(c)
+			err = removeAccessTokenAndCookies(c, server.Store, userID, accessToken)
+			if err != nil {
+				log.Error("fail to remove AccessToken and Cookies", zap.Error(err))
+			}
 			return echo.NewHTTPError(http.StatusUnauthorized, "Invalid or expired access token")
 		}
 
@@ -88,7 +93,10 @@ func JWTMiddleware(server *APIV1Service, next echo.HandlerFunc, secret string) e
 			return echo.NewHTTPError(http.StatusInternalServerError, "Failed to get user access tokens.").WithInternal(err)
 		}
 		if !validateAccessToken(accessToken, accessTokens) {
-			removeAccessTokenAndCookies(c)
+			err = removeAccessTokenAndCookies(c, server.Store, userID, accessToken)
+			if err != nil {
+				log.Error("fail to remove AccessToken and Cookies", zap.Error(err))
+			}
 			return echo.NewHTTPError(http.StatusUnauthorized, "Invalid access token.")
 		}
 
diff --git a/store/user_setting.go b/store/user_setting.go
index bac59b72..7408ad4a 100644
--- a/store/user_setting.go
+++ b/store/user_setting.go
@@ -123,3 +123,30 @@ func (s *Store) GetUserAccessTokens(ctx context.Context, userID int32) ([]*store
 	accessTokensUserSetting := userSetting.GetAccessTokens()
 	return accessTokensUserSetting.AccessTokens, nil
 }
+
+// RemoveUserAccessToken remove the access token of the user.
+func (s *Store) RemoveUserAccessToken(ctx context.Context, userID int32, token string) error {
+	oldAccessTokens, err := s.GetUserAccessTokens(ctx, userID)
+	if err != nil {
+		return err
+	}
+
+	newAccessTokens := make([]*storepb.AccessTokensUserSetting_AccessToken, 0, len(oldAccessTokens))
+	for _, t := range oldAccessTokens {
+		if token != t.AccessToken {
+			newAccessTokens = append(newAccessTokens, t)
+		}
+	}
+
+	_, err = s.UpsertUserSettingV1(ctx, &storepb.UserSetting{
+		UserId: userID,
+		Key:    storepb.UserSettingKey_USER_SETTING_ACCESS_TOKENS,
+		Value: &storepb.UserSetting_AccessTokens{
+			AccessTokens: &storepb.AccessTokensUserSetting{
+				AccessTokens: newAccessTokens,
+			},
+		},
+	})
+
+	return err
+}