Spring

AOP로 jwt token 검증하기 - 커스텀 어노테이션, ThreadLocal 사용

oyatplum 2025. 1. 3. 17:03

들어가기 전 문제 상황

@PostMapping("/logout")
@Operation(summary = "로그아웃 API", description = "로그아웃")
public ApiResponse<String> logout(@RequestHeader("Authorization") String token) {
    if (!token.startsWith("Bearer ")) {
        return ApiResponse.onFailure("400", "유효하지 않은 토큰 형식입니다.", null);
    }

    String tokenWithoutBearer = token.substring(7);
    try {
        Long userId = jwtUtil.getUserIdFromToken(tokenWithoutBearer);
        if (userId == null) {
            return ApiResponse.onFailure("400", "존재하지 않거나 만료된 토큰입니다.", null);
        }
        userService.logout(userId);
        return ApiResponse.onSuccess("로그아웃에 성공했습니다.");
    } catch (Exception e) {
        return ApiResponse.onFailure("400", e.getMessage(), null);
    }
}
@PostMapping("/mytravels/memo/{travelId}")
    @Operation(summary = "메모 기록 API", description = "특정 여행기에서의 여행조각 추가")
    public ApiResponse<TravelResponseDto.CreateTripPieceResultDto> createTripPieceMemo(@RequestBody TravelRequestDto.MemoDto request, @PathVariable("travelId") Long travelId, @RequestHeader("Authorization") String token){
        ...
        String tokenWithoutBearer = token.substring(7);
        TripPiece tripPiece = travelService.createMemo(travelId, request, tokenWithoutBearer);
        return ApiResponse.onSuccess(TravelConverter.toCreateTripPieceResultDto(tripPiece));
    }
@Transactional
    public TripPiece createMemo(Long travelId, TravelRequestDto.MemoDto request, String token) {
        Long userId = jwtUtil.getUserIdFromToken(token);
        User user = userRepository.findById(userId).orElseThrow(() -> new RuntimeException("user not found"));
        ...
}

 

리팩토링 이전 토큰을 검증하는 로직이 controller와 service에서 난잡하게 작성되고 있었다.

팀원끼리 상의하지 않은 상태에서 진행되었고 통일 방식을 모르기도 했었다.

단순히 jwtUtil에서 토큰의 유효성을 검증하고는 있었지만, Bearer를 제거하는 방식이나 jwtUtil을 통해 유효성을 검증하는 방식 등 호출 위치 모두 제각각 달랐다.

 

따라서 이 과정을 커스텀 어노테이션을 사용해 하나의 로직으로 통일하여 사용하면 효율적일 것 같다는 생각이 들었다.

 

그러다 aop와 threadLocal을 함께 사용하면 된다는 것을 알게 됐다.

 

 


AOP란

aspect-oriented programming

즉, 말그대로 관점 지향 프로그래밍이다.

흩어진 관심을 aspect로 모듈화한다고 생각하면 된다.

 

(aop에 대한 자세한 설명을 아래 포스팅 참고!)

https://engkimbs.tistory.com/entry/%EC%8A%A4%ED%94%84%EB%A7%81AOP

 


 

AOP로 jwt token 검증하기 (커스텀 어노테이션, ThreadLocal 사용)

1. 커스텀 어노테이션 정의 (@ValidateToken)

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface ValidateToken {
}

 

위와 같이 커스텀 어노테이션을 작성했고 타켓은 METHOD로 메서드 레벨에서만 사용 가능하고, @Retention(RetentionPolicy.RUNTIME)을 통해 런타임에도 어노테이션 정보가 유지되도록 했다.

 

2. 토큰 검증 로직 (TokenValidationAspect)

@Aspect
@Component
@RequiredArgsConstructor
public class TokenValidationAspect {
    private final JWTUtil jwtUtil;
    private final UserRepository userRepository;

    @Before("@annotation(umc.TripPiece.validation.annotation.ValidateToken)")
    public void validateToken() {
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        if (attributes == null){
            throw new NotFoundHandler(ErrorStatus.REQUEST_INFO_NOT_FOUND);
        }
        HttpServletRequest request = attributes.getRequest();

        String token = request.getHeader("Authorization");
        if (token == null || !token.startsWith("Bearer ")) {
            throw new BadRequestHandler(ErrorStatus.INVALID_TOKEN);
        }

        String tokenWithoutBearer = token.substring(7);
        Long userId = jwtUtil.getUserIdFromToken(tokenWithoutBearer);

        boolean userExists = userRepository.existsById(userId);
        if(!userExists) {
            throw new NotFoundHandler(ErrorStatus.NOT_FOUND_USER);
        }

        UserContext.setUserId(userId);
    }
    @After("@annotation(umc.TripPiece.validation.annotation.ValidateToken)")
    public void clearUserContext() {
        UserContext.clear();
    }
}

 

