WebSocket에서 JWT 인증 구현

Chooooo·2024년 7월 30일
0

TIL

목록 보기
28/28

👍 WebSocket JWT 인증

이번 실시간 채팅 기능을 구현하면서 WebSocket 사용 시 JWT 인증을 처리하는 방법에 있어서 어떻게 적용했는지, 어떤 문제를 마주했는지 작성해보겠다.

WebSocket 연결 시 JWT 토큰을 사용한 인증이 제대로 이루어지지 않은 것이 문제여서 이를 해결하고자 이상한 짓을 너무 많이 했음

WebSocket 설정

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

    private final ChatPreHandler chatPreHandler;

    /**
     * TODO : 메시지 브로커 구성.
     * ! 메시지를 클라이언트에게 브로드캐스트할 때 사용할 주제 정의.
     * ! 클라이언트에서 서버로 메시지를 보낼 때 사용할 애플리케이션 목적지 접두사 설정
     * * -> 라우팅 규칙 정의하는 것
     */
    @Override
    public void configureMessageBroker(MessageBrokerRegistry config) {
        // ! sub으로 시작하는 목적지로 메시지 전달(해당 주소를 구독하는 클라이언트에게 메시지 보냄)
        config.enableSimpleBroker("/sub", "/user");
        // ! 클라이언트가 서버로 메시지 보낼 때 사용할 접두사 지정 -> 클라이언트가 "/pub"으로 시작하는 목적지로 메시지를 보내면 @MessageMapping이 달린 메서드로 라우팅
        config.setApplicationDestinationPrefixes("/pub");
        config.setUserDestinationPrefix("/user");
    }

    /**
     * TODO : STOMP 엔드포인트 등록
     * ! 클라이언트가 WebSocket 연결을 맺기 위해 사용할 엔드포인트 URL 지정.
     * ! /stomp/chat으로 클라이언트가 WebSocjet 연결을 맺기 위해 접속할 URL 경로
     * ? withSockJS()는 WebSocket을 지원하지 않는 브라우저에서도 동작하도록. 폴백 옵션 제공
     * * -> 이 메서드는 클라이언트가 서버와 WebSocket 연결을 시작하는 진입점 역할
     * @param registry
     */
    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/stomp/chat")
                .setAllowedOriginPatterns("*");
    }

    @Override
    public void configureWebSocketTransport(WebSocketTransportRegistration registration) {
        registration.setMessageSizeLimit(8192)
                    .setSendBufferSizeLimit(8192)
                    .setSendTimeLimit(10000);
    }

    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {
        registration.interceptors(chatPreHandler);
    }
}

WebSocket 메시지 브로커를 활성화하고 엔드 포인트, prefix, CORS 설정, 인증 인터셉터를 등록하는 설정 클래스

Spring Security

웹소켓 엔드포인트는 인증을 진행하지 않고 허용(permitAll())하도록 설정하자.
테스트 도구인 Apic을 사용할 경우는 CORS 설정을 전체 허용으로 해둬야 테스트가 가능하다.

CORS

특정한 origin만 허용하기 위해 사용하는 CORS(Cross Origin Resource Sharing)

클라이언트의 origin과 서버의 origin이 다르기 때문에 특정한 origin만 허용시켜 연결이 잘 동작하게 해야한다.
-> 그러기 위해서는 setAllowedOriginPatterns()를 사용해야 한다.

스프링 부트 2.x 버전에서는 setAllowedOrigins로도 가능했지만 부트 3 이상부터는 setAllowedOriginPatterns만 사용해야 CORS가 가능한 것 같았다.

👍 JWT 인증 기반의 프로젝트에서의 인증

JWT를 기반으로 하는 프로젝트의 인증은 클라이언트에서 보내는 모든 요청에 JWT를 넣어 전송하고, 서버에서 이를 검증해 통신을 주고 받으며 사용한다.

🧐 문제점 발생

-> 채팅 기능 개발하던 중에 소켓을 연결할 때 JWT 토큰이 인식되지 않고 null이 계속 들어오는 이슈 발생

특히 여러번 코드를 다른 방법으로 수정해도 토큰이 서버에 전달되지 않거나 null로 인식되는 문제가 있었음.
인증된 사용자 정보(memberId)가 WebSocket 세션에 저장되지 않아 후속 메시지 처리 시 인증 정보를 사용할 수 없었다.

왜 발생했는가 ?

팀원이 설정한 Spring Security의 권한을 모두 없애도 null이 들어옴
-> 찾아본 결과 웹소켓의 경우 헤더의 토큰을 검사하던 프로토콜이 HTTP와 달라 발생하는 것이었다.

Interceptor : HandshakeInterceptor vs ChannelInterceptor

둘다 모두 interceptor로 비슷하게 동작하지만, HandshakeInterceptor

