DefaultHandshakeHandler를 이용하여 Websocket session에 유저 정보 담기

사건 발달

현재 어플리케이션은 사용자가 웹소켓 서버를 통해 전송하는 모든 메시지와 관련 정보를 데이터베이스(DB)에 저장하는 방식으로 동작하고 있었다.

이로 인해 사용자가 메시지를 전송할 때마다 매번 사용자의 정보를 DB에서 조회하거나 해당 사용자의 정보를 메시지에 포함하여 전송해야 하는 번거로움이 있었다. 따라서 서버의 최적화와 보안상의 이유로 더 효율적인 방법을 고려해야만 했다.

이에 따라, 사용자의 정보를 SecurityContextHolder에 저장된 Principal 객체에 담아 Websocket의 세션(session)에 저장하기로 결정하였다. 그러나 Websocket의 세션 객체에는 setPrincipal과 같은 메서드가 없어 Principal 객체를 세션에 저장하기 위해서는 TCP 3-way handshake 단계에서 Principal 객체를 전달하는 방법을 사용해야만 했다.

 

 

 

WS 프로토콜의 기본 동작 방식

웹소켓은 하나의 TCP 접속에 전이중 통신 채널을 제공하는 컴퓨터 통신 프로토콜이다. 즉 TCP의 3- way handshake 과정을 거처 사용자를 특정하고 Spring Websocket에서는 사용자들에게 고유의 세션을 발급하여 해당 세션에 직접적으로 데이터를 전송하게된다. 

 

STOMP 프로토콜을 사용하면 메세지를 헤더와 바디로 분리하여 작성이 가능한 것 같다. 하지만 프로젝트 진행 기간이 짧아 STOMP까지 공부할 시간이 없고 해당 어플리케이션의 특성상높은 수준의 보안이 필요하지 않아 클라이언트가 Handshake 과정을 거칠때 토큰을 전달받아 인증을 마치고 해당 토큰에 담긴 사용자의 정보를 고유의 Websocket session에 담아 사용하기로 했다.

 

 

DefaultHandshakeHandler

Handshake 과정을 개발자가 정의하기 위해서는 DefaultHandshakeHandler클래스를 상속받아 사용해준다.

 

@RequiredArgsConstructor
@Configuration
@EnableWebSocket
public class WebSockConfig implements WebSocketConfigurer {
    private final WebSocketHandler webSocketHandler;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(webSocketHandler, "/ws/chat")
                .setHandshakeHandler(new CustomHandshakeHandler())
                .setAllowedOrigins("*");
    }
}

먼저 웹소켓 설정을 위한 WebSocketConfigurer의 WebSocketHandlerRegistry에 setHandshakeHandler 메서드를 이용해 사용자 정의 핸들러를 등록해준다.

 

 

 

이후 해당 클래스에 DefaultHandshakeHandler를 상속받아 determinUser 메서드를 정의해주면 된다.

public class CustomHandshakeHandler extends DefaultHandshakeHandler {

    @Override
    protected Principal determineUser(ServerHttpRequest request, 
                                      WebSocketHandler wsHandler, 
                                      Map<String, Object> attributes) {
        Principal member = (Principal) SecurityContextHolder.getContext().getAuthentication();
        return member;
    }
}

determineUser 메서드는 Principal 객체를 반환해 주어야 한다.

 

request의 헤더를 통해 토큰을 가저와 복호화를 하여 사용자의 정보에 접근해도 되지만 이 경우에는 이미 스프링 시큐리티를 통해 토큰인증이 완료가 된 사용자만을 접근을 허용하게 되있어 따로 예외처리를 하지 않고 바로 시큐리티 컨텍스트 홀더에 접근해 Principal로 캐스팅해준 뒤 반환해 주었다. 

 

 

이후 WebSocketHandler에는 사용자 정보가 담겨있는 session을 전달 받을 수 있다.

@Slf4j
@RequiredArgsConstructor
@Component
@Transactional
public class WebSockChatHandler extends TextWebSocketHandler {
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception, IOException{
        Authentication authentication = (Authentication) session.getPrincipal();
        MemberContextInform member =  (MemberContextInform) authentication.getPrincipal();
    }
}

getPrincipal() 메서드를 사용하면 Principal 객체로 사용자 정보를 받을 수있다. 해당 어플리케이션은 컨텍스트 홀더에서 사용자 정보를 가저왔으므로 Authentication으로 캐스팅 할 수 있었다. 만약 위와같이 MemberContextInform과 같은 DTO를 사용한다면 다시한번 캐스팅하여 사용한다.

 

 

 

마무리

처음에는 WebSocketSession에 setPrincipal 같은 메서드가 존재 할 줄 알았다. 하지만 웹소켓 세션에 사용자를 정의해 주기 위해서는 DefaultHandshakeHandler의 determinUser 메서드를 사용해야만 했다.

 

사실 위의 방식으로 WebSocketSession에 사용자의 정보를 담는 것은 너무 쉬웠다.

 

하지만 클라이언트에서 아래와 같은 방법으로 토큰을 넘겨주게 되면 request header의 sec-websocket-protocol로 토큰 값이 넘어오는데 띄어쓰기가 불가능하고 만약 Refresh 토큰도 같이 넘겨주게되면 배열도 아닌 그냥 콤마(,)로만 분리가 되어 개발자가 해당값을 직접적으로 분리하여 Access 토큰과 Refresh 토큰을 분류해 주어야 했다. 

const client = new WebSocket(
      "ws://localhost:8080/ws/chat/roomId="+ room,
      token
    );

 

 

뭔가 더 세련된 방법이 있을거라 생각하고 더 찾아보았지만 STOMP 프로토콜을 사용하지 않는 이상 Spring Websocket으로는 한계가 명확했다... 

 

 

더군다나 구글링을 할 수록 웹소켓 서버는 노드JS를 사용하는 경우가 더 많은 듯 했다. 아마도 실시간 채팅 특성상 자바의 안정성보다는 JS와 노드의 높은 호환성과 빠른 속도가 더 중요하기 때문이 아닐까라 생각이 든다. 

 

어쨋든 다음에 웹소켓을 개발하여야 한다면 그냥 노드JS를 사용하는 걸로....

 

수정사항

스택오버플로우의 추천으로 protocol에 토큰을 넣어 보냈지만 아무래도 프로토콜 섹션에 토큰을 보내는게 맞지않는다 생각하여 토큰을 쿠키에 저장하여 보내도록 바꿧다. 

 

클라이언트에서 토큰을 쿠키에 저장한 뒤 웹소켓에 연결을 시도하면 자동적으로 쿠키 정보가 웹소켓 서버로 전달된다.

해당 쿠키를 JWT를 이용하여 복호화 작업을 한뒤 사용자 인증을 마첬다.

 

EC2에 배포이후 보안성을 이유로 쿠키를 통해 특정 정보를 보내지 못하게 되었다. 결국 쿼리파라미터로 토큰을 전달방식을 사용하게 되었다.