Spring WebSocket 예외 처리 - @MessageExceptionHandler, StompSubProtocolErrorHandler

서론

Spring WebSockets 라이브러리를 사용하면서 예외 발생시 클라이언트에게 적절히 응답을 보내줄 필요가 있었습니다.

 

그래서 공식문서, 깃허브, StackOverFlow 등 여러 방법을 찾아 예외 처리 코드를 구현했습니다.

 

레퍼런스가 많지 않아서 웹소켓 예외를 처리하는 적절한 방법을 정리해보았습니다.

 

웹소켓에서 발생하는 예외

웹소켓에서 발생하는 예외는 크게 초기 연결 시 발생하는 인증 예외와 비즈니스 로직에 대한 검증 예외로 나눌 수 있습니다.

 

초기 연결시 발생하는 예외 - 인터셉터에서 발생하는 예외 처리

웹소켓 초기 연결시 인터셉터에서 클라이언트로부터 전달받은 토큰을 검증하고, 토큰이 유효하지 않을 경우나 만료될 경우 예외를 발생시킵니다. 

(더 자세한 내용은 아래 글에서 확인해보세요 ! )

https://shout-to-my-mae.tistory.com/430

Spring WebSocket 애플리케이션에 Spring Security 적용하기(simpUser, Interceptor, handler) - 0518 업데이트

 

인터셉터 예시 코드

@Component
@RequiredArgsConstructor
public class JwtChannelInterceptor implements ChannelInterceptor {

    private final AuthService authService;

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

        if (accessor != null && StompCommand.CONNECT.equals(accessor.getCommand())) {
            String jwtToken = extractJwtToken(accessor);
            if (jwtToken == null) {
                return message;
            }

            Authentication authentication = authService.createAuthenticationByToken(jwtToken);
            accessor.setUser(authentication);
        }

        return message;
    }
    ...
 }

 

 

preSend 메서드에서 웹소켓 연결 시 전달되는 메시지를 가로채고, StompHeaderAccessor를 사용하여 헤더 정보에 접근합니다.

이후 JWT 토큰을 추출하고 인증 서비스(AuthService)를 통해 토큰을 검증합니다.

검증이 성공하면 인증 정보를 설정하고, 실패할 경우 예외가 발생합니다.

 

인터셉터에서 발생하는 예외 처리 - 에러 핸들러 구현

@Slf4j
@RequiredArgsConstructor
@Configuration
public class StompErrorHandler extends StompSubProtocolErrorHandler {

    private final ObjectMapper objectMapper;

    @Override
    public Message<byte[]> handleClientMessageProcessingError(Message<byte[]> clientMessage, Throwable ex) {
        if (ex instanceof MessageDeliveryException) {
            Throwable cause = ex.getCause();
            if (cause instanceof AccessDeniedException) {
                return sendErrorMessage(new ErrorResponse(1201, "Access denied"));
            }

            if (cause instanceof InvalidAuthorizationHeaderException) {
                return sendErrorMessage(new ErrorResponse(1003, cause.getMessage()));
            }

            if (isJwtException(cause)) {
                return sendErrorMessage(new ErrorResponse(1201, cause.getMessage()));
            }
        }
        return super.handleClientMessageProcessingError(clientMessage, ex);
    }

    private boolean isJwtException(Throwable ex) {
        return ex instanceof JwtSignatureException
                || ex instanceof JwtExpiredException
                || ex instanceof JwtUnsupportedJwtException
                || ex instanceof JwtIllegalArgumentException;
    }

    private Message<byte[]> sendErrorMessage(ErrorResponse errorResponse) {
        StompHeaderAccessor headers = StompHeaderAccessor.create(StompCommand.ERROR);
        headers.setMessage(errorResponse.message());
        headers.setLeaveMutable(true);

        try {
            String json = objectMapper.writeValueAsString(errorResponse);
            return MessageBuilder.createMessage(json.getBytes(StandardCharsets.UTF_8),
                    headers.getMessageHeaders());
        } catch (JsonProcessingException e) {
            log.error("Failed to convert ErrorResponse to JSON", e);
            return MessageBuilder.createMessage(errorResponse.message().getBytes(StandardCharsets.UTF_8),
                    headers.getMessageHeaders());
        }
    }
}

 

인터셉터에서 발생하는 예외를 처리하기 위해 StompSubProtocolErrorHandler를 상속받아 구현했습니다.

 

MessageDeliveryException이 발생한 경우, 원인 예외를 확인하여 적절한 에러 메시지를 반환합니다.

