WebSocket 과 @AuthenticationPrincipal

MoonJaeGyeong·2024년 12월 6일
Springwebsocket

WebSocket

목록 보기
1/1

개요

보통 나는 API를 개발해서 로그인한 유저, 인증 객체를 가져오기 위해 @AuthenticationPrincipal 이라는 어노테이션을 통해 가져오는 편이다. 그래서 웹소켓에서도 동일하게 해당 어노테이션을 통해 인증 객체의 이메일을 가져오고자 하였는데

해당 코드를 실행하니 아래와 같은 에러가 발생하였다

2024-12-06T10:36:18.732+09:00 ERROR 8256 --- [nboundChannel-6] t.d.g.e.WebSocketExceptionHandler        : WebSocket Exception: Could not read JSON: Cannot construct instance of `team9.ddang.member.oauth2.CustomOAuth2User`, problem: attributes cannot be empty
 at [Source: REDACTED (`StreamReadFeature.INCLUDE_SOURCE_IN_LOCATION` disabled); line: 4, column: 1]


문제점

그럼 뭐가 문제일까? Seucrity 에서 HTTP 와 WS 가 완전히 다르게 생성되기 때문이다. 헤더 접근 방식 부터 다른데WS 에서는 HTTP 형식의 헤더를 지원하지 않는다.

그래서 인증객체를 SecurityContext 에 등록해주지도 않았으니 당연하게도 해당 어노테이션으로는 원하는 인증객체를 불러올 수 없는 것 이다.



해결 과정

1. 헤더에 직접 접근

나는 우선 접근을 인증 객체를 얻지 않더라도 액세스 토큰에 포함되어있는 이메일에 접근하고자 하였다.

 	@MessageMapping("/api/v1/walk-alone")
    public void startWalk(SimpMessageHeaderAccessor headerAccessor, @RequestBody @Valid StartWalkRequest startWalkRequest) {
        String token = jwtService.extractAccessToken(headerAccessor).orElseThrow(() -> new AuthenticationException(TOKEN_NOT_FOUND));
        String email = jwtService.extractEmail(token).orElseThrow(() -> new AuthenticationException(TOKEN_DO_NOT_EXTRACT_EMAIL));
        walkLocationService.startWalk(email , startWalkRequest.toService());
    }

그래서 위와 같은 코드가 나오게 되었다.

jwtService 도 의존하고 있고 굉장히 컨트롤러에 많이 붙여놔서 굉장히 보기 안좋은 코드였다. 그래서 해당 부분은 고쳐야겠다고 생각했고 어떻게 하는게 좋을까 생각했는데,

웹소켓이 내가 맡은 부분 말고도 다른 부분에서 사용을 해서 AOP 로 email 을 추출하여 쓰레드 로컬에 저장하는 방식으로 구현하였다.



2. 쓰레드 로컬을 통해 저장

AuthenticationContext.java

public class AuthenticationContext {
    private static final ThreadLocal<String> emailHolder = new ThreadLocal<>();
    public static void setEmail(String email) {
        emailHolder.set(email);
    }
    public static String getEmail() {
        return emailHolder.get();
    }
    public static void clear() {
        emailHolder.remove();
    }
}

JwtAuthenticationAspect.java

@Component
@Aspect
public class JwtAuthenticationAspect {
    private final JwtService jwtService;
    public JwtAuthenticationAspect(JwtService jwtService) {
        this.jwtService = jwtService;
    }
    @Around("@annotation(ExtractEmail)")
    public Object authenticateToken(ProceedingJoinPoint joinPoint) throws Throwable {
        Object[] args = joinPoint.getArgs();
        SimpMessageHeaderAccessor headerAccessor = findHeaderAccessor(args);
        if (headerAccessor == null) {
            throw new IllegalArgumentException("SimpMessageHeaderAccessor not found in method arguments.");
        }
        
        String token = jwtService.extractAccessToken(headerAccessor)
                .orElseThrow(() -> new AuthenticationException(TOKEN_NOT_FOUND));
        String email = jwtService.extractEmail(token)
                .orElseThrow(() -> new AuthenticationException(TOKEN_DO_NOT_EXTRACT_EMAIL));

        AuthenticationContext.setEmail(email);
        return joinPoint.proceed();
    }
    private SimpMessageHeaderAccessor findHeaderAccessor(Object[] args) {
        for (Object arg : args) {
            if (arg instanceof SimpMessageHeaderAccessor) {
                return (SimpMessageHeaderAccessor) arg;
            }
        }
        return null;
    }
}
    @MessageMapping("/api/v1/walk-alone")
    @ExtractEmail
    public void startWalk(SimpMessageHeaderAccessor headerAccessor ,@RequestBody @Valid StartWalkRequest startWalkRequest) {
        walkLocationService.startWalk(AuthenticationContext.getEmail() , startWalkRequest.toService());
    }

