Spring WebSocket STOMP 채팅 서버 구현하기 (with. JWT, Exception Handling)

===== 채팅서버 구현하기 시리즈  =====

• ✔️ [Web-Network] - 채팅 서버 설계를 위한 배경지식 정리 (HTTP, WebSocket, WebRTC)
• ✔️ [Spring/Spring Boot] - Spring WebSocket 공식문서 가이드 살펴보기
• 👉🏻 [Spring/Spring Boot] - Spring WebSocket STOMP 채팅 서버 구현하기 (with. JWT, Exception Handling)
• ✔️ [Spring/Spring Boot] - Spring Data mongoDB + mysql 사용하기 (with. queries)

---

 

spring 으로 간단하게 채팅기능을 구현해 기록해보고자 합니다. 구현을 중점으로해서 가장 좋은 방법은 아닐겁니다(?)

만들고자 했던 서비스는, 카카오톡과 같은 다중 채팅방이 가능한 1:1 채팅 서비스 입니다

 

 

이야기 해볼건 크게 3가지가 될 것 같습니다.

1. spring websocket stomp 로 채팅 구현하기 
2. websocket 연결전 tcp handshake 과정에서 JWT 인증하기
3. spring websocket Exception 에러 핸들링

( + 로컬에서 websocket 테스트 방법)

 

---

1. 📡 웹소켓 STOMP 채팅 구현

사실 이전 게시글에서 공식문서를 봤기 때문에, 빠르고 간결하게 코드만 슥슥 기록하고 넘어가겠습니다,,,ㅎㅎ

 

✔️ build.gradle

// spring websocket
implementation 'org.springframework.boot:spring-boot-starter-websocket'
implementation 'org.webjars:webjars-locator-core'
implementation 'org.webjars:sockjs-client:1.5.1'	//앱은 없어도 됨
implementation 'org.webjars:stomp-websocket:2.3.4'

 

✔️  WebSocketConfig

• StompHandler 는 websocket stomp 로 연결하는 흐름에 대한 제어를 위한 interceptor 입니다. JWT 인증을 위해 사용했습니다.
• StompExceptionHandler 는 websocket 연결 시 터지는 exception 을 핸들링하기 위한 클래스 입니다.

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

	//둘 다 아래에서 exception handler 에 필요함
    private final StompHandler stompHandler;
    private final StompExceptionHandler stompExceptionHandler;

    @Override
    public void configureMessageBroker(MessageBrokerRegistry registry) {
        //해당 파라미터의 접두사가 붙은 목적지(구독자)에 메시지를 보낼
        registry.enableSimpleBroker("/sub");
        registry.setApplicationDestinationPrefixes("/pub");
    }

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        //엔드포인트 추가 등록
        registry
            .setErrorHandler(stompExceptionHandler)		//exception handler를 위한 것
            .addEndpoint("/websocket-endpoint")
            .addInterceptors()
            .setAllowedOriginPatterns("*");
    }

	//tcp handshake시 jwt 인증용
    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {

        registration.interceptors(stompHandler);
    }
}

 

✔️  Controller

• 다중 채팅방을 구현하는 방법은 생각보다 간단했습니다. websocket 으로 연결 후 puslisher(보내는 사람) 이 어디다 보낼지 채팅방 번호를 명시하고, subscriber(받는 사람)도 채팅방 번호를 구독하고 → 들어온 그대로 발송해주면 됩니다.
• websocket send 의 목적지로 들어오는 pathvariable의 변수값이 필요하다면 @DestinationVariable 어노테이션을 사용해 아래처럼 받을 수 있습니다.

@Slf4j
@Controller
@RequiredArgsConstructor
public class ChatController {

    private final ChatService chatService;

    @MessageMapping("/chat/{roomNo}")
    @SendTo("/sub/chat/{roomNo}")
    public ChatResponse broadcasting(final ChatRequest request,
        @DestinationVariable(value = "roomNo") final String chatRoomNo) {

        log.info("{roomNo : {}, request : {}}", chatRoomNo, request);

        return chatService.recordHistory(chatRoomNo, request);
    }

}

👏🏻 이렇게 하면 다중 spring websocket 을 이용한 다중 채팅방 구현이 끝났습니다!

chatservice 는 mongodb를 이용해 채팅내용을 저장하는 기능만 하기에 여기서는 넘기고 다음 게시물에서 작성하겠습니다

 

