[Spring Boot][WebSocket + STOMP] 웹소켓 JWT 인증 및 파싱

1.  개요

모의 투자 사이트를 개발하면서 토스 증권을 참고해보았을 때 실제 주문도 STOMP 를 통해서 진행되는 것을 확인할 수 있었다. WebSocket 이라는 것 자체가 양방향 통신이 가능한 프로토콜이기에 이미 연결되어있는 웹소켓을 사용하지 않을 이유가 없다. 

 

모의 투자 사이트에서는 HTTP API 에 대해서 JWT 토큰을 사용해 사용자를 인증하고 있기에 STOMP 에서도 JWT 를 사용하고자 한다. 

 

2. JWT 헤더 파싱

const headers = {
  Authorization: `Bearer ${JWT_TOKEN}`,
};

try {
  const message = JSON.parse(messageJson); // JSON으로 변환
  stompClient.send(destinationPath, headers, JSON.stringify(message)); // 메시지 발행
} catch (error) {
  stompClient.send(destinationPath, headers, message);
}

 

클라이언트에서는 stompClient 를 사용해서 쉽게 헤더에 jwt 토큰을 담을 수 있다. 

 

@Component
@RequiredArgsConstructor
@Slf4j
public class StompHandler implements ChannelInterceptor {
	@Override
	public void postSend(Message<?> message, MessageChannel channel, boolean sent) {
		StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
		String sessionId = accessor.getSessionId();

		switch (Objects.requireNonNull(accessor.getCommand())) {
			case CONNECT -> log.info("CONNECT: " + message);
			case CONNECTED -> log.info("CONNECTED: " + message);
			case DISCONNECT -> log.info("DISCONNECT: " + message);
			case SUBSCRIBE -> log.info("SUBSCRIBE: " + message);
			case UNSUBSCRIBE -> log.info("UNSUBSCRIBE: " + sessionId);
			case SEND -> log.info("SEND: " + sessionId);
			case MESSAGE -> log.info("MESSAGE: " + sessionId);
			case ERROR -> log.info("ERROR: " + sessionId);
			default -> log.info("UNKNOWN: " + sessionId);
		}

	}
}

 

이후에는 이미 사용하고 있던 StompHandler 를 그대로 사용하였다. 

 

(1) postSend 와 preSend

클라이언트 메시지 발송 -> [서버 preSend] -> 서버 메시지 처리 -> [서버 postSend]

서버 메시지 발송 -> [서버 preSend] -> 클라이언트로 전송 -> [서버 postSend]

 

첫 번째를 서버로 들어오기에 인바운드 메시지, 두 번째를 서버에서 나가는 메시지이므로 아웃바운드 메시지라 한다. 

여기서 주문 접수를 위해서는 인바운드 메시지를 사용할 예정이고, 서버에서 메시지를 처리하기 전에 preSend 에서 Header 에 대한 처리가 필요하다.

 

@Component
@RequiredArgsConstructor
@Slf4j
public class StompHandler implements ChannelInterceptor {
	private final TokenService tokenService;
	private final UserRepository userRepository;
	private final UserDetailsServiceImpl userDetailsService;

	@Override
	public Message<?> preSend(Message<?> message, MessageChannel channel) {
		StompHeaderAccessor accessor = StompHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
		String sessionId = accessor.getSessionId();

		if (Objects.requireNonNull(accessor.getCommand()) == StompCommand.SEND) {
			String tokenHeader = accessor.getFirstNativeHeader("Authorization");
			if (tokenHeader == null) {
				log.error("No token found in message from session: " + sessionId);
				throw new InvalidJwtException("No token found");
			}

			String token = tokenHeader.replace("Bearer ", "");
			String userEmail = tokenService.validateAccessToken(token);

			UserDetails userDetails = userDetailsService.loadUserByUsername(userEmail);

			accessor.setUser(new UsernamePasswordAuthenticationToken(userDetails, null, userDetails.getAuthorities()));

		}

		return message;
	}

	...
}

 

Header 를 처리하기 위해서는 StompHeaderAccessor 를 사용한다. 

public class StompHeaderAccessor extends SimpMessageHeaderAccessor {
	...
}