@Before는 메서드가 실행되기 전에 동작하는 부분이다. 즉, @ValidateToken 어노테이션이 붙은 메서드가 실행되기 전에 작동한다.

 

참고로

@Before("@annotation(ValidateToken)")
    public void validateToken(JoinPoint joinPoint) {
        // 메서드 인자에서 Authorization 헤더 추출
        Object[] args = joinPoint.getArgs();
        String token = Arrays.stream(args)
                              .filter(arg -> arg instanceof String && ((String) arg).startsWith("Bearer "))
                              .map(arg -> (String) arg)
                              .findFirst()
                              .orElseThrow(() -> new BadRequestHandler(ErrorStatus.INVALID_TRAVEL_PARARM));

        // Bearer 제거
        String tokenWithoutBearer = token.substring(7);

 

위와 같이 JoinPoint를 사용해서 메서드 파라미터로 Authorization 헤더를 추출해 Bearer 토큰을 찾는 방식도 있지만

HttpServletRequset를 이용하여 요청 객체에서 직접 헤더를 가져왔다.

 

이후 토큰에서 Bearer를 제거한 뒤 userRepository에서 사용자 존재 여부를 확인했다.

 

이후 ThreadLocal를 사용하여 UserContext를 통해 요청 스레드별로 userId를 저장했다.

 

3. 유저 컨텍스트 (UserContext)

public class UserContext {
    private static final ThreadLocal<Long> userIdHolder = new ThreadLocal<>();

    public static void setUserId(Long userId) {
        userIdHolder.set(userId);
    }

    public static Long getUserId() {
        return userIdHolder.get();
    }

    public static void clear() {
        userIdHolder.remove();
    }
}

 

위와 같이 UserContext class를 ThreadLocal를 사용해 작성했다.

(자세한 설명은 아래 포스팅 참고)

https://velog.io/@semi-cloud/%EC%8A%A4%ED%94%84%EB%A7%81-%EA%B3%A0%EA%B8%89%ED%8E%B81-ThreadLocal

 

ThreadLocal은 사용 후 값을 제거해줘야 하기 때문에

TokenValidationAspect 하단을 보면

@After("@annotation(umc.TripPiece.validation.annotation.ValidateToken)")
public void clearUserContext() {
    UserContext.clear();
}

이와 같이 @After로 @ValidateToken 어노테이션이 붙은 메서드 실행 이후에는 UserContext.clear()로 스레드 값을 제거했다.

 

 

이는 위와 같이 로그를 통해 확인해줬다.

 

 

4. Controller, Service에서 사용

따라서

@PostMapping(value = "/mytravels", consumes = "multipart/form-data")
@ValidateToken
@Operation(summary = "여행 생성 API", description = "여행 시작하기")
public ApiResponse<TravelResponseDto.Create> createTravel(@Valid @RequestPart("data") TravelRequestDto.Create request, @RequestPart("thumbnail") MultipartFile thumbnail){
    if (thumbnail == null || thumbnail.isEmpty()) {
        throw new BadRequestHandler(ErrorStatus.MISSING_TRAVEL_THUMBNAIL);
    }
    TravelResponseDto.Create response = travelService.createTravel(request, thumbnail);

    return ApiResponse.onSuccess(response);
}
 @Transactional
    public TravelResponseDto.Create createTravel(TravelRequestDto.Create request, MultipartFile thumbnail) {
        Long userId = UserContext.getUserId();
        User user = userRepository.findById(userId).orElseThrow(() -> new NotFoundHandler(ErrorStatus.NOT_FOUND_USER));

위와 같이 파라미터로 토큰을 받지 않고  @ValidateToken 어노테이션을 사용하면 토큰 유효성을 검증할 수 있다.

또한, 이렇게 반환받은 userId는 service에서 UserContext.getUserId()를 통해 가져와 사용할 수 있다.

 

 

 

위의 과정을 통해

1. 중앙화된 토큰 검증 - AOP 활용

: @ValidateToken 어노테이션을 활용하여 비즈니스 로직과 토큰 검증 로직이 분리됨

2. 재사용성

: @ValidateToken 으로 필요한 컨트롤러에서 해당 로직 사용 가능

3. 유지보수 용이

: 토큰 검증 로직 변경 시 TokenValidationAspect만 수정

와 같은 장점이 가능해졌고..

 

aop에 대해 알아보다 Interceptor에 대해서도 알게 되었는데.. 이를 이용한 리팩토링은 다음 포스팅에서...

 

'Spring' 카테고리의 다른 글

Spring Security 사용하기  (2) 2025.01.09
Interceptor로 토큰 전역 처리하기  (0) 2025.01.07
커스텀 어노테이션으로 중복 코드 개선하기  (2) 2025.01.03
스프링 기초4  (0) 2024.02.01
스프링 기초3  (3) 2024.01.28