---

2. WebSocket + JWT 사용하기

그래도 별도의 api 서버에서 jwt 를 유저에게 발급해 주고있고, 아무나 아무방에 막 연결하는건 아닌거같아 채팅메세지를 보내는데 최소한의 인증절차를 걸쳤으면 했고,

websocket 도 처음에는 tcp 를 이용해 연결하기 때문에 딱 처음 연결될 때 JWT 를 이용해  딱 1번! 유효한 유저인가!를 판단하고자 했습니다.

 

✔️ WebSocketConfig

• 위에 Websocket Config 에서 configureClientInboundChannel(ChannelRegistration ~) 메소드를 재정의 해주었습니다.
• configureClientInboundChannel 메소드는 Client 로 부터 들어오는 메시지를 처리하는 MessageChannel 을 구성하는 역할을 하는 메소드입니다.
• registration.intreceptor() 메소드를 사용해 STOMP 메세지 처리를 구성하는 메세지 채널에 custom 한 인터셉터를 추가 구성하여 채널의 현재 인터셉터 목록에 추가하는 단계를 거칩니다.

@Override
public void configureClientInboundChannel(ChannelRegistration registration) {

    registration.interceptors(stompHandler);
}

 

간단하게 그냥 웹소켓 연결전, 연결후 추가적인 작업을위한 interceptor를 추가 구성할 수 있게 해줍니다~

 

✔️ StompHandler

• @Order 어노테이션은 구성된 인터셉터들 간의 작업 우선순위를 지정할 수 있게 해줍니다. HIGHEST_PRECEDENCE 는 가장 높은 우선순위이며, 거기에 +99를 해줘서 더 높은 우선순위를 주었습니다.
• presend( ) 는 메시지가 실제로 채널로 전송되기 전에 호출된다고 합니다. 즉, publisher 가 send 하기 전에 일어난다고 이해가 됩니다.
• StompHeaderAccessor.wrap으로 message를 감싸면 STOMP의 헤더에 직접 접근할 수 있습니다. 클라이언트에서 첫 연결시 헤더에 token 을 담아주면 인증 절차가 진행됩니다.

@Component
@RequiredArgsConstructor
@Order(Ordered.HIGHEST_PRECEDENCE + 99)
public class StompHandler implements ChannelInterceptor {

    private final JwtUtils jwtUtils;

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

        final StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);

        // websocket 연결시 헤더의 jwt token 유효성 검증
        if (StompCommand.CONNECT == accessor.getCommand()) {
            final String authorization = jwtUtils.extractJwt(accessor);

            jwtUtils.validateToken(authorization);
        }
        return message;
    }
}

 

✔️ JwtUtils

• presend( ) 메소드로 메세지가 채널에 전송되기 전(웹소켓 연결 전) jwt 인증 절차를 걸치고 유요하지 않은 토큰이라면 exception을 터트립니다.

@Slf4j
@Component
public class JwtUtils {

    private final Key jwtSecretKey;

    public JwtUtils(@Value("${jwt.secret-key}") String secretKey) {
        byte[] keyBytes = Decoders.BASE64.decode(secretKey);
        jwtSecretKey = Keys.hmacShaKeyFor(keyBytes);
    }

    public String extractJwt(final StompHeaderAccessor accessor) {
        return accessor.getFirstNativeHeader("Authorization");
    }

    // jwt 인증
    public void validateToken(final String token) {
        try {
            Jwts.parserBuilder().setSigningKey(jwtSecretKey).build().parseClaimsJws(token);
        } catch (io.jsonwebtoken.security.SecurityException | MalformedJwtException e) {
            throw UnauthorizedException.of(e.getClass().getName(),"잘못된 JWT 서명입니다.");
        } catch (ExpiredJwtException e) {
            throw UnauthorizedException.of(e.getClass().getName(),"만료된 JWT 토큰입니다.");
        } catch (UnsupportedJwtException e) {
            throw UnauthorizedException.of(e.getClass().getName(),"지원되지 않는 JWT 토큰입니다.");
        } catch (IllegalArgumentException e) {
            throw UnauthorizedException.of(e.getClass().getName(),"JWT 토큰이 잘못되었습니다.");
        }
    }
}

 

---

3. ⛔️ WebSocket STOMP 에러 Exception Handling

이제 터진 exception을 핸들링 해야겠죠

 

✔️ WebSocketConfig

