Customize error message in stomp websocket spring boot 2

Question

I use Spring Boot 2 WebSocket + SockJS + STOMP.

I have JWT authorization when client is connecting in ChannelInterceptor in preSend method.

@Component
public class WebSocketChannelInterceptor implements ChannelInterceptor {

    private final TokenUtils tokenUtils;

    private final AuthenticationManager websocketAuthenticationManager;

    private static final Logger LOGGER = LoggerFactory.getLogger(WebSocketChannelInterceptor.class);

    public WebSocketChannelInterceptor(TokenUtils tokenUtils, AuthenticationManager websocketAuthenticationManager) {
        this.tokenUtils = tokenUtils;
        this.websocketAuthenticationManager = websocketAuthenticationManager;
    }

    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {

        LOGGER.info("WEBSOCKETCHANNELINTERCEPTOR -> "+message.toString());

        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);

        if (accessor != null && StompCommand.CONNECT.equals(accessor.getCommand())) {
            List<String> headers = accessor.getNativeHeader(AUTHORIZATION);
            accessor.setUser(websocketAuthenticationManager.authenticate(new JWTTokenAuthentication(tokenUtils.resolveToken(headers != null ? headers.get(0) : null))));
        }

        return message;
    }

    @Override
    public boolean preReceive(MessageChannel channel) {
        LOGGER.info("preReceive");
        return true;
    }

}

AuthenticationManager:

@Component("websocketAuthenticationManager")
@AllArgsConstructor
public class WebsocketAuthenticationManager implements AuthenticationManager {

    private final TokenUtils tokenUtils;

    private final ClientRepository clientRepository;

    private final OperatorRepository operatorRepository;

    @Override
    @Transactional(readOnly = true)
    public Authentication authenticate(Authentication authentication) throws AuthenticationException {

        JWTTokenAuthentication jwtTokenAuthentication = (JWTTokenAuthentication) authentication;

        String token = jwtTokenAuthentication.getToken();

        if(tokenUtils.validateToken(token)) {

            RoleType role = RoleType.get(tokenUtils.get("role", token));

            if(role == null) {
                throw throwBadCredentialsException("Unknown role");
            }

            switch (role) {
                case OPERATOR:

                    Optional<OperatorEntity> operator = operatorRepository.findByLogin(tokenUtils.get("login", token));

                    if(operator.isPresent()) {
                        return new WebsocketOperatorDetails(token, OperatorDetails.builder().role(role.toString()).operator(operator.get()).build());
                    } else {
                        throw throwBadCredentialsException("Operator not found");
                    }

                case CLIENT:

                    Optional<ClientEntity> client = clientRepository.findByHash(tokenUtils.get("hash", token));

                    if(client.isPresent()) {
                        return new WebsocketClientDetails(token, ClientDetails.builder().role(role.toString()).client(client.get()).build());
                    } else {
                        throw throwBadCredentialsException("Client not found");
                    }
                case OWNER:
                    throw throwBadCredentialsException("Access denied");
                default:
                    throw throwBadCredentialsException("Unknown role");
            }

        } else {
            throw throwBadCredentialsException("Invalid token");
        }

    }

    private MessagingException throwBadCredentialsException(String message) {

        Map<String, Object> headers = new HashMap<>();

        headers.put("status", StatusType.UNAUTHORIZED);
        headers.put("message", message);

        return new MessagingException(new MutableMessage<>(new MessageHeaders(headers)));
    }

}

In AuthenticationManager i trowing error if somthing wrong.

On client side i get next content:

{  
   "command":"ERROR",
   "headers":{  
      "content-length":"0",
      "message":"Failed to send message to ExecutorSubscribableChannel[clientInboundChannel]; nested exception is org.springframework.security.authentication.AuthenticationCredentialsNotFoundException\\c Invalid token"
   },
   "_binaryBody":{  

   },
   "isBinaryBody":true,
   "escapeHeaderValues":false,
   "skipContentLengthHeader":false,
   "_body":""
}

How change the code to send error message with payload?

Thanks.

Answer 1

There is a dedicated setter 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)
            }
        })
    }
}

Answer 2

You can construct a MutableMessage with payload and headers:

new MutableMessage<String>("payload", headers);

Look the constructor inside class MutableMessage:

public MutableMessage(T payload, Map<String, Object> headers) {
    Assert.notNull(payload, "payload must not be null");
    this.payload = payload;

    this.headers = new MutableMessageHeaders(headers);

    if (headers != null) {
        this.headers.put(MessageHeaders.ID, headers.get(MessageHeaders.ID));
        this.headers.put(MessageHeaders.TIMESTAMP, headers.get(MessageHeaders.TIMESTAMP));
    }
}