웹소켓에서 유저 정보를 가져오려면?

❓ 문제상황

채팅 기능을 구현하는데, 어떤 사용자가 어떤 메세지를 입력했는지 보여주려고 한다. 메세지 내용 전달에는 성공했으나, 어떤 사용자가 그 내용을 입력했는지 즉 채팅방의 유저 정보를 전달하는데 어려움을 겪고 있었다.

💡 문제 해결을 위한 사고의 흐름

1. 웹소켓 또한 웹서버를 통해 웹브라우저와 통신한다. 하지만 서블릿 세션과 웹소켓 세션은 다르다.
2. 그렇다면 HttpSession에 담긴 정보들을 웹소켓 통신에서 가져올 수 있는 방법은 무엇일까 생각해 보았다.
   1. interceptor를 통해 handshake를 하기 전에 HttpSession 정보를 가져와서 WebSocketSession이나 StompHeaderAccessor 에 그 값을 넣어주기
   2. 기존에 사용하고 있던 ChannelInterceptor 의 핸들러를 사용해 presend나 postsend 단계에서 HttpSession정보를 StompHeaderAccessor 에 넣어주기

이런 식으로 생각을 해보다가 HttpSessionHandshakeInterceptor 의 존재를 알게 되었다.

🔍 HttpSessionHandshakeInterceptor.class

자바 공식 문서에 따르면, 이 클래스에 대한 설명은 다음과 같다.

💡 An interceptor to copy information from the HTTP session to the "handshake attributes" map to made available via `WebSocketSession.getAttributes()`. Copies a subset or all HTTP session attributes and/or the HTTP session id under the key `HTTP_SESSION_ID_ATTR_NAME`.

HttpSessionHandshakeInterceptor 의 기능은 HttpSession에 있던 정보들을 copy 해주는 것이다. 그래서 적합한 기능을 찾았다고 생각해서 해당 인터셉터를 추가하고 핸들러를 추가해주었다.

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(echoHandler, "/portfolio")
                .addInterceptors(new HttpSessionHandshakeInterceptor())
                .withSockJS();
    }

테스트로 WebSocketSession을 통해 SpringSecurityContext의 값을 가져올 수 있었고 유저명 정보도 가져올 수 있음을 확인했다.

@Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        SecurityContextImpl o = (SecurityContextImpl) session.getAttributes().get("SPRING_SECURITY_CONTEXT");
        UserDetailsImpl principal = (UserDetailsImpl) o.getAuthentication().getPrincipal();
        System.out.println("username = " + principal.getUsername());
    }

❓ 문제 발생

하지만 여기서 문제는 이렇게 되면 Front에서 Socket을 시작하는 버튼과 STOMP Message Broker를 시작하는 버튼이 따로 분리가 될 수 밖에 없는 문제가 발생했다.

💡 새로운 Solution

그래서 2-a에서 이야기한 것처럼 기존의 StompEndpoints에 MyHttpSessionHandshakeInterceptor 를 추가하고 HttpSessionHandshakeInterceptor 를 상속받은 클래스를 만들어서 handshake 전에 HttpSession 값을 가져오기로 했다.

@Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/portfolio")
                .addInterceptors(new MyHttpSessionHandShakeInterceptor())
                .withSockJS();
    }

public class MyHttpSessionHandShakeInterceptor extends HttpSessionHandshakeInterceptor {
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        HttpSession session = getSession(request);
        SecurityContextImpl o = (SecurityContextImpl) session.getAttribute("SPRING_SECURITY_CONTEXT");
        UserDetailsImpl principal = (UserDetailsImpl) o.getAuthentication().getPrincipal();
        System.out.println("principal.getUsername() = " + principal.getUsername());
        return true;
    }

    @Nullable
    private HttpSession getSession(ServerHttpRequest request) {
        if (request instanceof ServletServerHttpRequest) {
            ServletServerHttpRequest serverRequest = (ServletServerHttpRequest) request;
            return serverRequest.getServletRequest().getSession(isCreateSession());
        }
        return null;
    }
}

❓ 문제 발생

하지만 이렇게 구현을 했을 때, 2-b 단계에서처럼 header 값에 정보를 넣어주기 위해서는 그 값들을 그 다음 인터셉터로 전달하는 과정이 필요한데 그 부분에서 해답이 떠오르지 않았다.

아니면 프론트 단에 전달하는 message의 payload 안에 그 값을 넣으려고 @MessageMapping 된 Controller 단에서 방법을 찾아보려고 했지만 그 또한 방법을 찾지 못했다.

그러던 중 공식문서에서 MessageMapping 이 가질 수 있는 여러 arguments 들을 재확인하게되었다.

💡 최종 Solution 발견

이미 @Payload, @DestinationVariable 은 사용하고 있었고, java.security.Principal 라는 것이 눈에 들어왔다. 이에 대한 설명은 Reflects the user logged in at the time of the WebSocket HTTP handshake. 라고 되어 있었다.

Spring Security 프레임워크를 쓰다보니 SecurityContextHolder 안에 SecurityContext가, 그리고 그 안에 Authentication이 들어가고 그를 통해 Principal 을 얻어 현재 로그인된 사용자의 정보를 추출할 수 있다. 나는 이걸 Session 값을 통해서 얻으려고 했는데, STOMP의 Message Mapping 기능에서 파라미터로 Principal 을 바로 쓸 수 있도록 제공을 해주니… 참 최종적으로 여러 삽질을 하다가 굉장히 쉽게 문제가 해결된 셈이다. (공식문서를 처음부터 꼼꼼하게 읽었으면 더 좋았을 걸…)

@MessageMapping("/chat/{chatroom}")
@SendTo("/topic/{chatroom}")
public ChatMessage handle(@Payload ChatMessage chatMessage, 
													@DestinationVariable String chatroom, 
													java.security.Principal principal) {
    chatMessage.addWriter(principal.getName());
    return chatMessage;
}

하지만 이번 계기를 통해 웹소켓의 동작원리, 일반 WebSocket과 SockJs의 차이, STOMP의 역할 뿐 아니라 웹소켓 VS 서블릿 세션의 통신 방법 등 HTTP 프로토콜과 ws 프로토콜에 대해 더 명확하게 알 수 있는 시간이 되었다.