[WebSocket] 웹 소켓 연결 시 검증은 어디서 해야 할까?

 

개요

프로젝트에서 WebSocket을 통해 채팅방을 성공적으로 구현하였다. 하지만 요구사항은 모임에 참여한 사람만 채팅방의 참여가 가능해야하지만, 현재는 채팅방은 누구나 WebSocket을 연결하여 채팅을 송신하거나 수신할 수 있는 상태이다. 

WebSocket 연결 시, 연결에 대한 인증/인가 로직은 어디다 두어야할지 고민한 흔적을 적어보려 한다.

 

 

 

WebSocket 연결 흐름

1. WebSocket 연결

출처 : https://velog.io/@kmh916/Spring-Boot-%EC%B1%84%ED%8C%85-%EA%B5%AC%ED%98%84SockJSStomp-0.-WebSocketSockJSStomp-%EB%9E%80

WebSocket 연결 시, HTTP 연결이 한 번 일어나며 HandShake 과정도 일어난다. 이 과정이 잘 통과되면, 다른 한 쪽의 연결이 끊기기 전까지 양방향 통신이 가능하다

 

2. Stomp 연결

해당 프로젝트에서 Stomp를 이용하여 채팅방을 고도화 하였다. Stomp는 WebSocket 위에서 동작하는 프로토콜이며, 다음과 같은 요청 상태들이 있다

• Connect - WebSocket HandShake 이후 Stomp 연결
• Subscribe - sender를 구독하여 수신할 수 있다.
• Send - 메시지를 송신
• Disconnect - WebSocket에 연결이 끊김

 

 

연결 검증은 어디서?

정확한 요구사항은 다음과 같다.

채팅방 입장(WebSocket 연결) 시, 검증된 사용자가 접근할 수 있다.
추가적으로 사용자는 해당 모임의 참여중이어야 한다.

WebSocket 연결 시, 모임의 정보와 사용자의 정보를 갖고와 가입 여부를 검증해야 한다. 해당 로직을 해결할 수 있는 곳이 군데 있다

 

 

 

1. Security FilterChain

 

WebSocket 연결 시도 시, HTTP 요청이 한 번 전송되어 HandShake가 일어나며 SpringSecurity Filter도 거친다. 하지만 해당 HTTP 요청은 Custom Header 설정이 불가능하기에 사용자 정보를 헤더에 넣어줄 수 없다..!!

Spring Security를 이용한 인증은 패스!!

 

 

2. Stomp Connect 연결 시점

Spring은 Stomp에 대해 다양한 기능을 제공하고 있다. 그 중 Client로부터 오는 요청을 catch하는 Interceptor를 통해 요청을 처리하기 전 작업을 할 수 있다

 

필자는 WebSocket 연결 후 Stomp Connect 요청을 catch하여 참여 검증과 예외 처리를 하려 하였다.

코드는 다음과 같다

 

 

코드

StompAuthenticationHandler - 인증 핸들러

@RequiredArgsConstructor
@Component
public class StompAuthenticationHandler implements ChannelInterceptor {

    private final JwtTokenProvider jwtTokenProvider;
    private final ParticipantRepository participantRepository;

    private Long getUserInfo(String token) {
        if (StringUtils.hasText(token)) {
            String accessToken = token.split(" ")[1];
            TokenPayload payLoad = jwtTokenProvider.getPayLoad(accessToken);
            return payLoad.userId();
        } else {
            throw new ValidationException(AuthErrorCode.AUTH_REQUIRED);
        }
    }

    //클라이언트에서 요청을 보낼 때 실행
    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {

        //Stomp 헤더 접근
        StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
        StompCommand command = accessor.getCommand();

        //첫 연결이거나 메시지를 보낼 때 검증
        if (command == StompCommand.CONNECT) {
            //AccessToken 검증
            String token = accessor.getFirstNativeHeader(AUTHORIZATION);
            Long userInfo = getUserInfo(token);

            //참여 여부 확인
            String roomId = accessor.getFirstNativeHeader("roomId");
            boolean isParticipant = participantRepository.existsByUserIdAndRoomId(userInfo,
                Long.parseLong(roomId));

            if (!isParticipant) {
                throw new ValidationException(RoomErrorCode.INVALID_PARTICIPANT);
            }
        }

        return message;
    }
}

 

 

StompExceptionHandler - 예외 핸들러

@Component
public class StompExceptionHandler extends StompSubProtocolErrorHandler {

    //예외 메시지 생성
    private Message<byte[]> prepareErrorMessage(String errorMessage, String code) {

        //헤더 유형 -> 에러
        StompHeaderAccessor accessor = StompHeaderAccessor.create(StompCommand.ERROR);
        accessor.setMessage(code);

        //헤더 내용 변경 불가 설정
        accessor.setLeaveMutable(true);

        //에러 메시지 생성 후 전달
        return MessageBuilder.createMessage(errorMessage.getBytes(StandardCharsets.UTF_8),
            accessor.getMessageHeaders());
    }

    //웹 소켓 통신 중 발생한 예외 -> MessageDeliveryException으로 배달
    private Throwable convertThrowException(Throwable throwable) {
        if (throwable instanceof MessageDeliveryException) {
            return throwable.getCause();
        }
        return throwable;
    }

