Leverage Spring Session's IndexResolver.

Resolves #104.

Author: gregturn

src/main/java/org/springframework/session/data/mongo/AbstractMongoSessionConverter.java

@@ -30,8 +30,10 @@
 import org.springframework.data.mongodb.core.index.IndexOperations;
 import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.lang.Nullable;
+import org.springframework.session.DelegatingIndexResolver;
 import org.springframework.session.FindByIndexNameSessionRepository;
-import org.springframework.session.Session;
+import org.springframework.session.IndexResolver;
+import org.springframework.session.PrincipalNameIndexResolver;
 
 import com.mongodb.DBObject;
 
@@ -49,6 +51,8 @@ public abstract class AbstractMongoSessionConverter implements GenericConverter
 	private static final Log LOG = LogFactory.getLog(AbstractMongoSessionConverter.class);
 	private static final String SPRING_SECURITY_CONTEXT = "SPRING_SECURITY_CONTEXT";
 
+	private IndexResolver<MongoSession> indexResolver = new DelegatingIndexResolver<>(new PrincipalNameIndexResolver<>());
+
 	/**
 	 * Returns query to be executed to return sessions based on a particular index.
 	 * 
@@ -80,15 +84,9 @@ protected void ensureIndexes(IndexOperations sessionCollectionIndexes) {
 				.ensureIndex(new Index(EXPIRE_AT_FIELD_NAME, Sort.Direction.ASC).named(EXPIRE_AT_FIELD_NAME).expire(0));
 	}
 
-	protected String extractPrincipal(Session expiringSession) {
-
-		String resolvedPrincipal = AuthenticationParser.extractName(expiringSession.getAttribute(SPRING_SECURITY_CONTEXT));
-
-		if (resolvedPrincipal != null) {
-			return resolvedPrincipal;
-		} else {
-			return expiringSession.getAttribute(FindByIndexNameSessionRepository.PRINCIPAL_NAME_INDEX_NAME);
-		}
+	protected String extractPrincipal(MongoSession expiringSession) {
+		return this.indexResolver.resolveIndexesFor(expiringSession)
+				.get(FindByIndexNameSessionRepository.PRINCIPAL_NAME_INDEX_NAME);
 	}
 
 	public Set<ConvertiblePair> getConvertibleTypes() {
@@ -116,4 +114,8 @@ public Object convert(Object source, TypeDescriptor sourceType, TypeDescriptor t
 	protected abstract DBObject convert(MongoSession session);
 
 	protected abstract MongoSession convert(Document sessionWrapper);
+
+	public void setIndexResolver(IndexResolver<MongoSession> indexResolver) {
+		this.indexResolver = Assert.requireNonNull(indexResolver, "indexResolver must not be null!");
+	}
 }

src/main/java/org/springframework/session/data/mongo/AuthenticationParser.java

@@ -31,11 +31,13 @@
 import org.springframework.core.serializer.support.SerializingConverter;
 import org.springframework.core.type.AnnotationMetadata;
 import org.springframework.data.mongodb.core.MongoOperations;
+import org.springframework.session.IndexResolver;
 import org.springframework.session.config.SessionRepositoryCustomizer;
 import org.springframework.session.config.annotation.web.http.SpringHttpSessionConfiguration;
 import org.springframework.session.data.mongo.AbstractMongoSessionConverter;
 import org.springframework.session.data.mongo.JdkMongoSessionConverter;
 import org.springframework.session.data.mongo.MongoIndexedSessionRepository;
+import org.springframework.session.data.mongo.MongoSession;
 import org.springframework.util.StringUtils;
 import org.springframework.util.StringValueResolver;
 
@@ -57,6 +59,7 @@ public class MongoHttpSessionConfiguration extends SpringHttpSessionConfiguratio
 	private StringValueResolver embeddedValueResolver;
 	private List<SessionRepositoryCustomizer<MongoIndexedSessionRepository>> sessionRepositoryCustomizers;
 	private ClassLoader classLoader;
+	private IndexResolver<MongoSession> indexResolver;
 
 	@Bean
 	public MongoIndexedSessionRepository mongoSessionRepository(MongoOperations mongoOperations) {
@@ -66,10 +69,19 @@ public MongoIndexedSessionRepository mongoSessionRepository(MongoOperations mong
 
 		if (this.mongoSessionConverter != null) {
 			repository.setMongoSessionConverter(this.mongoSessionConverter);
+
+			if (this.indexResolver != null) {
+				this.mongoSessionConverter.setIndexResolver(this.indexResolver);
+			}
 		} else {
 			JdkMongoSessionConverter mongoSessionConverter = new JdkMongoSessionConverter(new SerializingConverter(),
 					new DeserializingConverter(this.classLoader),
 					Duration.ofSeconds(MongoIndexedSessionRepository.DEFAULT_INACTIVE_INTERVAL));
+
+			if (this.indexResolver != null) {
+				mongoSessionConverter.setIndexResolver(this.indexResolver);
+			}
+
 			repository.setMongoSessionConverter(mongoSessionConverter);
 		}
 
@@ -129,4 +141,8 @@ public void setEmbeddedValueResolver(StringValueResolver resolver) {
 		this.embeddedValueResolver = resolver;
 	}
 
+	@Autowired(required = false)
+	public void setIndexResolver(IndexResolver<MongoSession> indexResolver) {
+		this.indexResolver = indexResolver;
+	}
 }

src/main/java/org/springframework/session/data/mongo/config/annotation/web/http/MongoHttpSessionConfiguration.java

@@ -32,10 +32,12 @@
 import org.springframework.core.type.AnnotationMetadata;
 import org.springframework.data.mongodb.core.MongoOperations;
 import org.springframework.data.mongodb.core.ReactiveMongoOperations;
+import org.springframework.session.IndexResolver;
 import org.springframework.session.config.ReactiveSessionRepositoryCustomizer;
 import org.springframework.session.config.annotation.web.server.SpringWebSessionConfiguration;
 import org.springframework.session.data.mongo.AbstractMongoSessionConverter;
 import org.springframework.session.data.mongo.JdkMongoSessionConverter;
+import org.springframework.session.data.mongo.MongoSession;
 import org.springframework.session.data.mongo.ReactiveMongoSessionRepository;
 import org.springframework.util.StringUtils;
 import org.springframework.util.StringValueResolver;
@@ -58,6 +60,8 @@ public class ReactiveMongoWebSessionConfiguration extends SpringWebSessionConfig
 
 	@Autowired(required = false) private MongoOperations mongoOperations;
 	private ClassLoader classLoader;
+	private IndexResolver<MongoSession> indexResolver;
+
 
 	@Bean
 	public ReactiveMongoSessionRepository reactiveMongoSessionRepository(ReactiveMongoOperations operations) {
@@ -66,10 +70,20 @@ public ReactiveMongoSessionRepository reactiveMongoSessionRepository(ReactiveMon
 
 		if (this.mongoSessionConverter != null) {
 			repository.setMongoSessionConverter(this.mongoSessionConverter);
+
+			if (this.indexResolver != null) {
+				this.mongoSessionConverter.setIndexResolver(this.indexResolver);
+			}
+
 		} else {
 			JdkMongoSessionConverter mongoSessionConverter = new JdkMongoSessionConverter(new SerializingConverter(),
 					new DeserializingConverter(this.classLoader),
 					Duration.ofSeconds(ReactiveMongoSessionRepository.DEFAULT_INACTIVE_INTERVAL));
+
+			if (this.indexResolver != null) {
+				mongoSessionConverter.setIndexResolver(this.indexResolver);
+			}
+			
 			repository.setMongoSessionConverter(mongoSessionConverter);
 		}
 
@@ -146,4 +160,9 @@ public void setSessionRepositoryCustomizers(
 			ObjectProvider<ReactiveSessionRepositoryCustomizer<ReactiveMongoSessionRepository>> sessionRepositoryCustomizers) {
 		this.sessionRepositoryCustomizers = sessionRepositoryCustomizers.orderedStream().collect(Collectors.toList());
 	}
+
+	@Autowired(required = false)
+	public void setIndexResolver(IndexResolver<MongoSession> indexResolver) {
+		this.indexResolver = indexResolver;
+	}
 }

src/main/java/org/springframework/session/data/mongo/config/annotation/web/reactive/ReactiveMongoWebSessionConfiguration.java

@@ -33,9 +33,12 @@
 import org.springframework.data.mongodb.core.MongoOperations;
 import org.springframework.data.mongodb.core.index.IndexOperations;
 import org.springframework.mock.env.MockEnvironment;
+import org.springframework.session.IndexResolver;
 import org.springframework.session.config.SessionRepositoryCustomizer;
 import org.springframework.session.data.mongo.AbstractMongoSessionConverter;
+import org.springframework.session.data.mongo.JacksonMongoSessionConverter;
 import org.springframework.session.data.mongo.MongoIndexedSessionRepository;
+import org.springframework.session.data.mongo.MongoSession;
 import org.springframework.test.util.ReflectionTestUtils;
 
 /**
@@ -156,6 +159,32 @@ public void sessionRepositoryCustomizer() {
 		assertThat(sessionRepository).hasFieldOrPropertyWithValue("maxInactiveIntervalInSeconds", 10000);
 	}
 
+	@Test
+	void customIndexResolverConfigurationWithDefaultMongoSessionConverter() {
+
+		registerAndRefresh(MongoConfiguration.class, CustomIndexResolverConfigurationWithDefaultMongoSessionConverter.class);
+
+		MongoIndexedSessionRepository repository = this.context.getBean(MongoIndexedSessionRepository.class);
+		IndexResolver<MongoSession> indexResolver = this.context.getBean(IndexResolver.class);
+		
+		assertThat(repository).isNotNull();
+		assertThat(indexResolver).isNotNull();
+		assertThat(repository).extracting("mongoSessionConverter").hasFieldOrPropertyWithValue("indexResolver", indexResolver);
+	}
+
+	@Test
+	void customIndexResolverConfigurationWithProvidedMongoSessionConverter() {
+
+		registerAndRefresh(MongoConfiguration.class, CustomIndexResolverConfigurationWithProvidedMongoSessionConverter.class);
+
+		MongoIndexedSessionRepository repository = this.context.getBean(MongoIndexedSessionRepository.class);
+		IndexResolver<MongoSession> indexResolver = this.context.getBean(IndexResolver.class);
+
+		assertThat(repository).isNotNull();
+		assertThat(indexResolver).isNotNull();
+		assertThat(repository).extracting("mongoSessionConverter").hasFieldOrPropertyWithValue("indexResolver", indexResolver);
+	}
+
 	private void registerAndRefresh(Class<?>... annotatedClasses) {
 
 		this.context.register(annotatedClasses);
@@ -264,4 +293,30 @@ public SessionRepositoryCustomizer<MongoIndexedSessionRepository> sessionReposit
 		}
 	}
 
+	@Configuration
+	@EnableMongoHttpSession
+	static class CustomIndexResolverConfigurationWithDefaultMongoSessionConverter {
+
+		@Bean
+		@SuppressWarnings("unchecked")
+		public IndexResolver<MongoSession> indexResolver() {
+			return mock(IndexResolver.class);
+		}
+	}
+
+	@Configuration
+	@EnableMongoHttpSession
+	static class CustomIndexResolverConfigurationWithProvidedMongoSessionConverter {
+
+		@Bean
+		public AbstractMongoSessionConverter mongoSessionConverter() {
+			return new JacksonMongoSessionConverter();
+		}
+
+		@Bean
+		@SuppressWarnings("unchecked")
+		public IndexResolver<MongoSession> indexResolver() {
+			return mock(IndexResolver.class);
+		}
+	}
 }