• 다시 config 로 돌아와 endpoint를 등록할 때, setErrorHandler()를 이용해 해당 endpoint에 에러 핸들러를 등록해줍니다.

@Override
public void registerStompEndpoints(StompEndpointRegistry registry) {
    //엔드포인트 추가 등록
    registry
        .setErrorHandler(stompExceptionHandler)
        .addEndpoint("/websocket-endpoint")
        .addInterceptors()
        .setAllowedOriginPatterns("*");
}

 

✔️ StompExceptionHandler

💡 stomp의 exception handler는 사실 구성해주기 나름이라고 생각합니다.
저는 가장 간단하게만 사용해주었고 중요한 부분은 handlerClientMessageProcessingError( ) 메소드와  HandlerInternal( ) 메소드입니다.

1. websocket endpoint 에서 에러가 터지면 원래는 → 아래 코드에서 상속받은 StompSubProtocolErrorHandler 의 handlerClientMessageProcessingError( ) 를 호출하고, 처리가 끝나면 HandlerInternal( )를 마지막으로 호출해 메세지를 전송하는 흐름을 가집니다.
2. 그렇기 때문에 handlerClientMessageProcessingError( ) 를 오버라이딩해 터진 exception이 내가 핸들링하기 원했던 exception 이다! 를 체크해 별도의 핸들러 로직을 호출하고 원하는 메세지를 client에게 반환하도록 구성해 주었습니다.

 

 

@Component
public class StompExceptionHandler extends StompSubProtocolErrorHandler {

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

    public StompExceptionHandler() {
        super();
    }

    @Override
    public Message<byte[]> handleClientMessageProcessingError(Message<byte[]> clientMessage,
        Throwable ex) {

        final Throwable exception = converterTrowException(ex);

        if (exception instanceof UnauthorizedException) {
            return handleUnauthorizedException(clientMessage, exception);
        }

        return super.handleClientMessageProcessingError(clientMessage, ex);

    }

    private Throwable converterTrowException(final Throwable exception) {
        if (exception instanceof MessageDeliveryException) {
            return exception.getCause();
        }
        return exception;
    }

    private Message<byte[]> handleUnauthorizedException(Message<byte[]> clientMessage,
        Throwable ex) {

        return prepareErrorMessage(clientMessage, ex.getMessage(), HttpStatus.UNAUTHORIZED.name());

    }

    private Message<byte[]> prepareErrorMessage(final Message<byte[]> clientMessage,
        final String message, final String errorCode) {

        final StompHeaderAccessor accessor = StompHeaderAccessor.create(StompCommand.ERROR);
        accessor.setMessage(errorCode);
        accessor.setLeaveMutable(true);

        setReceiptIdForClient(clientMessage, accessor);

        return MessageBuilder.createMessage(
            message != null ? message.getBytes(StandardCharsets.UTF_8) : EMPTY_PAYLOAD,
            accessor.getMessageHeaders()
        );
    }

    private void setReceiptIdForClient(final Message<byte[]> clientMessage,
        final StompHeaderAccessor accessor) {

        if (Objects.isNull(clientMessage)) {
            return;
        }

        final StompHeaderAccessor clientHeaderAccessor = MessageHeaderAccessor.getAccessor(
            clientMessage, StompHeaderAccessor.class);

        final String receiptId =
            Objects.isNull(clientHeaderAccessor) ? null : clientHeaderAccessor.getReceipt();

        if (receiptId != null) {
            accessor.setReceiptId(receiptId);
        }
    }

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

        return MessageBuilder.createMessage(errorPayload, errorHeaderAccessor.getMessageHeaders());

//        return super.handleInternal(errorHeaderAccessor, errorPayload, cause, clientHeaderAccessor);
    }
}

 

+

websocket stomp 는 postman으로 테스트해볼수가 없습니다.

저는 apic 이라는 별도의 프로그램을 사용하여 로컬 환경에서 웹소켓 테스트를 할 수 있었습니다.

websocket 을 사용할때 endpoing url 은 http 대신 ws 를, https 대신 wss 를 사용하시면 됩니다 :)

 

 

 

끝!

 

---

참고

• stomphandler : https://docs.spring.io/spring-integration/archive/1.0.0.M6/reference/html/ch02s05.html
• spring docs - websocket support : https://docs.spring.io/spring-framework/docs/4.3.x/spring-framework-reference/html/websocket.html