코드 설명을 조금 하자면 `@ExtractEmail` 이라는 커스텀 어노테이션이 붙은 컨트롤러의 `headerAccesor` 에 접근 하여 액세스 토큰을 추출하고 이메일을 추출하여 그 값을 쓰레드로컬 에 저장하면 그것을 컨트롤러에서 접근하는 방식이다.

그래서 이렇게 좀 깔끔해진 코드가 생성되기는 하였다. 하지만 쓰레드 로컬을 사용하는 것도 그렇고 혹시 다른 방법으로 인증객체를 좀 더 깔끔하게 저장하는 방법은 없을까 고민이 생겼다.



3. Interceptor 를 통한 인증 객체 저장

StompHandler.java

.
.
.

 private void handleConnect(StompHeaderAccessor accessor) {
        String token = extractToken(accessor);

        if (token == null || !jwtService.isTokenValid(token)) {
            throw new IllegalArgumentException("유효하지 않은 토큰입니다.");
        }

        String email = jwtService.extractEmail(token)
                .orElseThrow(() -> new IllegalArgumentException("토큰에서 이메일을 추출할 수 없습니다."));

        Authentication authentication = new UsernamePasswordAuthenticationToken(email, null, Collections.emptyList());
        accessor.setUser(authentication);

        System.out.println("User connected: " + email);
    }

.
.
.



WebSocketConfig.java

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

    private final StompHandler stompHandler;

    @Override
    public void configureMessageBroker(MessageBrokerRegistry config) {
        config.enableSimpleBroker("/sub", "/queue");
        config.setApplicationDestinationPrefixes("/pub");
        config.setUserDestinationPrefix("/user");
    }

    @Override
    public void registerStompEndpoints(final StompEndpointRegistry registry) {
        registry
                .addEndpoint("/ws")
                .setAllowedOriginPatterns("*");
                //.withSockJS();
    }
    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {
        registration.interceptors(stompHandler);
    }
}

위와 WebSocketConfig 에 interceptors 를 등록해 놓고 accessor 에 Authentication 인증 객체를 저장해놓으면 해당 객체를 컨트롤러에서도 Principal 로 불러올 수 있다고 찾다보니 나왔었다...

하지만 왜인지 모르겠지만 accessor.setUser() 로 저장된 객체가 다음 요청에는 전혀 남아있지 않았었다. 그래서 결국 쓰레드 로컬이 아닌 Session 에 등록해놓고 사용하는 꼴이 되었다. 이것도 방법이라면 방법이지만 더 좋은 방법이 있을 것 같다.



결론

3가지 모두 다 가능한 방법이고 실제로 적용하고 있는 곳도 있을 것으로 생각된다. 하지만 인증 객체를

Spring 공식 문서를 보면

스프링 공식 문서 (WebSocket)

@EnableWebSocketSecurity 어노테이션을 통해 accessor.setUser() 에 등록되는 'simpUser' 를 인증 객체로 등록할 수 있다고 적혀있는 것으로 보아 Principal 을 통해 인증 객체를 받아올 수 있을 것 같다. 하지만 현재 시간이 부족하기도 하므로 이건 다음에 이어서 할 예정이다.

profile
MoonJaeGyeong
내 맘대로 끄적이는 개발 블로그

0개의 댓글