GH-8600: Fix WebSocket removeRegistration

Fixes #8600

When we register a dynamic WebSocket endpoint and use a `WebSocketHandlerDecoratorFactory`
such an endpoint is not removed on an `IntegrationFlow` destruction.
The actual `WebSocketHandler` is decorated, however we still use an initial one
for condition.

* Refactor `IntegrationWebSocketContainer` to expose a `protected` setter for the
`WebSocketHandler` which is called from the `ServerWebSocketContainer` after decoration

**Cherry-pick to `6.0.x` & `5.5.x`**

spring-integration-websocket/src/main/java/org/springframework/integration/websocket/ClientWebSocketContainer.java

@@ -208,7 +208,7 @@ public void stop(Runnable callback) {
 	 * <p>
 	 * Opened {@link WebSocketSession} is populated to the wrapping {@link ClientWebSocketContainer}.
 	 * <p>
-	 * The {@link #webSocketHandler} is used to handle {@link WebSocketSession} events.
+	 * The {@link #getWebSocketHandler()} is used to handle {@link WebSocketSession} events.
 	 */
 	private final class IntegrationWebSocketConnectionManager extends ConnectionManagerSupport {
 
@@ -258,8 +258,7 @@ protected void openConnection() {
 			}
 			ClientWebSocketContainer.this.headers.setSecWebSocketProtocol(getSubProtocols());
 			CompletableFuture<WebSocketSession> future =
-					this.client.execute(ClientWebSocketContainer.this.webSocketHandler,
-							ClientWebSocketContainer.this.headers, getUri());
+					this.client.execute(getWebSocketHandler(), ClientWebSocketContainer.this.headers, getUri());
 
 			future.whenComplete((session, throwable) -> {
 				if (throwable == null) {

spring-integration-websocket/src/main/java/org/springframework/integration/websocket/IntegrationWebSocketContainer.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2014-2022 the original author or authors.
+ * Copyright 2014-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -67,8 +67,7 @@ public abstract class IntegrationWebSocketContainer implements DisposableBean {
 
 	protected final Log logger = LogFactory.getLog(getClass()); // NOSONAR
 
-	protected final WebSocketHandler webSocketHandler = new IntegrationWebSocketHandler(); // NOSONAR
-
+	private WebSocketHandler webSocketHandler = new IntegrationWebSocketHandler();
 	protected final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>(); // NOSONAR
 
 	private final List<String> supportedProtocols = new ArrayList<>();
@@ -104,6 +103,15 @@ public void addSupportedProtocols(String... protocols) {
 		}
 	}
 
+	/**
+	 * Replace a default {@link WebSocketHandler} with provided, e.g. via decoration factories.
+	 * @param handler the actual {@link WebSocketHandler} to replace.
+	 * @since 5.5.18
+	 */
+	protected void setWebSocketHandler(WebSocketHandler handler) {
+		this.webSocketHandler = handler;
+	}
+
 	public WebSocketHandler getWebSocketHandler() {
 		return this.webSocketHandler;
 	}

spring-integration-websocket/src/main/java/org/springframework/integration/websocket/ServerWebSocketContainer.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2014-2022 the original author or authors.
+ * Copyright 2014-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -150,11 +150,12 @@ public TaskScheduler getSockJsTaskScheduler() {
 
 	@Override
 	public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
-		WebSocketHandler webSocketHandler = this.webSocketHandler;
+		WebSocketHandler webSocketHandler = getWebSocketHandler();
 
 		if (this.decoratorFactories != null) {
 			for (WebSocketHandlerDecoratorFactory factory : this.decoratorFactories) {
 				webSocketHandler = factory.decorate(webSocketHandler);
+				setWebSocketHandler(webSocketHandler);
 			}
 		}
 

spring-integration-websocket/src/test/java/org/springframework/integration/websocket/dsl/WebSocketDslTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2021-2022 the original author or authors.
+ * Copyright 2021-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,6 +16,8 @@
 
 package org.springframework.integration.websocket.dsl;
 
+import java.util.concurrent.atomic.AtomicReference;
+
 import jakarta.websocket.DeploymentException;
 import org.junit.jupiter.api.Test;
 
@@ -37,6 +39,7 @@
 import org.springframework.test.annotation.DirtiesContext;
 import org.springframework.test.context.junit.jupiter.SpringJUnitConfig;
 import org.springframework.web.context.support.AnnotationConfigWebApplicationContext;
+import org.springframework.web.socket.WebSocketHandler;
 import org.springframework.web.socket.client.WebSocketClient;
 import org.springframework.web.socket.client.standard.StandardWebSocketClient;
 import org.springframework.web.socket.server.HandshakeHandler;
@@ -46,6 +49,9 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
 import static org.awaitility.Awaitility.await;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
 
 @SpringJUnitConfig(classes = WebSocketDslTests.ClientConfig.class)
 @DirtiesContext
@@ -61,13 +67,19 @@ public class WebSocketDslTests {
 	IntegrationFlowContext integrationFlowContext;
 
 	@Test
-	void testDynamicServerEndpointRegistration() {
+	void testDynamicServerEndpointRegistration() throws Exception {
 		// Dynamic server flow
 		AnnotationConfigWebApplicationContext serverContext = this.server.getServerContext();
 		IntegrationFlowContext serverIntegrationFlowContext = serverContext.getBean(IntegrationFlowContext.class);
+		AtomicReference<WebSocketHandler> decoratedHandler = new AtomicReference<>();
 		ServerWebSocketContainer serverWebSocketContainer =
 				new ServerWebSocketContainer("/dynamic")
 						.setHandshakeHandler(serverContext.getBean(HandshakeHandler.class))
+						.setDecoratorFactories(handler -> {
+							WebSocketHandler spy = spy(handler);
+							decoratedHandler.set(spy);
+							return spy;
+						})
 						.withSockJs();
 
 		WebSocketInboundChannelAdapter webSocketInboundChannelAdapter =
@@ -106,6 +118,8 @@ void testDynamicServerEndpointRegistration() {
 				.extracting(Message::getPayload)
 				.isEqualTo("dynamic test");
 
+		verify(decoratedHandler.get()).handleMessage(any(), any());
+
 		dynamicServerFlow.destroy();
 
 		await() // Looks like endpoint is removed on the server side somewhat async