Merge pull request #589 from lowcoder-org/update-refresh-token-handling

FalkWolsky · web-flow · commit dff211958537 · 2023-12-20T18:12:04.000+01:00
Rework Oauth Refresh Token Handling
diff --git a/client/VERSION b/client/VERSION
@@ -1 +1 @@
-2.2.1
+dev
diff --git a/server/api-service/lowcoder-domain/src/main/java/org/lowcoder/domain/user/model/User.java b/server/api-service/lowcoder-domain/src/main/java/org/lowcoder/domain/user/model/User.java
@@ -46,6 +46,8 @@ public class User extends HasIdAndAuditing implements BeforeMongodbWrite, AfterM
 
     private Boolean isEnabled = true;
 
+    private String activeAuthId;
+
     // used in form login
     @JsonProperty(access = JsonProperty.Access.WRITE_ONLY)
     private String password;
diff --git a/server/api-service/lowcoder-domain/src/main/java/org/lowcoder/domain/user/service/UserServiceImpl.java b/server/api-service/lowcoder-domain/src/main/java/org/lowcoder/domain/user/service/UserServiceImpl.java
@@ -210,7 +210,10 @@ public Mono<Boolean> addNewConnection(String userId, Connection connection) {
     @Override
     public Mono<User> addNewConnectionAndReturnUser(String userId, Connection connection) {
         return findById(userId)
-                .doOnNext(user -> user.getConnections().add(connection))
+                .doOnNext(user -> {
+                    user.getConnections().add(connection);
+                    user.setActiveAuthId(connection.getAuthId());
+                })
                 .flatMap(repository::save);
     }
 
diff --git a/server/api-service/lowcoder-server/src/main/java/org/lowcoder/api/authentication/request/oauth2/request/KeycloakRequest.java b/server/api-service/lowcoder-server/src/main/java/org/lowcoder/api/authentication/request/oauth2/request/KeycloakRequest.java
@@ -94,6 +94,7 @@ protected Mono<AuthToken> refreshAuthToken(String refreshToken) {
                             .accessToken(MapUtils.getString(map, "access_token"))
                             .expireIn(MapUtils.getIntValue(map, "expires_in"))
                             .refreshToken(MapUtils.getString(map, "refresh_token"))
+                            .refreshTokenExpireIn(MapUtils.getIntValue(map, "refresh_expires_in"))
                             .build();
                     return Mono.just(authToken);
                 });
diff --git a/server/api-service/lowcoder-server/src/main/java/org/lowcoder/api/authentication/service/AuthenticationApiServiceImpl.java b/server/api-service/lowcoder-server/src/main/java/org/lowcoder/api/authentication/service/AuthenticationApiServiceImpl.java
@@ -242,6 +242,8 @@ public void updateConnection(AuthUser authUser, User user) {
         oldConnection.setAuthConnectionAuthToken(
                 Optional.ofNullable(authUser.getAuthToken()).map(ConnectionAuthToken::of).orElse(null));
         oldConnection.setRawUserInfo(authUser.getRawUserInfo());
+
+        user.setActiveAuthId(oldConnection.getAuthId());
     }
 
     @SuppressWarnings("OptionalGetWithoutIsPresent")
diff --git a/server/api-service/lowcoder-server/src/main/java/org/lowcoder/api/framework/filter/UserSessionPersistenceFilter.java b/server/api-service/lowcoder-server/src/main/java/org/lowcoder/api/framework/filter/UserSessionPersistenceFilter.java
@@ -1,15 +1,18 @@
 package org.lowcoder.api.framework.filter;
 
 import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang3.tuple.Triple;
 import org.lowcoder.api.authentication.request.AuthRequest;
 import org.lowcoder.api.authentication.request.AuthRequestFactory;
 import org.lowcoder.api.authentication.request.oauth2.OAuth2RequestContext;
 import org.lowcoder.api.authentication.service.AuthenticationApiServiceImpl;
 import org.lowcoder.api.home.SessionUserService;
 import org.lowcoder.domain.authentication.AuthenticationService;
-import org.lowcoder.domain.authentication.FindAuthConfig;
 import org.lowcoder.domain.authentication.context.AuthRequestContext;
 import org.lowcoder.domain.user.model.AuthUser;
+import org.lowcoder.domain.user.model.Connection;
+import org.lowcoder.domain.user.model.User;
+import org.lowcoder.domain.user.service.UserService;
 import org.lowcoder.sdk.util.CookieHelper;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebFilter;
@@ -18,8 +21,7 @@
 
 import javax.annotation.Nonnull;
 import java.time.Instant;
