How to globally handle Spring WebSockets/Spring Messaging exception?

Question

Question
Is there a way to globally handle Spring Messaging MessageDeliveryException caused by error (usualy insufficient authorities) in Spring WebSocket module?

Use case
I have implemented Spring WebSockets over STOMP to support ws connection in my webapp. To secure websocket endpoint I have created interceptor that authorizes user to start STOMP session at STOMP CONNECT time (as suggested in Spring documentation here in 22.4.11 section):

@Component
public class StompMessagingInterceptor extends ChannelInterceptorAdapter {

    // Some code not important to the problem

    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor headerAccessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);

        switch (headerAccessor.getCommand()) {
            // Authenticate STOMP session on CONNECT using jwt token passed as a STOMP login header - it's working great
            case CONNECT:
                authorizeStompSession(headerAccessor);
                break;
        }

        // Returns processed message
        return message;
    }

    // Another part of code not important for the problem
}

and included spring-security-messaging configuration to add some fine-grained control over authorities when messaging:

@Configuration
public class WebSocketSecurityConfig extends AbstractSecurityWebSocketMessageBrokerConfigurer {

    @Override
    protected void configureInbound(MessageSecurityMetadataSourceRegistry messages) {
        messages
            .simpTypeMatchers(
                SimpMessageType.CONNECT,
                SimpMessageType.DISCONNECT,
                SimpMessageType.HEARTBEAT
            ).authenticated()
            .simpSubscribeDestMatchers("/queue/general").authenticated()
            .simpSubscribeDestMatchers("/user/queue/priv").authenticated()
            .simpDestMatchers("/app/general").authenticated()
            .simpDestMatchers("/user/*/queue/priv").hasAuthority("ADMIN")
            .anyMessage().denyAll();
    }

    @Override
    protected boolean sameOriginDisabled() {
        return true;
    }
}

First of all - this configuration works as expected, the problem is when some security exception happens during websocket communication (say user without admin authority tries to send message on "/user/{something}/queue/priv" endpoint) it will end in org.springframework.messaging.MessageDeliveryException being rised and:

• Full exception stack trace being written down to my server log
• Returning STOMP ERROR frame containing part of stack trace as it's message field.

What I would like to do is catching (if possible globally) DeliveryException, checking what caused it and accoridingly to that create my own message for returning in STOMP ERROR frame (lets say with some error code like just 403 to mimic HTTP) and instead of throwing original exception further just logging some warning with my logger. Is it possible?

What I tried
When looking for solution I found some people using @MessageExceptionHandler to catch messaging exceptions, Spring 4.2.3 (which is version I use) documentation mentions it only once here in 25.4.11 section. I tried to use it like this:

@Controller
@ControllerAdvice
public class WebSocketGeneralController {

    ...

    @MessageExceptionHandler
    public WebSocketMessage handleException(org.springframework.messaging.MessageDeliveryException e) {
        WebSocketMessage errorMessage = new WebSocketMessage();
        errorMessage.setMessage(e.getClass().getName());
        return errorMessage;
    }
}

but it seems like method isn't called at any point (tried catching different exceptions, just Exception including - no results). What else should I look into?

Answer 1

@ControllerAdvice and @MessageExceptionHandler are working on business-logic level (like @MessageMapping or SimpMessagingTemplate).

To handle STOMP exceptions, you need to set STOMP error handler in STOMP registry:

@Configuration
@EnableWebSocketMessageBroker
class WebSocketConfiguration : WebSocketMessageBrokerConfigurer {

    override fun configureMessageBroker(registry: MessageBrokerRegistry) {
        // ...
    }

    override fun registerStompEndpoints(registry: StompEndpointRegistry) {
        registry.addEndpoint("/ws")

        // Handle exceptions in interceptors and Spring library itself.
        // Will terminate a connection and send ERROR frame to the client.
        registry.setErrorHandler(object : StompSubProtocolErrorHandler() {
            override fun handleInternal(
                errorHeaderAccessor: StompHeaderAccessor,
                errorPayload: ByteArray,
                cause: Throwable?,
                clientHeaderAccessor: StompHeaderAccessor?
            ): Message<ByteArray> {
                errorHeaderAccessor.message = null
                val message = "..."
                return MessageBuilder.createMessage(message.toByteArray(), errorHeaderAccessor.messageHeaders)
            }
        })
    }
}