    //예외 발생 시 실행되는 메서드
    @Override
    public Message<byte[]> handleClientMessageProcessingError(Message<byte[]> clientMessage,
        Throwable ex) {

        //예외 변환
        Throwable exception = convertThrowException(ex);

        if (exception instanceof ValidationException ve) {
            //예외 메시지 생성
            return prepareErrorMessage(ve.getMessage(), ve.getCode());
        }

        return super.handleClientMessageProcessingError(clientMessage, ex);
    }
}

 

Stomp는 요청 시, Custom Header가 가능하다.

헤더에 사용자 정보와 모임 정보를 넣어주고, 요청 처리 전 인증 핸들러에서 검증하는 로직을 구상했다.

 

토큰이 올바르지 않거나, 사용자가 모임의 참가하지 않았다면 예외를 던지고

예외를 처리하는 핸들러를 통해 응답을 하도록 하였다

 

 

결과

결론적으로 해당 방법은 내가 생각한 결과가 아니였다.

 

당연하지만, 일단 예외 메시지가 HTTP 응답이 아닌 Stomp 메시지로 이루어진다.

 

 

또한 실제로 Connect가 실패되었지만, HandShake 자체는 성공했기에 200 응답 코드를 뱉는다.

필자는 클라이언트에게 확실하게 응답해주고 싶었다.

 

 

3. HTTP HandShake 시점

Spring은 놀랍게도 WebSocket 연결 HandShake 시 요청을 Interceptor해주는 클래스를 지원한다..!

WebSocket 연결 시, 검증이 성공하면 HandShake가 일어나고,

검증 예외가 발생하면 HTTP HandShake가 일어나지 않고, 4XX 응답을 내려줄 수 있어서 아주 이상적인 방법이다!!

 

단점은 앞에서도 말했듯, Custom Header 세팅이 불가능하다는 점이다..!! 필자는 사용자와 모임 정보들을 QueryParameter로 넘겨주기로 하였다!

 

 

WebSocketHandShakeHandler

@Component
@RequiredArgsConstructor
public class WebSocketHandShakeHandler implements HandshakeInterceptor {

    private final JwtTokenProvider jwtTokenProvider;
    private final ParticipantRepository participantRepository;
    private final ObjectMapper mapper = new ObjectMapper();

    private String toJson(ErrorResponse response) throws JsonProcessingException {
        return mapper.writeValueAsString(response);
    }

    private void sendJson(HttpServletResponse response, String resultJson) throws IOException {
        response.setStatus(HttpStatus.UNAUTHORIZED.value());
        response.setContentType("application/json;charset=UTF-8");
        response.setContentLength(resultJson.getBytes().length);
        response.getWriter().write(resultJson);
    }

    //HandShake 일어나기 전 실행
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
        WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {

        if (request instanceof ServletServerHttpRequest httpRequest &&
            response instanceof ServletServerHttpResponse httpResponse) {
            HttpServletRequest httpServletRequest = httpRequest.getServletRequest();
            HttpServletResponse httpServletResponse = httpResponse.getServletResponse();

            //정보 추출
            String token = httpServletRequest.getParameter("user");
            String roomId = httpServletRequest.getParameter("roomId");

            //AccessToken 검증
            Long userInfo;
            try {
                TokenPayload payLoad = jwtTokenProvider.getPayLoad(token);
                userInfo = payLoad.userId();
            } catch (ValidationException ve) {
                sendJson(httpServletResponse,
                    toJson(new ErrorResponse(ve.getMessage(), ve.getCode())));

                return false;
            }

            //참여 여부 확인
            boolean isParticipant = participantRepository.existsByUserIdAndRoomId(userInfo,
                Long.parseLong(roomId));

            if (!isParticipant) {
                RoomErrorCode invalidParticipant = RoomErrorCode.INVALID_PARTICIPANT;
                sendJson(httpServletResponse,
                    toJson(new ErrorResponse(invalidParticipant.getMessage(),
                        invalidParticipant.getCode())));

                return false;
            }
        }

        return true;
    }

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

    }
}

 

QueryString을 통해 정보를 받고, HandShake 전 검증 로직을 수행한다. 

인자로 ServerHttpResponse를 제공하니, 올바르지 않은 사용자가 접근 시 직접 response를 세팅하여 응답한다. 

 

 

HTTP 응답으로 에러 코드와 메시지가 잘 내려오는 것을 확인할 수 있다!

 

 

 

정리

사실 이 방법이 아직 올바른지는 확실하지 않다.

AccessToken을 URL에 담아 옮겨주는 것이 인증에 대 취약점이 될 수도 있으며, 프론트엔드에서 해당 에러 응답을 처리하지 못 할 수도 있을 것 같아 좀 더 나은 방법을 고려해야 할 것 같다..

 

그렇더라도 WebSocket 연결 과정과 Spring에서 제공해주는 기능들을 사용해볼 수 있어 좋았으며, WebSocket을 사용하며 추후 타 기능 구현 시 써먹을 수 있을 것 같다!

---

참고

https://velog.io/@tlatldms/Socket-%EC%9D%B8%EC%A6%9D-with-API-Gateway-Refresh-JWT