From 456d6e78fe07968a8157e639a3ab5fff70487148 Mon Sep 17 00:00:00 2001
From: Aarti Gupta <aartiguptaa@gmail.com>
Date: Fri, 17 Apr 2020 16:03:28 -0700
Subject: [PATCH] Add support for customizing RSocketMessageHandler

See gh-21081
---
 .../RSocketMessageHandlerCustomizer.java       | 18 ++++++++++++++++++
 .../RSocketMessagingAutoConfiguration.java     | 14 +++++++++++---
 2 files changed, 29 insertions(+), 3 deletions(-)
 create mode 100644 spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/rsocket/RSocketMessageHandlerCustomizer.java

diff --git a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/rsocket/RSocketMessageHandlerCustomizer.java b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/rsocket/RSocketMessageHandlerCustomizer.java
new file mode 100644
index 0000000000..978dffd5db
--- /dev/null
+++ b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/rsocket/RSocketMessageHandlerCustomizer.java
@@ -0,0 +1,18 @@
+package org.springframework.boot.autoconfigure.rsocket;
+
+import org.apache.catalina.connector.Connector;
+import org.springframework.boot.web.embedded.tomcat.ConfigurableTomcatWebServerFactory;
+import org.springframework.messaging.rsocket.RSocketStrategies;
+import org.springframework.messaging.rsocket.annotation.support.RSocketMessageHandler;
+import org.springframework.util.RouteMatcher;
+
+/**
+ * @author Aarti Gupta
+ * Callback interface that can be used to customize a RSocketMessageHandler {@link Connector}.
+ */
+@FunctionalInterface
+public interface RSocketMessageHandlerCustomizer {
+
+	RSocketMessageHandler setRouteMatcher(RouteMatcher handler);
+
+}
diff --git a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/rsocket/RSocketMessagingAutoConfiguration.java b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/rsocket/RSocketMessagingAutoConfiguration.java
index 09a3898180..d73fc35bf7 100644
--- a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/rsocket/RSocketMessagingAutoConfiguration.java
+++ b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/rsocket/RSocketMessagingAutoConfiguration.java
@@ -19,6 +19,9 @@ package org.springframework.boot.autoconfigure.rsocket;
 import io.rsocket.RSocketFactory;
 import io.rsocket.transport.netty.server.TcpServerTransport;
 
+import java.util.Objects;
+import java.util.stream.Collectors;
+import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -28,6 +31,7 @@ import org.springframework.context.annotation.Configuration;
 import org.springframework.messaging.rsocket.RSocketRequester;
 import org.springframework.messaging.rsocket.RSocketStrategies;
 import org.springframework.messaging.rsocket.annotation.support.RSocketMessageHandler;
+import org.springframework.util.RouteMatcher;
 
 /**
  * {@link EnableAutoConfiguration Auto-configuration} for Spring RSocket support in Spring
@@ -37,16 +41,20 @@ import org.springframework.messaging.rsocket.annotation.support.RSocketMessageHa
  * @since 2.2.0
  */
 @Configuration(proxyBeanMethods = false)
-@ConditionalOnClass({ RSocketRequester.class, RSocketFactory.class, TcpServerTransport.class })
+@ConditionalOnClass({RSocketRequester.class, RSocketFactory.class, TcpServerTransport.class})
 @AutoConfigureAfter(RSocketStrategiesAutoConfiguration.class)
 public class RSocketMessagingAutoConfiguration {
 
+
 	@Bean
 	@ConditionalOnMissingBean
-	public RSocketMessageHandler messageHandler(RSocketStrategies rSocketStrategies) {
+	public RSocketMessageHandler messageHandler(RSocketStrategies rSocketStrategies, ObjectProvider<RSocketMessageHandlerCustomizer> customizers) {
 		RSocketMessageHandler messageHandler = new RSocketMessageHandler();
 		messageHandler.setRSocketStrategies(rSocketStrategies);
-		return messageHandler;
+		RSocketMessageHandlerCustomizer rSocketMessageHandlerCustomizer = customizers.getIfAvailable();
+		return rSocketMessageHandlerCustomizer.setRouteMatcher(rSocketStrategies.routeMatcher());
 	}
 
+
+
 }