[Chatty] Handshake Intercetpor 적용하기

문제점

채팅 메시지를 Stomp Handler를 사용해서 토큰 인증을 할 경우에, 토큰 인증이 실패되었을때, 에러메시지를 사용자에게 알려주려면 에러 메시지를 날려주는 채널을 따로 생성을해서 구독을 하는 방식으로 해야했다.

 

내가 생각한 해결 방법

사람이 많아지게되는 경우 각 채널마다 에러 메시지를 날려주는 채널을 추가적으로 만들어주게 되면 서버 자원 측면에서 비효율적이라고 생각했다. 그래서 이거에 대한 대안은 Stomp Handler에서 토큰 검증을 해주는 것이 아니라 http 통신을 할경우 3-way handshake를 하게 되는데, 이때 이 handshake를 가로채는 handshake Interceptor를 사용해서 토큰 검증을 해주도록 수정하려고 한다.

 

 

StompHandler에서 토큰 검증을 해주는 로직

import com.chatty.validator.TokenValidator;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.stereotype.Component;

@Slf4j
@Component
@RequiredArgsConstructor
public class StompHandler implements ChannelInterceptor {

    private final TokenValidator tokenValidator;

    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        log.info("Stomp Handler 실행");
        StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
        log.info("accessor : {}",accessor);
        log.info("stompCommand : {}",accessor.getCommand());
        if(StompCommand.CONNECT.equals(accessor.getCommand()) || StompCommand.SEND.equals(accessor.getCommand())){ // 연결할때
            String accessToken = accessor.getFirstNativeHeader("Authorization");
            tokenValidator.validateAccessToken(accessToken);
            log.info("토큰 검증 완료");
        }
        log.info("message : {}",message);
        return message;
    }
}

 

handshake Interceptor로 토큰을 검증해주는 로직

import com.chatty.handler.ChatHandler;
import com.chatty.jwt.JwtTokenProvider;
import com.chatty.validator.TokenValidator;
import lombok.RequiredArgsConstructor;
import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;

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

    private final TokenValidator tokenValidator;
    private final JwtTokenProvider jwtTokenProvider;
    private final ChatHandler chatHandler;

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {

        registry.addEndpoint("/ws")
                .addInterceptors(new WebSocketMatchInterceptor(tokenValidator, jwtTokenProvider))
                .setAllowedOriginPatterns("*");

        registry.addEndpoint("/signaling")
                .addInterceptors(new WebSocketMatchInterceptor(tokenValidator, jwtTokenProvider))
                .setAllowedOrigins("*");
    }

    @Override
    public void configureMessageBroker(MessageBrokerRegistry registry) {
        // 메시지를 구독하는 요청 url
        registry.enableSimpleBroker("/sub");

        // 메시지를 발행하는 요청 url
        registry.setApplicationDestinationPrefixes("/pub"); // Controller 객체의 MessageMapping 메서드 라우팅
    }
}

 

기존에 사용했던 Stomp Hanlder를 제거하고,

HandShakeInterceptor를 구현하는 WebSocketMatchInterceptor를 사용해준다.

import com.chatty.exception.CustomException;
import com.chatty.jwt.JwtTokenProvider;
import com.chatty.validator.TokenValidator;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

@Slf4j
@RequiredArgsConstructor
public class WebSocketMatchInterceptor implements HandshakeInterceptor {

    private final TokenValidator tokenValidator;
    private final JwtTokenProvider jwtTokenProvider;

    @Override
    public boolean beforeHandshake(final ServerHttpRequest request, final ServerHttpResponse response, final WebSocketHandler wsHandler, final Map<String, Object> attributes) throws Exception {
        String token = request.getHeaders().getFirst("Authorization");

        if (token != null && token.startsWith("Bearer ")) {

            try {
                tokenValidator.validateAccessToken(token);
            } catch (CustomException e){
                log.error("valid 에러 코드 발생");
                response.setStatusCode(HttpStatus.UNAUTHORIZED);
                return false;
            }
        }

        String mobileNumber = jwtTokenProvider.getMobileNumber(token.substring(7)); // key값으로 저장
        if (WebSocketConnectionManager.isConnected(mobileNumber)) {
            log.error("이미 연결이 존재합니다.");
            response.setStatusCode(HttpStatus.BAD_REQUEST);
            return false;
        }

        attributes.put("mobileNumber", mobileNumber);
        return true;
    }

    @Override
    public void afterHandshake(final ServerHttpRequest request, final ServerHttpResponse response, final WebSocketHandler wsHandler, final Exception exception) {
    }
}

 

 

해결 완료~!