chore: update workspace store

api/v1/auth.go

@@ -122,7 +122,32 @@ func (s *APIV1Service) registerAuthRoutes(g *echo.Group, secret string) {
 	})
 
 	g.POST("/auth/logout", func(c echo.Context) error {
+		ctx := c.Request().Context()
 		RemoveTokensAndCookies(c)
+		accessToken := findAccessToken(c)
+		userID, _ := getUserIDFromAccessToken(accessToken, 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.UpsertUserSetting(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)
+			}
+		}
 		c.Response().WriteHeader(http.StatusOK)
 		return nil
 	})

api/v1/jwt.go

@@ -48,15 +48,6 @@ func findAccessToken(c echo.Context) string {
 	return accessToken
 }
 
-func audienceContains(audience jwt.ClaimStrings, token string) bool {
-	for _, v := range audience {
-		if v == token {
-			return true
-		}
-	}
-	return false
-}
-
 // JWTMiddleware validates the access token.
 func JWTMiddleware(s *APIV1Service, next echo.HandlerFunc, secret string) echo.HandlerFunc {
 	return func(c echo.Context) error {
@@ -69,8 +60,8 @@ func JWTMiddleware(s *APIV1Service, next echo.HandlerFunc, secret string) echo.H
 			return next(c)
 		}
 
-		token := findAccessToken(c)
-		if token == "" {
+		accessToken := findAccessToken(c)
+		if accessToken == "" {
 			// When the request is not authenticated, we allow the user to access the shortcut endpoints for those public shortcuts.
 			if util.HasPrefixes(path, "/s/") && method == http.MethodGet {
 				return next(c)
@@ -78,36 +69,16 @@ func JWTMiddleware(s *APIV1Service, next echo.HandlerFunc, secret string) echo.H
 			return echo.NewHTTPError(http.StatusUnauthorized, "Missing access token")
 		}
 
-		claims := &auth.ClaimsMessage{}
-		_, err := jwt.ParseWithClaims(token, claims, func(t *jwt.Token) (any, error) {
-			if t.Method.Alg() != jwt.SigningMethodHS256.Name {
-				return nil, errors.Errorf("unexpected access token signing method=%v, expect %v", t.Header["alg"], jwt.SigningMethodHS256)
-			}
-			if kid, ok := t.Header["kid"].(string); ok {
-				if kid == "v1" {
-					return []byte(secret), nil
-				}
-			}
-			return nil, errors.Errorf("unexpected access token kid=%v", t.Header["kid"])
-		})
+		userID, err := getUserIDFromAccessToken(accessToken, secret)
 		if err != nil {
-			return echo.NewHTTPError(http.StatusUnauthorized, errors.Wrap(err, "Invalid or expired access token"))
-		}
-		if !audienceContains(claims.Audience, auth.AccessTokenAudienceName) {
-			return echo.NewHTTPError(http.StatusUnauthorized, fmt.Sprintf("Invalid access token, audience mismatch, got %q, expected %q.", claims.Audience, auth.AccessTokenAudienceName))
-		}
-
-		// We either have a valid access token or we will attempt to generate new access token.
-		userID, err := util.ConvertStringToInt32(claims.Subject)
-		if err != nil {
-			return echo.NewHTTPError(http.StatusUnauthorized, "Malformed ID in the token.").WithInternal(err)
+			return echo.NewHTTPError(http.StatusUnauthorized, "Invalid or expired access token")
 		}
 
 		accessTokens, err := s.Store.GetUserAccessTokens(ctx, userID)
 		if err != nil {
 			return echo.NewHTTPError(http.StatusInternalServerError, "Failed to get user access tokens.").WithInternal(err)
 		}
-		if !validateAccessToken(token, accessTokens) {
+		if !validateAccessToken(accessToken, accessTokens) {
 			return echo.NewHTTPError(http.StatusUnauthorized, "Invalid access token.")
 		}
 
@@ -128,6 +99,30 @@ func JWTMiddleware(s *APIV1Service, next echo.HandlerFunc, secret string) echo.H
 	}
 }
 
+func getUserIDFromAccessToken(accessToken, secret string) (int32, error) {
+	claims := &auth.ClaimsMessage{}
+	_, err := jwt.ParseWithClaims(accessToken, claims, func(t *jwt.Token) (any, error) {
+		if t.Method.Alg() != jwt.SigningMethodHS256.Name {
+			return nil, errors.Errorf("unexpected access token signing method=%v, expect %v", t.Header["alg"], jwt.SigningMethodHS256)
+		}
+		if kid, ok := t.Header["kid"].(string); ok {
+			if kid == "v1" {
+				return []byte(secret), nil
+			}
+		}
+		return nil, errors.Errorf("unexpected access token kid=%v", t.Header["kid"])
+	})
+	if err != nil {
+		return 0, errors.Wrap(err, "Invalid or expired access token")
+	}
+	// We either have a valid access token or we will attempt to generate new access token.
+	userID, err := util.ConvertStringToInt32(claims.Subject)
+	if err != nil {
+		return 0, errors.Wrap(err, "Malformed ID in the token")
+	}
+	return userID, nil
+}
+
 func validateAccessToken(accessTokenString string, userAccessTokens []*storepb.AccessTokensUserSetting_AccessToken) bool {
 	for _, userAccessToken := range userAccessTokens {
 		if accessTokenString == userAccessToken.AccessToken {