feat: JSON Session serialization on Redis (#15368)

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

@@ -1,5 +1,7 @@
 package com.appsmith.server.configurations;
 
+import com.appsmith.server.domains.UserSession;
+import com.fasterxml.jackson.databind.json.JsonMapper;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
@@ -8,13 +10,19 @@ import org.springframework.data.redis.connection.ReactiveRedisConnectionFactory;
 import org.springframework.data.redis.core.ReactiveRedisOperations;
 import org.springframework.data.redis.core.ReactiveRedisTemplate;
 import org.springframework.data.redis.listener.ChannelTopic;
+import org.springframework.data.redis.serializer.GenericJackson2JsonRedisSerializer;
 import org.springframework.data.redis.serializer.Jackson2JsonRedisSerializer;
 import org.springframework.data.redis.serializer.JdkSerializationRedisSerializer;
 import org.springframework.data.redis.serializer.RedisSerializationContext;
 import org.springframework.data.redis.serializer.RedisSerializer;
 import org.springframework.data.redis.serializer.StringRedisSerializer;
+import org.springframework.data.redis.util.ByteUtils;
+import org.springframework.security.core.context.SecurityContext;
+import org.springframework.security.core.context.SecurityContextImpl;
 import org.springframework.session.data.redis.config.annotation.web.server.EnableRedisWebSession;
 
+import java.util.Arrays;
+
 @Configuration
 @Slf4j
 // Setting the maxInactiveInterval to 30 days
@@ -32,6 +40,11 @@ public class RedisConfig {
         return new ChannelTopic("appsmith:queue");
     }
 
+    @Bean
+    public RedisSerializer<Object> springSessionDefaultRedisSerializer() {
+        return new JSONSessionRedisSerializer();
+    }
+
     @Primary
     @Bean
     ReactiveRedisOperations<String, String> reactiveRedisOperations(ReactiveRedisConnectionFactory factory) {
@@ -47,6 +60,7 @@ public class RedisConfig {
 
     // Lifted from below and turned it into a bean. Wish Spring provided it as a bean.
     // RedisWebSessionConfiguration.createReactiveRedisTemplate
+
     @Bean
     ReactiveRedisTemplate<String, Object> reactiveRedisTemplate(ReactiveRedisConnectionFactory factory) {
         RedisSerializer<String> keySerializer = new StringRedisSerializer();
@@ -57,4 +71,40 @@ public class RedisConfig {
         return new ReactiveRedisTemplate<>(factory, serializationContext);
     }
 
+    private static class JSONSessionRedisSerializer implements RedisSerializer<Object> {
+
+        private final JdkSerializationRedisSerializer fallback = new JdkSerializationRedisSerializer();
+
+        private final GenericJackson2JsonRedisSerializer jsonSerializer = new GenericJackson2JsonRedisSerializer(new JsonMapper());
+
+        private static final byte[] SESSION_DATA_PREFIX = "appsmith-session:".getBytes();
+
+        @Override
+        public byte[] serialize(Object t) {
+            if (t instanceof SecurityContext) {
+                final UserSession session = UserSession.fromToken(((SecurityContext) t).getAuthentication());
+                final byte[] bytes = jsonSerializer.serialize(session);
+                return bytes == null ? null : ByteUtils.concat(SESSION_DATA_PREFIX, bytes);
+            }
+
+            return fallback.serialize(t);
+        }
+
+        @Override
+        public Object deserialize(byte[] bytes) {
+            if (ByteUtils.startsWith(bytes, SESSION_DATA_PREFIX)) {
+                final byte[] data = Arrays.copyOfRange(bytes, SESSION_DATA_PREFIX.length, bytes.length);
+                final UserSession session = jsonSerializer.deserialize(data, UserSession.class);
+
+                if (session == null) {
+                    throw new IllegalArgumentException("Could not deserialize user session, got null");
+                }
+
+                return new SecurityContextImpl(session.makeToken());
+            }
+
+            return fallback.deserialize(bytes);
+        }
+    }
+
 }

app/server/appsmith-server/src/main/java/com/appsmith/server/domains/UserSession.java

@@ -0,0 +1,112 @@
+package com.appsmith.server.domains;
+
+import lombok.Data;
+import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
+import org.springframework.security.core.Authentication;
+import org.springframework.security.core.GrantedAuthority;
+import org.springframework.security.oauth2.client.authentication.OAuth2AuthenticationToken;
+
+import java.util.Collection;
+import java.util.Set;
+
+/**
+ * UserSession is a POJO class that represents a user's session. It is serialized to JSON and stored in Redis. That
+ * means that this class doesn't have to be serializable, and the serialVersionUID is not required. This class can
+ * change/evolve in the future, as long as pre-existing JSON session data can be safely deserialized.
+ */
+@Data
+public class UserSession {
+
+    private String userId;
+
+    private String email;
+
+    private LoginSource source;
+
+    private UserState state;
+
+    private Boolean isEnabled;
+
+    private String currentWorkspaceId;
+
+    private Set<String> workspaceIds;
+
+    private String tenantId;
+
+    private Object credentials;
+
+    private Collection<? extends GrantedAuthority> authorities;
+
+    private String authorizedClientRegistrationId;
+
+    private static final String PASSWORD_PROVIDER = "password";
+
+    private static final Set<String> ALLOWED_OAUTH_PROVIDERS = Set.of("google", "github");
+
+    /**
+     * We don't expect this class to be instantiated outside this class. Remove this constructor when needed.
+     */
+    private UserSession() {}
+
+    /**
+     * Given an authentication token, typically from a Spring Security context, create a UserSession object. This
+     * UserSession object can then be serialized to JSON and stored in Redis.
+     * @param authentication The token to create the UserSession from. Usually an instance of UsernamePasswordAuthenticationToken or Oauth2AuthenticationToken.
+     * @return A UserSession object representing the user's session, with details from the given token.
+     */
+    public static UserSession fromToken(Authentication authentication) {
+        final UserSession session = new UserSession();
+        final User user = (User) authentication.getPrincipal();
+
+        session.userId = user.getId();
+        session.email = user.getEmail();
+        session.source = user.getSource();
+        session.state = user.getState();
+        session.isEnabled = user.isEnabled();
+        session.currentWorkspaceId = user.getCurrentWorkspaceId();
+        session.workspaceIds = user.getWorkspaceIds();
+        session.tenantId = user.getTenantId();
+
+        session.credentials = authentication.getCredentials();
+        session.authorities = authentication.getAuthorities();
+
+        if (authentication instanceof OAuth2AuthenticationToken) {
+            session.authorizedClientRegistrationId = ((OAuth2AuthenticationToken) authentication).getAuthorizedClientRegistrationId();
+        } else if (authentication instanceof UsernamePasswordAuthenticationToken) {
+            session.authorizedClientRegistrationId = PASSWORD_PROVIDER;
+        } else {
+            throw new IllegalArgumentException("Unsupported authentication type: " + authentication.getClass().getName());
+        }
+
+        return session;
+    }
+
+    /**
+     * Performs the reverse of fromToken method. Given a UserSession object, create a Spring Security authentication
+     * token. This authentication token can then be wrapped in a SecurityContext and used as the user's session.
+     * @return A Spring Security authentication token representing the user's session. Usually an instance of UsernamePasswordAuthenticationToken or Oauth2AuthenticationToken.
+     */
+    public Authentication makeToken() {
+        final User user = new User();
+
+        user.setId(userId);
+        user.setEmail(email);
+        user.setSource(source);
+        user.setState(state);
+        user.setIsEnabled(isEnabled);
+        user.setCurrentWorkspaceId(currentWorkspaceId);
+        user.setWorkspaceIds(workspaceIds);
+        user.setTenantId(tenantId);
+
+        if (PASSWORD_PROVIDER.equals(authorizedClientRegistrationId)) {
+            return new UsernamePasswordAuthenticationToken(user, credentials, authorities);
+
+        } else if (ALLOWED_OAUTH_PROVIDERS.contains(authorizedClientRegistrationId)) {
+            return new OAuth2AuthenticationToken(user, authorities, authorizedClientRegistrationId);
+
+        }
+
+        throw new IllegalArgumentException("Invalid registration ID " + authorizedClientRegistrationId);
+    }
+
+}