/*
 * 메시징 프로토콜의 메시지 헤더를 다루는 기본 클래스
*/
public class SimpMessageHeaderAccessor extends NativeMessageHeaderAccessor {
	...
    public void setUser(@Nullable Principal principal) {
        this.setHeader("simpUser", principal);
        if (this.userCallback != null) {
            this.userCallback.accept(principal);
        }

    }
    ...
}

/*
 * 외부 소스에서 온 헤더들의 저장과 접근을 지원하는 MessageHeaderAccessor 의 하위 클래스
 * 외부 소스란 클라이언트(브라우저, 앱 등)에서 전송한 헤더를 의미
 * 이러한 헤더들을 외부 소스 헤더(Native Headers)로 취급한다
*/
public class NativeMessageHeaderAccessor extends MessageHeaderAccessor {
	...
}

 

결국 클라이언트에서 주입된 AccessToken 은 NativeMessageHeader 에서 찾을 수 있다. 그리고 헤더는 같은 이름으로 여러 값이 존재할 수 있기 때문에 getFirstNativeHeader() 를 통해 하나의 토큰 값만 가져오게 된다. 

 

Token 파싱 과정은 일반적인 헤더의 파싱 과정과 동일하게 진행된다. 

 

Spring Security 같은 경우는 Context Holder 에 로그인된 유저 정보를 넣어놓지만 웹소켓같은 경우에는 HTTP 와 달리 지속적인 연결을 가지므로 다른 맥락을 가지고 있다. HTTP 는 각 요청이 독립적이기에 매 요청마다 하나의 스레드가 할당되고, 요청이 끝나면 SecurityContext 도 정리된다. 반변 웹소켓은 하나의 session Id 에 대해서도 여러 스레드에서 동시에 처리될 수 있고, SecurityContext 정리 시점이 불분명하기에 일반적으로 권장되는 방법은 아니라고 한다.

 

일단 JWT 를 그래도 사용한다면 Token 으로 파싱된 이후 검증된 유저를 Message 에 넘겨줘야한다. 이후 Principal 객체를 헤더에 담음으로써 Message 를 넘겨받는 컨트롤러단에서 이를 사용할 수 있다.

(2) MessageMapping

@RestController
@RequiredArgsConstructor
public class BuyOrderMessageController {
	private final BuyOrderServiceImpl buyOrderServiceImpl;

	@MessageMapping("/order/buy")
	public void getStockPriceByWebsocket(@Payload BuyOrderRequest buyOrderRequest, Authentication authentication) {
		User user = ((UserDetailsImpl) authentication.getPrincipal()).getUser();

		buyOrderServiceImpl.createMarketOrder(user, buyOrderRequest);
	}
}

 

@MessageMapping 에서는 위 처럼 body 의 값과 헤더의 Principal 을 Authentication 에 매핑시킬 수 있다. 

 

 

SimpAnnotationMethodMessageHanlder 는 @MessageMapping 메서드를 찾고 실행하는 역할을 한다. 그렇다면 여기에서 헤더의 값을 매핑시켜줄테니 메서드를 더 살펴보면, 해당 클래스는 init 할 때 PrincipalMethodArgumentResolver 를 가지고 초기화한다.

 

여기서 추가된 Resolver 가 아까 SimpMessageHeaderAccessor 에 setUser 로 주입한 Principal 값을 가져오는 것을 확인할 수 있다. 

3. 결과 

STOMP 테스터를 사용해서 테스트해봤다. 

// 메시지 전송 버튼 클릭 시
$("#sendBtn").click(function () {
  const destinationPath = $("#destinationPath").val(); // 대상 경로 가져오기
  const messageJson = $("#message").val(); // JSON 형태의 메시지 가져오기

  const headers = {
    Authorization: `Bearer ${JWT_TOKEN}`,
  };

  try {
    const message = JSON.parse(messageJson); // JSON으로 변환
    stompClient.send(destinationPath, headers, JSON.stringify(message)); // 메시지 발행
  } catch (error) {
    stompClient.send(destinationPath, headers, message);
  }
});

 

물론 SEND 시에 JWT 를 헤더에 추가하는 과정도 거쳤다.

 

 

 

Order 테이블에도 주문이 추가되는 모습을 확인할 수 있다.