-import java.util.LinkedList;
-import java.util.List;
+import java.util.Optional;
 
 import static org.lowcoder.api.authentication.util.AuthenticationUtils.toAuthentication;
 import static org.lowcoder.domain.authentication.AuthenticationService.DEFAULT_AUTH_CONFIG;
@@ -29,6 +31,8 @@
 public class UserSessionPersistenceFilter implements WebFilter {
 
     private final SessionUserService service;
+
+    private final UserService userService;
     private final CookieHelper cookieHelper;
 
     private final AuthenticationService authenticationService;
@@ -37,9 +41,10 @@ public class UserSessionPersistenceFilter implements WebFilter {
 
     private final AuthRequestFactory<AuthRequestContext> authRequestFactory;
 
-    public UserSessionPersistenceFilter(SessionUserService service, CookieHelper cookieHelper, AuthenticationService authenticationService,
+    public UserSessionPersistenceFilter(SessionUserService service, UserService userService, CookieHelper cookieHelper, AuthenticationService authenticationService,
                                         AuthenticationApiServiceImpl authenticationApiService, AuthRequestFactory<AuthRequestContext> authRequestFactory) {
         this.service = service;
+        this.userService = userService;
         this.cookieHelper = cookieHelper;
         this.authenticationService = authenticationService;
         this.authenticationApiService = authenticationApiService;
@@ -52,48 +57,88 @@ public Mono<Void> filter(@Nonnull ServerWebExchange exchange, WebFilterChain cha
         String cookieToken = cookieHelper.getCookieToken(exchange);
         return service.resolveSessionUserFromCookie(cookieToken)
                 .switchIfEmpty(chain.filter(exchange).then(Mono.empty()))
-                .doOnNext(user -> {
+                .map(user -> {
 
-                    List<String> tokensToRemove = new LinkedList<>();
+                    Connection activeConnection = null;
+                    String orgId = null;
 
-                    user.getConnections().forEach(connection -> {
-                        if(!connection.getAuthId().equals(DEFAULT_AUTH_CONFIG.getId())) {
-                            Instant next5Minutes = Instant.now().plusSeconds( 300 );
-                            if(connection.getAuthConnectionAuthToken().getExpireAt() == 0) {
-                                return;
-                            }
-                            boolean isAccessTokenExpiryNear = (connection.getAuthConnectionAuthToken().getExpireAt()*1000) <= next5Minutes.toEpochMilli();
-                            if(isAccessTokenExpiryNear) {
-                                connection.getOrgIds().forEach(orgId -> {
-                                    authenticationService.findAuthConfigByAuthId(orgId, connection.getAuthId())
-	                                	.doOnSuccess(findAuthConfig -> {
-	                                        if(findAuthConfig == null) {
-	                                            return;
-	                                        }
-	                                        OAuth2RequestContext oAuth2RequestContext = new OAuth2RequestContext(orgId, null, null);
-	                                        oAuth2RequestContext.setAuthConfig(findAuthConfig.authConfig());
-	                                        AuthRequest authRequest = authRequestFactory.build(oAuth2RequestContext).block();
-	                                        try {
-	                                            AuthUser authUser = authRequest.refresh(connection.getAuthConnectionAuthToken().getRefreshToken()).block();
-	                                            authUser.setAuthContext(oAuth2RequestContext);
-	                                            authenticationApiService.updateConnection(authUser, user);
-	                                        } catch (Exception e) {
-	                                            log.error("Failed to refresh access token. Removing user sessions/tokens.");
-	                                            tokensToRemove.addAll(connection.getTokens());
-	                                        }
-	                                	});
-                                });
+                    Optional<Connection> activeConnectionOptional = user.getConnections()
+                            .stream()
+                            .filter(connection -> connection.getAuthId().equals(user.getActiveAuthId()))
+                            .findFirst();
+
+                    if(!activeConnectionOptional.isPresent()) {
+                        return Triple.of(user, activeConnection, orgId);
+                    }
+
+                    activeConnection = activeConnectionOptional.get();
+
+                    if(!activeConnection.getAuthId().equals(DEFAULT_AUTH_CONFIG.getId())) {
+                        if(activeConnection.getAuthConnectionAuthToken().getExpireAt() == 0) {
+                            return Triple.of(user, activeConnection, orgId);
+                        }
+                        boolean isAccessTokenExpired = (activeConnection.getAuthConnectionAuthToken().getExpireAt()*1000) < Instant.now().toEpochMilli();
+                        if(isAccessTokenExpired) {
+
+                            Optional<String> orgIdOptional = activeConnection.getOrgIds().stream().findFirst();
+                            if(!orgIdOptional.isPresent()) {
+                                return Triple.of(user, activeConnection, orgId);
                             }
+                            orgId = orgIdOptional.get();
                         }
-                    });
+                    }
 
-                    tokensToRemove.forEach(token -> {
-                        service.removeUserSession(token).block();
-                    });
+                    return Triple.of(user, activeConnection, orgId);
 
-                })
+                }).flatMap(this::refreshOauthToken)
                 .flatMap(user -> chain.filter(exchange).contextWrite(withAuthentication(toAuthentication(user)))
                         .then(service.extendValidity(cookieToken))
                 );
     }
+
+    private Mono<User> refreshOauthToken(Triple<User, Connection, String> triple) {
+
+        User user = triple.getLeft();
+        Connection connection = triple.getMiddle();
+        String orgId = triple.getRight();
+
+        if (connection == null || orgId == null) {
+            return Mono.just(user);
+        }
+
+        OAuth2RequestContext oAuth2RequestContext = new OAuth2RequestContext(triple.getRight(), null, null);
+
+        return authenticationService
+                .findAuthConfigByAuthId(orgId, connection.getAuthId())
+                .switchIfEmpty(Mono.empty())
+                .flatMap(findAuthConfig -> {
+
+                    Mono<AuthRequest> authRequestMono = Mono.empty();
+
+                    if(findAuthConfig == null) {
+                        return authRequestMono;
+                    }
+                    oAuth2RequestContext.setAuthConfig(findAuthConfig.authConfig());
+
+                    return authRequestFactory.build(oAuth2RequestContext);
+                }).flatMap(authRequest -> {
+                    if(authRequest == null) {
+                        return Mono.just(user);
+                    }
+                    try {
+                        AuthUser authUser = authRequest.refresh(connection.getAuthConnectionAuthToken().getRefreshToken()).block();
+                        authUser.setAuthContext(oAuth2RequestContext);
+                        authenticationApiService.updateConnection(authUser, user);
+                        return userService.update(user.getId(), user);
+                    } catch (Exception e) {
+                        log.error("Failed to refresh access token. Removing user sessions/tokens.");
+                        connection.getTokens().forEach(token -> {
+                            service.removeUserSession(token).block();
+                        });
+                    }
+                    return Mono.just(user);
+                });
+
+    }
+
 }
diff --git a/server/api-service/lowcoder-server/src/main/java/org/lowcoder/api/framework/security/SecurityConfig.java b/server/api-service/lowcoder-server/src/main/java/org/lowcoder/api/framework/security/SecurityConfig.java
@@ -10,6 +10,7 @@
 import org.lowcoder.domain.authentication.AuthenticationService;
 import org.lowcoder.domain.authentication.context.AuthRequestContext;
 import org.lowcoder.domain.user.model.User;
+import org.lowcoder.domain.user.service.UserService;
 import org.lowcoder.infra.constant.NewUrl;
 import org.lowcoder.sdk.config.CommonConfig;
 import org.lowcoder.sdk.util.CookieHelper;
@@ -50,6 +51,9 @@ public class SecurityConfig {
     @Autowired
     private SessionUserService sessionUserService;
 
+    @Autowired
+    private UserService userService;
+
     @Autowired
     private AccessDeniedHandler accessDeniedHandler;
 
@@ -153,7 +157,7 @@ SecurityWebFilterChain securityWebFilterChain(ServerHttpSecurity http) {
                 .accessDeniedHandler(accessDeniedHandler)
         );
 
-        http.addFilterBefore(new UserSessionPersistenceFilter(sessionUserService, cookieHelper, authenticationService, authenticationApiService, authRequestFactory), SecurityWebFiltersOrder.AUTHENTICATION);
+        http.addFilterBefore(new UserSessionPersistenceFilter(sessionUserService, userService, cookieHelper, authenticationService, authenticationApiService, authRequestFactory), SecurityWebFiltersOrder.AUTHENTICATION);
         http.addFilterBefore(new APIKeyAuthFilter(sessionUserService, cookieHelper, jwtUtils), SecurityWebFiltersOrder.AUTHENTICATION);
 
         return http.build();

Original file line number	Diff line number	Diff line change
`@@ -242,6 +242,8 @@ public void updateConnection(AuthUser authUser, User user) {`
`242`	`242`	`oldConnection.setAuthConnectionAuthToken(`
`243`	`243`	`Optional.ofNullable(authUser.getAuthToken()).map(ConnectionAuthToken::of).orElse(null));`
`244`	`244`	`oldConnection.setRawUserInfo(authUser.getRawUserInfo());`
	`245`	`+`
	`246`	`+ user.setActiveAuthId(oldConnection.getAuthId());`
`245`	`247`	`}`
`246`	`248`
`247`	`249`	`@SuppressWarnings("OptionalGetWithoutIsPresent")`