public interface HandshakeInterceptor {
    boolean beforeHandshake(ServerHttpRequest var1, ServerHttpResponse var2, WebSocketHandler var3, Map<String, Object> var4) throws Exception;

    void afterHandshake(ServerHttpRequest var1, ServerHttpResponse var2, WebSocketHandler var3, @Nullable Exception var4);
}

ServerHttpRequest, ServerHttpResponse, WebSocketHandler를 사용함.

HandshakeInterceptor에서 Connect, DisConnect 프레임을 직접 가져와서 사용할 수는 없다고 알게 되었다.

그렇기에 웹 소켓을 연결하기 전에 JWT 인증을 마친 뒤 정상적으로 동작하는 것이 목표이기에 채널과 메시지 자체를 얻어와 처리할 수 있는 ChannelInterceptor를 사용

ChannelInterceptor 구현체

ChannelInterceptor를 구현하여 WebSocket 메시지를 직접 가로채고 처리할 수 있게 했다.

public interface ChannelInterceptor {
    @Nullable
    default Message<?> preSend(Message<?> message, MessageChannel channel) {
        return message;
    }

    default void postSend(Message<?> message, MessageChannel channel, boolean sent) {
    }

    default void afterSendCompletion(Message<?> message, MessageChannel channel, boolean sent, @Nullable Exception ex) {
    }

    default boolean preReceive(MessageChannel channel) {
        return true;
    }

    @Nullable
    default Message<?> postReceive(Message<?> message, MessageChannel channel) {
        return message;
    }

    default void afterReceiveCompletion(@Nullable Message<?> message, MessageChannel channel, @Nullable Exception ex) {
    }
}

Solution : ChannelInterceptor 구현

ChannelInterceptor를 구현하여 WebSocket 메시지를 직접 가로채고 처리할 수 있게 했다.

@Component
@RequiredArgsConstructor
public class StompHandler implements ChannelInterceptor {
    private final JwtTokenProvider jwtTokenProvider;

    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
        System.out.println("message:" + message);
        System.out.println("헤더 : " + message.getHeaders());
        System.out.println("토큰" + accessor.getNativeHeader("Authorization"));
        if (StompCommand.CONNECT.equals(accessor.getCommand())) {
            jwtTokenProvider.validateToken(Objects.requireNonNull(accessor.getFirstNativeHeader("Authorization")).substring(7));
        }
        return message;
    }
}

StompHeaderAccessor.wrap으로 message를 감싸면 STOMP의 헤더에 직접 접근할 수 있다. -> 즉, HTTP 프로토콜처럼 헤더에서 꺼내올 수 있게됨

위에서 작성한 클라이언트에서 보낸 JWT가 들어있는 헤더 Authorization

StompHeaderAccessor.getNativeHeader("Authorization") 메서드를 통해 받아올 수 있고

받아온 헤더의 값은 JWT가 된다. 받은 JWT를 검증해 정상적으로 소켓을 사용할 수 있도록 동작시키기 가능해짐

이제 이 ChannelInterceptor 구현한 후에 WebSocketConfig에서 사용할 수 있도록 코드를 추가해 인터셉터로 등록하면 JWT & WebSocket 인증이 마무리 된다.

우리 프로젝트의 ChannelInterceptor를 구현한 ChatPreHandler

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

    private final JWTUtil jwtUtil;

    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
        // ! 연결 요청시 JWT 검증
        if (StompCommand.CONNECT.equals(accessor.getCommand())) {
            // ! Authorization 헤더 추출
            String token = accessor.getFirstNativeHeader("Authorization");
            log.info("token = {}", token);
            if (token != null && token.startsWith("Bearer ")) {
                token = token.substring(7); // ! 추출
                try {
                    // ! JWT 토큰 검증
                    Map<String, Object> claims = jwtUtil.validateToken(token);
                    log.info("claims = {}", claims); // ! 전체 클레임 로깅
                    Long memberId = Long.parseLong(claims.get("memberId").toString());
                    log.info("memberId = {}", memberId);
                    accessor.setUser(new Principal() {
                        @Override
                        public String getName() {
                            return memberId.toString();
                        }
                    });
                } catch (Exception e) {
                    // ! 토큰 검증 실패
                    log.info("error = {}", e.getMessage());
                    return null;
                }
            }
        } else {
            return null; // ! 토큰 없거나 형식 잘못됨
        }

        return message;

    }
}

이제 해당 인터셉터를 WebSocketConfig에 추가하자.

	@Override
    public void configureClientInboundChannel(ChannelRegistration registration) {
        registration.interceptors(chatPreHandler);
    }

