chore: Custom cooke resolver and logout on session invalidation

Author: nidhi-nair

app/server/appsmith-server/src/main/java/com/appsmith/server/configurations/CustomCookieWebSessionIdResolver.java

@@ -0,0 +1,34 @@
+package com.appsmith.server.configurations;
+
+import org.springframework.web.server.ServerWebExchange;
+import org.springframework.web.server.session.CookieWebSessionIdResolver;
+
+import java.time.Duration;
+
+import static java.time.temporal.ChronoUnit.DAYS;
+
+/**
+ * This class is a custom implementation of the CookieWebSessionIdResolver class.
+ * It allows us to set the SameSite attribute of the session cookie based on
+ * the organization configuration for cross site embedding.
+ */
+public class CustomCookieWebSessionIdResolver extends CookieWebSessionIdResolver {
+
+    public static final String LAX = "Lax";
+
+    public CustomCookieWebSessionIdResolver() {
+        // Set default cookie attributes
+
+        // Setting the max age to 30 days so that the cookie doesn't expire on browser close
+        // If the max age is not set, some browsers will default to deleting the cookies on session close.
+        this.setCookieMaxAge(Duration.of(30, DAYS));
+        this.addCookieInitializer((builder) -> builder.path("/"));
+    }
+
+    @Override
+    public void setSessionId(ServerWebExchange exchange, String id) {
+        // Add the appropriate SameSite attribute based on the exchange attribute
+        addCookieInitializer((builder) -> builder.sameSite(LAX).secure(true));
+        super.setSessionId(exchange, id);
+    }
+}

app/server/appsmith-server/src/main/java/com/appsmith/server/configurations/SecurityConfig.java

@@ -53,11 +53,9 @@
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebFilterChain;
 import org.springframework.web.server.adapter.ForwardedHeaderTransformer;
-import org.springframework.web.server.session.CookieWebSessionIdResolver;
 import org.springframework.web.server.session.WebSessionIdResolver;
 import reactor.core.publisher.Mono;
 
-import java.time.Duration;
 import java.util.HashSet;
 import java.util.List;
 
@@ -73,7 +71,6 @@
 import static com.appsmith.server.constants.Url.USAGE_PULSE_URL;
 import static com.appsmith.server.constants.Url.USER_URL;
 import static com.appsmith.server.constants.ce.UrlCE.CONSOLIDATED_API_URL;
-import static java.time.temporal.ChronoUnit.DAYS;
 
 @EnableWebFluxSecurity
 @EnableReactiveMethodSecurity
@@ -271,13 +268,7 @@ public SecurityWebFilterChain securityWebFilterChain(ServerHttpSecurity http) {
      */
     @Bean
     public WebSessionIdResolver webSessionIdResolver() {
-        CookieWebSessionIdResolver resolver = new CookieWebSessionIdResolver();
-        // Setting the max age to 30 days so that the cookie doesn't expire on browser close
-        // If the max age is not set, some browsers will default to deleting the cookies on session close.
-        resolver.setCookieMaxAge(Duration.of(30, DAYS));
-        resolver.addCookieInitializer((builder) -> builder.path("/"));
-        resolver.addCookieInitializer((builder) -> builder.sameSite("Lax"));
-        return resolver;
+        return new CustomCookieWebSessionIdResolver();
     }
 
     private User createAnonymousUser() {

app/server/appsmith-server/src/main/java/com/appsmith/server/exceptions/GlobalExceptionHandler.java

@@ -54,6 +54,15 @@ public class GlobalExceptionHandler {
 
     private final SessionUserService sessionUserService;
 
+    @ExceptionHandler(IllegalStateException.class)
+    public Mono<Void> handleSessionInvalidation(ServerWebExchange exchange, IllegalStateException ex) {
+        if (ex.getMessage().contains("Session was invalidated")) {
+            exchange.getResponse().setStatusCode(HttpStatus.UNAUTHORIZED);
+            return exchange.getResponse().setComplete();
+        }
+        return Mono.error(ex);
+    }
+
     /**
      * This function only catches the AppsmithException type and formats it into ResponseEntity<ErrorDTO> object
      * Ideally, we should only be throwing AppsmithException from our code. This ensures that we can standardize