STOMP Interceptor 예외를 처리해주자

원래 STOMP는 예외가 클라이언트로 전달이 안되는 줄 알았는데
이번에 채팅을 연결하면서 클라이언트 개발자와 이야기를 하면서
이게 스프링 시큐리티처럼 예외를 catch를 하는 바람에 클라이언트로 전달이 안된다는 걸 알았다.

스프링 공식 깃허브에 관련된 코드들이 있어서 활용한다.

@Component
public class StompErrorHandler extends StompSubProtocolErrorHandler {

    private static final byte[] EMPTY_PAYLOAD = new byte[0];

    @Override
    public Message<byte[]> handleClientMessageProcessingError(Message<byte[]> clientMessage, Throwable ex) {
        StompHeaderAccessor accessor = StompHeaderAccessor.create(StompCommand.ERROR);
        accessor.setMessage(ex.getMessage());
        accessor.setLeaveMutable(true);

        StompHeaderAccessor clientHeaderAccessor = null;
        if (clientMessage != null) {
            clientHeaderAccessor = MessageHeaderAccessor.getAccessor(clientMessage, StompHeaderAccessor.class);
            if (clientHeaderAccessor != null) {
                String receiptId = clientHeaderAccessor.getReceipt();
                if (receiptId != null) {
                    accessor.setReceiptId(receiptId);
                }
            }
        }

        return handleInternal(accessor, EMPTY_PAYLOAD, ex, clientHeaderAccessor);
    }

    protected Message<byte[]> handleInternal(StompHeaderAccessor errorHeaderAccessor, byte[] errorPayload,
                                             Throwable cause, StompHeaderAccessor clientHeaderAccessor) {

        String errorCause = "";
        if (cause != null) errorCause = cause.getCause().toString();

        String fullErrorMessage = extractErrorCode(errorCause);
        byte[] newPayload = fullErrorMessage.getBytes(StandardCharsets.UTF_8);
        return MessageBuilder.createMessage(newPayload, errorHeaderAccessor.getMessageHeaders());
    }


    private String extractErrorCode(String input) {
        String[] parts = input.split(":");
        if (parts.length > 1) {
            return parts[1].trim();
        } else {
            // ':' 문자가 없는 경우 정의된 에러 메세지 반환
            return "Undefined exception";
        }
    }

이렇게 핸들러를 등록해주어야 한다.

그럼 이렇게 커스텀 예외의 메시지가 전송된다.
지금은 상태 코드가 없으니
상태 코드도 함께 첨부해주자.

   protected Message<byte[]> handleInternal(StompHeaderAccessor errorHeaderAccessor, byte[] errorPayload,
                                             Throwable cause, StompHeaderAccessor clientHeaderAccessor) {

        String errorCause = "";
        if (cause != null) {
            CustomException customException = (CustomException) cause.getCause();
            errorCause = ErrorResponse.toResponse(customException).toString();

        }
        byte[] newPayload = errorCause.getBytes(StandardCharsets.UTF_8);
        return MessageBuilder.createMessage(newPayload, errorHeaderAccessor.getMessageHeaders());
    }
    

메시지를 커스텀 예외로 변형시켜서 보내면

이렇게 예외 메시지를 클라이언트에서 받을 수 있다.