Can't unregister websocket when using decorators on ServerWebSocketContainer

[aperinot]

Impacted artifact / version
spring-integration-websocket version 5.5.16

Even if not tested, the issue seems to be present in 5.5.17, and latest 6.X.X branch

Describe the bug
When adding decorator factories to the ServerWebSocketContainer, it is not possible to unregister the websocket endpoint.

When registering, the ServerWebSocketContainer code is as follow:

       @Override
	public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
		WebSocketHandler webSocketHandler = this.webSocketHandler;

		if (this.decoratorFactories != null) {
			for (WebSocketHandlerDecoratorFactory factory : this.decoratorFactories) {
				webSocketHandler = factory.decorate(webSocketHandler);
			}
		}

		WebSocketHandlerRegistration registration = registry.addHandler(webSocketHandler, this.paths)
				.setHandshakeHandler(this.handshakeHandler)
				.addInterceptors(this.interceptors)
				.setAllowedOrigins(this.origins);

		configureSockJsOptionsIfAny(registration);
	}

As you can see, the addHandler is called with the latest returned webSocketHandler from the list of decoratorFactories.
If no decoratorFactories are set, the websocketHandler used is the default IntegrationWebSocketHandler.

However, the unregister code in IntegrationServletWebSocketHandlerRegistry is as follow:

      void removeRegistration(ServerWebSocketContainer serverWebSocketContainer) {
		List<String> patterns = this.dynamicRegistrations.remove(serverWebSocketContainer.getWebSocketHandler());
		if (this.dynamicHandlerMapping != null && !CollectionUtils.isEmpty(patterns)) {
			for (String pattern : patterns) {
				this.dynamicHandlerMapping.unregisterHandler(pattern);
			}
		}
	}

It tries to remove from the dynamicRegistrations by retrieving the handler with serverWebSocketContainer.getWebSocketHandler().
However the getWebSocketHandler() returns the default IntegrationWebSocketHandler, and not the decorated one, so the remove doesn't work, and it doesn't call the unregisterHandler() for the pattern.

To Reproduce

  @Autowired
  private IntegrationFlowContext integrationFlowContext;
 
  public void bugReproducer() {
    ServerWebSocketContainer serverWebSocketContainer = new ServerWebSocketContainer("/path")
        .setHandshakeHandler(new DefaultHandshakeHandler());

   // Adding a decorator triggers the bug
    serverWebSocketContainer.setDecoratorFactories(WebsocketHandlerDecorator::new);

   // The rest is some kind of boiler code, however i'm not sure how to make it shorter
    WebSocketOutboundMessageHandler webSocketOutboundMessageHandler = new WebSocketOutboundMessageHandler(
        serverWebSocketContainer);
    webSocketOutboundMessageHandler.afterPropertiesSet();

    StandardIntegrationFlow standardIntegrationFlow = IntegrationFlows.from(new DirectChannel())
        .split(new CustomSplitter(serverWebSocketContainer)).handle(webSocketOutboundMessageHandler).get();

    IntegrationFlowContext.IntegrationFlowRegistration dynamicServerFlow = integrationFlowContext
        .registration(standardIntegrationFlow).id("path").addBean(serverWebSocketContainer).register();
   
    // This call should kill the websocket endpoint previously registered
    dynamicServerFlow.destroy();
   }

  // Simple decorator for the default WebSocketHandler
private static class WebsocketHandlerDecorator implements WebSocketHandler
{
  private WebSocketHandler handler;

  private WebsocketHandlerDecorator(WebSocketHandler handler)
  {
    this.handler = handler;
  }

  @Override
  public void afterConnectionEstablished(WebSocketSession session) throws Exception
  {
    handler.afterConnectionEstablished(session);
  }

  @Override
  public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception
  {
    handler.handleMessage(session, message);
  }

  @Override
  public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception
  {
    handler.handleTransportError(session, exception);
  }

  @Override
  public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception
  {
    handler.afterConnectionClosed(session, closeStatus);
  }

  @Override
  public boolean supportsPartialMessages()
  {
    return handler.supportsPartialMessages();
  }
}

Expected behavior

The expected behavior shall be the unregistration of the endpoint.

Workaround

Here is a workaround to make it work as expected.

private static class CustomServerWebSocketContainer extends ServerWebSocketContainer
  {

    private WebSocketHandler lastHandler;

    public CustomServerWebSocketContainer(String... paths)
    {
      super(paths);
    }

    @Override
    public ServerWebSocketContainer setDecoratorFactories(WebSocketHandlerDecoratorFactory... factories)
    {

      WebSocketHandlerDecoratorFactory[] factoriesDecorator = Arrays.copyOf(factories, factories.length);

      if (factoriesDecorator.length > 0)
      {
      // We must decorate the decorator to save the last returned handler, so the getWebSocketHandler() returns that handler
        WebSocketHandlerDecoratorFactory lastFactory = factoriesDecorator[factoriesDecorator.length - 1];
        factoriesDecorator[factoriesDecorator.length - 1] = w ->
          {
            lastHandler = lastFactory.decorate(w);
            return lastHandler;
          };
      }

      return super.setDecoratorFactories(factoriesDecorator);
    }

    @Override
    public WebSocketHandler getWebSocketHandler()
    {
      if (lastHandler != null)
      {
        return lastHandler;
      }
      
      return super.getWebSocketHandler();
    }
  }