Handle a MethodArgumentNotValidException from a WebSocket message

Im using Spring WebSocket and I have a Model annotated with validation annotations and want to capture the validation message it has before saving. Consider this @Controller I am testing:

@Controller
public class MessageThreadController {
    /* A plain XHR post request with a JSON body to match MessageThread. */
    @RequestMapping(
        headers={"Accept=application/hal+json", "Content-Type=application/json", "X-Requested-With=XMLHttpRequest"},
        method={RequestMethod.POST},
        produces="application/hal+json",
        value="/api/unsecure/thread/create")
    public HttpEntity createPost(@Valid @RequestBody MessageThread messageThread) throws Exception {
        return ResponseEntity.ok("valid");
    }

    /* The same request above but for WebSocket.*/
    @MessageMapping("/api/secure/thread/create")
    @SendToUser("/api/secure/broadcast")
    public HttpEntity createMessage(@Valid @RequestBody MessageThread messageThread) throws Exception {
        return ResponseEntity.ok("valid");
    }
}

And I have a @ControllerAdvice to capture the Exceptions that may happen:

@ControllerAdvice
public class MessageThreadControllerAdvice {
    private static final Logger log = LoggerFactory.getLogger(MessageThreadControllerAdvice.class);

    /* MessageThreadController.createPost() validation exception are handled here. */
    @ExceptionHandler(org.springframework.web.bind.MethodArgumentNotValidException.class)
    public void methodArgumentNotValidExceptionHandler(
        org.springframework.web.bind.MethodArgumentNotValidException e) {
        log.trace("methodArgumentNotValidExceptionHandler");
    }

    /* MessageThreadController.createMessage() validation exception are handled here. */
    @ExceptionHandler(org.springframework.messaging.handler.annotation.support.MethodArgumentNotValidException.class)
    @SendToUser("/api/secure/broadcast")
    public void methodArgumentNotValidWebSocketExceptionHandler(
        org.springframework.messaging.handler.annotation.support.MethodArgumentNotValidException e) {
        log.trace("methodArgumentNotValidWebSocketExceptionHandler");
    }

    /* All exceptions are handled here. */
    @ExceptionHandler
    @SendToUser("/api/secure/broadcast")
    public void exceptionHandler(Exception e) {
        log.trace("exceptionHandler");
    }
}

And here is the object in question:

@Getter @Setter
public class MessageThread {
    @Size(min=1, message="{validation.userIds.empty}")
    @NotNull(message="{validation.userIds.null}")
    private Set<BigInteger> userIds;
}

If I were to send the following JSON body in both MessageThreadController.createPost() and MessageThreadController.createMessage():

{"userIds":[]}

the validation kicks in. The problem is that only the exception from MessageThreadController.createPost() gets handled while the exception from MessageThreadController.createMessage() is not handled:

>2020-08-21 11:54:00.651 ERROR 4444 --- [nboundChannel-4] .WebSocketAnnotationMethodMessageHandler : Unhandled exception from message handler method
>
> org.springframework.messaging.handler.annotation.support.MethodArgumentNotValidException: Could not resolve method parameter at index 0 in public void controller.MessageThreadController.createMessage(model.form.MessageThread) throws java.lang.Exception: 1 error(s): [Field error in object 'messageThread' on field 'userIds': rejected value [[]]; codes [Size.messageThread.userIds,Size.userIds,Size.java.util.Set,Size]; arguments [org.springframework.context.support.DefaultMessageSourceResolvable: codes [messageThread.userIds,userIds]; arguments []; default message [userIds],2147483647,1]; default message [size must be between 1 and 2147483647]]

How can I handle the MethodArgumentNotValidException that is thrown on a websocket message?

You can annotate your @ControllerAdvice method with the @MessageExceptionHandler annotation.

@ControllerAdvice
public class MessageThreadControllerAdvice {
    private static final Logger log = LoggerFactory.getLogger(MessageThreadControllerAdvice.class);

    /* MessageThreadController.createPost() validation exception are handled here. */
    @ExceptionHandler(org.springframework.web.bind.MethodArgumentNotValidException.class)
    public void methodArgumentNotValidExceptionHandler(
        org.springframework.web.bind.MethodArgumentNotValidException e) {
        log.trace("methodArgumentNotValidExceptionHandler");
    }

    /* MessageThreadController.createMessage() validation exception are handled here. */
    @MessageExceptionHandler(org.springframework.messaging.handler.annotation.support.MethodArgumentNotValidException.class)
    @SendToUser("/api/secure/broadcast")
    public void methodArgumentNotValidWebSocketExceptionHandler(
        org.springframework.messaging.handler.annotation.support.MethodArgumentNotValidException e) {
        log.trace("methodArgumentNotValidWebSocketExceptionHandler");
    }

    /* All exceptions are handled here. */
    @ExceptionHandler
    @SendToUser("/api/secure/broadcast")
    public void exceptionHandler(Exception e) {
        log.trace("exceptionHandler");
    }
}

This annotation could be also applied at the @Controller level.

Please, be aware that this annotation is supported at @ControllerAdvice level since Spring 4.2.