접근 거부 예외(AccessDeniedException)나 잘못된 인증 헤더 예외(InvalidAuthorizationHeaderException), JWT 관련 예외 등을 처리합니다.

 

StompSubProtocolErrorHandler

StompSubProtocolErrorHandler는 인터셉터에서 발생한 에러를 포함하여 STOMP 메시지 처리 과정에서 발생하는 모든 에러를 처리할 수 있습니다.

 

https://docs.spring.io/spring-framework/docs/current/javadoc-api/org/springframework/web/socket/messaging/StompSubProtocolErrorHandler.html

StompSubProtocolErrorHandler (Spring Framework 6.1.9 API)

 

에러 핸들러 등록

@Configuration  
@EnableWebSocketMessageBroker  
@RequiredArgsConstructor  
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {  

    private final StompErrorHandler stompErrorHandler;  

    @Override  
    public void registerStompEndpoints(StompEndpointRegistry registry) {  
        registry.addEndpoint(ENDPOINT)  
                .setAllowedOriginPatterns(allowedOrigins);  
//                .withSockJS();  
        registry.setErrorHandler(stompErrorHandler); // 에러 핸들러 등록
    }  
}

설정에 구현한 에러 핸들러까지 등록하면 끝입니다.

 

예외 메세지 응답 화면

 

비즈니스 로직에 대한 검증 예외 처리

웹소켓 메시지 처리 과정에서 비즈니스 로직에 대한 검증 예외가 발생할 수 있습니다.

이를 처리하기 위해 @MessageExceptionHandler를 활용했습니다.

 

@Slf4j
@ControllerAdvice
@RequiredArgsConstructor
public class WebSocketExceptionHandler {

    private final CustomWebSocketHandlerDecorator decorator;

    @MessageExceptionHandler(SocketException.class)
    @SendToUser("/queue/errors")
    public ErrorResponse handleSocketException(Message<?> message) throws IOException {

        removeSession(message);

        return new ErrorResponse(3000, "SOCKET_ERROR");
    }

    @MessageExceptionHandler
    @SendToUser("/queue/errors")
    public ErrorResponse handleCustomException(CustomException exception) {

        return new ErrorResponse(exception.getErrorCode(), exception.getMessage());
    }

    @MessageExceptionHandler(RuntimeException.class)
    @SendToUser("/queue/errors")
    public ErrorResponse handleIllegalArgumentException() {

        return new ErrorResponse(2000, "Runtime Exception");
    }

    @MessageExceptionHandler
    @SendToUser("/queue/errors")
    public ErrorResponse handleValidationException(MethodArgumentNotValidException ex) {

        return new ErrorResponse(
                1001,
                ex.getBindingResult()
                        .getAllErrors()
                        .stream()
                        .map(DefaultMessageSourceResolvable::getDefaultMessage)
                        .collect(Collectors.joining(", "))
        );
    }

    private void removeSession(Message<?> message) throws IOException {
        StompHeaderAccessor stompHeaderAccessor = StompHeaderAccessor.wrap(message);
        String sessionId = stompHeaderAccessor.getSessionId();
        log.info("session = {}, connection remove", sessionId);
        decorator.closeSession(sessionId);
    }
}

 

@MessageExceptionHandler

@MessageExceptionHandler는 특정 예외가 발생했을 때 해당 예외를 처리할 수 있습니다.

CustomException이나 검증 예외 등 예외가 발생하였을 때 적절한 에러 메세지를 반환하도록 설정했습니다.

 

@SendToUser

예외를 처리한 이후, @SendToUser를 통해 예외 메세지를 해당 사용자에게 전송할 수 있습니다.

@SendToUser("/queue/errors")를 사용하면 /user/queue/errors를 구독하는 클라이언트에게 에러 메세지가 보내집니다.

 

https://docs.spring.io/spring-framework/reference/web/websocket/stomp/user-destination.html#page-title

User Destinations :: Spring Framework

 

removeSession

소켓이 끊어짐과 같은(SocketException) 치명적인 예외의 경우 세션을 삭제해 주었습니다.

 

 

예외 메세지 응답 화면

 

마무리하며..

예외 처리가 까다롭고 레퍼런스가 많이 없어서 블로그 글을 작성해보았습니다.

도움이 되셨으면 좋겠네요 !!

 

추가적인 질문이나 의견이 있다면 댓글로 남겨주세요.

읽어주셔서 감사합니다!