실제 사용

	/**
     * * 특정 채팅방 들어올 때
     * <p>
     * ! 읽지 않은 메시지 카운트 초기화 , 최근 메시지 로딩
     * <p>
     * ! 특정 채팅방 들어가는 것은 개인화된 반응 -> convertAndSendToUser로 변경.
     * TODO : 현재 memberId 하드코딩 -> JWT 토큰 이용한 인증 이후에 반영
     *
     * @return
     */
    @MessageMapping("/chat.enterRoom/{chatRoomId}")
    public void enterChatRoom(@DestinationVariable Long chatRoomId,
                              @Payload MessagePaginationRequest pageRequest,
                              SimpMessageHeaderAccessor headerAccessor) {
        Long memberId = Long.parseLong(headerAccessor.getUser().getName());
        log.info("memberId = {}", memberId);
        Pageable pageable = PageRequest.of(0, pageRequest.getPageSize(), Sort.by(Sort.Direction.ASC, "timestamp"));
        ChatRoomMessageResponse response = chatService.enterChatRoom(chatRoomId, memberId, pageable);
        messagingTemplate.convertAndSendToUser(memberId.toString(), "/sub/chatroom." + chatRoomId, response);
    }

위와 같이 SimpMessageHeaderAccessor 를 사용하면 검증이 가능하다.

새로운 Problem

기존의 WebSocket JWT 인증 방식에서는 인증 정보가 제대로 유지되지 않아, 후속 메시지 처리 시 사용자 식별에 실패했다...

사실 토큰까지는 받아와서 memberId까지 뽑아내는 것까지 성공했는데, 담아서 넘길 때 이때 뭐 때문에 안되는지는 모르겠지만.. Controller 단에서 사용할 때 null이 계속 체킹됐다.

원인 : Cause

  1. 인증 정보의 불완전한 저장 : 기존 방식에서는 인증 정보를 WebSocket 세션에 효과적으로 저장하지 못했음
  2. 후속 메시지에 대한 인증 처리 부재 : CONNECT 이외의 메시지에 대한 이증 확인 로직이 없었다.
  3. Spring Security와의 불완전한 통합 .. ? : UsernamePasswordAuthenticationToken을 사용하지 않아 Spring Security의 인증 메커니즘과 제대로 연동되지 않았는데, 사실 이것 때문인지는 모르겠음

해결 : Solution

  1. 세션 속성을 활용한 인증 정보 저장
accessor.getSessionAttributes().put("AUTHENTICATED_MEMBER_ID", memberId);
  • 이 코드로 인증된 사용자 ID를 WebSocket 세션 속성에 저장한다.
  1. Spring Security 인증 객체 생성 및 설정
UsernamePasswordAuthenticationToken auth =
    new UsernamePasswordAuthenticationToken(memberId, null, Collections.emptyList());
accessor.setUser(auth);
  • 이 코드로 Spring Security의 인증 객체를 생성하고 WebSocket 세션에 설정한다.
  1. 후속 메시지에 대한 인증 확인
else {
    Object memberId = accessor.getSessionAttributes().get("AUTHENTICATED_MEMBER_ID");
    if (memberId != null) {
        log.info("Authenticated user for non-CONNECT message: {}", memberId);
    } else {
        log.error("유저가 인증되지 못함");
        return null;
    }
}
  • 이 코드로 CONNECT 이외의 모든 메시지에 대해 인증 상태를 확인한다.

ChatPreHandler(ChannelInterceptor) 수정

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

    private final JWTUtil jwtUtil;

    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
        // ! 연결 요청시 JWT 검증
        if (StompCommand.CONNECT.equals(accessor.getCommand())) {
            // ! Authorization 헤더 추출
            String token = accessor.getFirstNativeHeader("Authorization");
            log.info("token = {}", token);
            if (token != null && token.startsWith("Bearer ")) {
                token = token.substring(7); // ! 추출
                try {
                    // ! JWT 토큰 검증
                    Map<String, Object> claims = jwtUtil.validateToken(token);
                    log.info("claims = {}", claims); // ! 전체 클레임 로깅
                    Long memberId = Long.parseLong(claims.get("id").toString());
                    log.info("memberId = {}", memberId);
                    // 인증 정보를 세션 속성에 저장
                    accessor.getSessionAttributes().put("AUTHENTICATED_MEMBER_ID", memberId);

                    // 인증 객체 생성 및 설정
                    UsernamePasswordAuthenticationToken auth =
                            new UsernamePasswordAuthenticationToken(memberId, null, Collections.emptyList());
                    accessor.setUser(auth);

                    log.info("Authentication set for user: {}", memberId);
                } catch (Exception e) {
                    // ! 토큰 검증 실패
                    log.error("토큰 유효하지 않음 or 실패", e);
                    return null;
                }
            } else {
                log.error("유효한 토큰을 찾지 못함");
                return null; // ! 토큰 없거나 형식 잘못됨
            }
        } else {
            // CONNECT 이외의 명령에 대해 인증 정보 확인
            Object memberId = accessor.getSessionAttributes().get("AUTHENTICATED_MEMBER_ID");
            if (memberId != null) {
                log.info("Authenticated user for non-CONNECT message: {}", memberId);
            } else {
                log.error("유저가 인증되지 못함");
                return null;
            }
        }

        // ! 모든 경우에 대해 인증 통과
        return message;

    }

}

WebSocketChatController에서의 변경

 `@MessageMapping("/chat.enterRoom/{chatRoomId}")
    public void enterChatRoom(@DestinationVariable Long chatRoomId,
                              @Payload MessagePaginationRequest pageRequest,
                              SimpMessageHeaderAccessor headerAccessor) {
        Long memberId = Long.parseLong(headerAccessor.getSessionAttributes().get("AUTHENTICATED_MEMBER_ID").toString());

// ... 나머지 로직
    }
  • SimpMessageHeaderAccessor 를 사용하여 세션 속성에서 직접 인증 정보를 가져온다.

최종 결과

  1. 인증 정보가 WebSocket 세션 전체에 걸쳐 유지된다.
  2. 모든 후속 메시지에 대해 일관된 인증 확인이 가능해졌다.
  3. Spring Security와의 통합이 개선되어 보안성 향상
  4. WebSocketChatController에서 이증된 사용자 정보를 쉽게 가져와 사용할 수 있음

Key Differences:

  1. 세션 속성 활용: 새로운 방식은 WebSocket 세션 속성을 활용하여 인증 정보를 저장하고 검색
  2. Spring Security 통합: UsernamePasswordAuthenticationToken을 사용하여 Spring Security와 더 잘 통합된다.
  3. 전체 메시지 인증: CONNECT 이외의 메시지에 대해서도 인증 상태를 확인한다.
  4. 컨트롤러에서의 접근 방식: SimpMessageHeaderAccessor를 통해 세션 속성에 직접 접근하여 인증 정보를 가져온다.

토큰이 오지 않을 경우 명시적으로 인터셉터에서 메시지를 보내기는 어렵기 때문에 null을 반환하자.

  • 역할(ROLE)에 대해서도 처리가 필요하다면 ref 블로그 두번째를 보면서 생각해보자.

안됐던 이유 분석

문제 상황:

웹소켓 연결에서 JWT를 이용한 사용자 인증이 지속되지 않았습니다.
초기 연결 이후의 메시지에서 사용자 식별이 불가능했습니다.

원인 분석:

웹소켓의 특성:

웹소켓은 한 번 연결된 후 지속적으로 통신합니다.
초기 연결 이후의 메시지들은 다른 스레드에서 비동기적으로 처리될 수 있습니다.

Spring Security의 인증 메커니즘:

Principal과 Authentication 객체는 주로 HTTP 요청-응답 모델에 최적화되어 있습니다.
SecurityContextHolder는 ThreadLocal을 사용하여 인증 정보를 저장합니다.
웹소켓 메시지가 다른 스레드에서 처리될 때 ThreadLocal에 저장된 인증 정보에 접근할 수 없습니다.

인증 정보의 불일치:

초기 연결 시 설정된 인증 정보가 후속 메시지 처리 시 유실되었습니다.

해결 방안:

세션 속성을 활용한 인증 정보 저장:

웹소켓 세션 속성에 인증 정보(memberId)를 저장합니다.
세션 속성은 연결이 유지되는 동안 계속 사용 가능합니다.

모든 메시지에 대한 인증 확인:

CONNECT 명령 외의 모든 메시지에 대해 세션 속성에서 인증 정보를 확인합니다.

Spring Security 통합:

UsernamePasswordAuthenticationToken을 생성하여 Spring Security의 인증 체계와 통합합니다.

구현 내용:

ChatPreHandler 수정:

CONNECT 명령 시 JWT 토큰을 검증하고 인증 정보를 세션 속성에 저장합니다.
다른 모든 메시지에 대해 세션 속성에서 인증 정보를 확인합니다.

WebSocketChatController 수정:

SimpMessageHeaderAccessor를 통해 세션 속성에서 인증 정보를 가져옵니다.

결과:

웹소켓 연결 전체에 걸쳐 일관된 사용자 인증이 가능해졌습니다.
다중 스레드 환경에서도 안정적으로 인증 정보를 유지하고 접근할 수 있게 되었습니다.
Spring Security와의 통합으로 보안성이 향상되었습니다.

이러한 해결 방안을 통해 웹소켓에서 JWT를 이용한 안정적인 사용자 인증을 구현할 수 있었습니다.

ref.블로그
1. Spring Boot - WebSocket & JWT & Spring Security 토큰 인증 설명
2. Security의 롤체크까지 하는 블로그
3. JWT + WebSocket 인증 관련

profile
back-end, 지속 성장 가능한 개발자를 향하여

0개의 댓글