웹소켓 Subscribe시에 Interceptor에서 채팅방 권한을 조회하도록 함

지난번에 이렇게 해서 웹소켓 연결을 했는데 

Spring과 STOMP를 이용해서 웹소켓 연결하기

연결해두고 보니 아무나 특정 채팅방이나 서버를 Subscribe할 수 있는 구조라 보안 이슈가 있을 것 같았다

 

그래서 Interceptor에서 웹소켓 연결이 Subscribe인 경우에 채팅방에 속해있는지 / 서버에 속해있는지 검증하는 로직을 추가했다

 

그러기 위해서 검증하는 부분이 Service Layer에도 있고 Interceptor에도 있으면 코드가 중복되는것 같아 AccessVerifier라는 util 디렉토리 아래의 클래스를 하나 만들어 주었다

package kpring.chat.global.util

import kpring.chat.chatroom.repository.ChatRoomRepository
import kpring.chat.global.exception.ErrorCode
import kpring.chat.global.exception.GlobalException
import kpring.core.server.dto.ServerSimpleInfo
import org.springframework.stereotype.Component

@Component
class AccessVerifier(
  private val chatRoomRepository: ChatRoomRepository,
) {
  fun verifyServerAccess(
    servers: List<ServerSimpleInfo>,
    serverId: String,
  ) {
    servers.forEach { info ->
      if (info.id.equals(serverId)) {
        return
      }
    }
    throw GlobalException(ErrorCode.FORBIDDEN_SERVER)
  }

  fun verifyChatRoomAccess(
    chatRoomId: String,
    userId: String,
  ) {
    if (!chatRoomRepository.existsByIdAndMembersContaining(chatRoomId, userId)) {
      throw GlobalException(ErrorCode.FORBIDDEN_CHATROOM)
    }
  }
}

 

 

그 후 ChannelInterceptor를 상속받은 webSocketInterceptor인터셉터에 검증 로직을 넣어주었다

@Configuration
@EnableWebSocketMessageBroker
class WebSocketConfig(
  val authClient: AuthClient,
  val serverClient: ServerClient,
  val accessVerifier: AccessVerifier,
) : WebSocketMessageBrokerConfigurer {
  @Value("\\${url.front}")
  val frontUrl: String = ":63343"

  override fun registerStompEndpoints(registry: StompEndpointRegistry) {
    registry.addEndpoint("/ws").setAllowedOrigins(frontUrl).withSockJS()
  }

  override fun configureMessageBroker(config: MessageBrokerRegistry) {
    config.setApplicationDestinationPrefixes("/app")
    config.enableSimpleBroker("/topic")
  }

  override fun configureWebSocketTransport(registry: WebSocketTransportRegistration) {
    registry.setMessageSizeLimit(4 * 8192)
    registry.setTimeToFirstMessage(30000)
  }

  @Bean
  fun webSocketInterceptor(): ChannelInterceptor {
    return object : ChannelInterceptor {
      override fun preSend(
        message: Message<*>,
        channel: MessageChannel,
      ): Message<*> {
        val simpMessageType = SimpMessageHeaderAccessor.getMessageType(message.headers)
        return when (simpMessageType) {
          SimpMessageType.CONNECT -> handleConnectMessage(message)
          SimpMessageType.SUBSCRIBE -> {
            handleSubscribeMessage(message)
            message
          }
          else -> message
        }
      }
    }
  }

  private fun handleConnectMessage(message: Message<*>): Message<*> {
    val headerAccessor = SimpMessageHeaderAccessor.wrap(message)
    val token =
      headerAccessor.getFirstNativeHeader("Authorization")
        ?.removePrefix("Bearer ")
        ?: throw GlobalException(ErrorCode.MISSING_TOKEN)

    val userId =
      authClient.getTokenInfo(token).data?.userId
        ?: throw GlobalException(ErrorCode.INVALID_TOKEN)

    val principal = UsernamePasswordAuthenticationToken(userId, null, emptyList())
    SimpMessageHeaderAccessor.getAccessor(message, SimpMessageHeaderAccessor::class.java)?.user = principal

    return message
  }

  private fun handleSubscribeMessage(message: Message<*>) {
    val headerAccessor = SimpMessageHeaderAccessor.wrap(message)

    val token =
      headerAccessor.getFirstNativeHeader("Authorization")
        ?.removePrefix("Bearer ")
        ?: throw GlobalException(ErrorCode.MISSING_TOKEN)

    val contextId =
      headerAccessor.getFirstNativeHeader("contextId")
        ?: throw GlobalException(ErrorCode.MISSING_CONTEXTID)

    val context =
      headerAccessor.getFirstNativeHeader("context")
        ?: throw GlobalException(ErrorCode.MISSING_CONTEXT)

    val userId =
      authClient.getTokenInfo(token).data?.userId
        ?: throw GlobalException(ErrorCode.INVALID_TOKEN)

    when (context) {
      ChatType.ROOM.toString() -> {
        accessVerifier.verifyChatRoomAccess(contextId, userId)
      }
      ChatType.SERVER.toString() -> {
        val serverList =
          serverClient.getServerList(token, GetServerCondition())
            .body?.data ?: throw GlobalException(ErrorCode.SERVER_ERROR)
        accessVerifier.verifyServerAccess(serverList, contextId)
      }
      else -> throw GlobalException(ErrorCode.INVALID_CONTEXT)
    }
  }

  override fun configureClientInboundChannel(registration: ChannelRegistration) {
    registration.interceptors(webSocketInterceptor())
  }

  @Bean
  fun corsFilter(): CorsFilter {
    val config = CorsConfiguration()
    config.allowCredentials = true
    config.addAllowedOrigin(frontUrl)
    config.addAllowedHeader("*")
    config.addAllowedMethod("*")

    val source: UrlBasedCorsConfigurationSource = UrlBasedCorsConfigurationSource()
    source.registerCorsConfiguration("/**", config)
    return CorsFilter(source)
  }
}

 

 

프론트에서도 CONNECT면 jwt token만 보내고

SUBSCRIBE할 때에는 jwt token과 함께 해당 채팅방/서버 id와 채팅방인지 서버인지 type을 함께 보내도록 했다

function connectWebSocket() {
        const token = getToken();
        const chatroomId = getChatRoomId();

        if (!token || !chatroomId) {
            alert('Authorization token and Chatroom ID are required!');
            return;
        }

        const socket = new SockJS('<http://localhost:8081/ws>');
        stompClient = Stomp.over(socket);

        const connectHeaders = {
            'Authorization': `Bearer ${token}`
        };

        stompClient.connect(connectHeaders, function (frame) {
            console.log('Connected: ' + frame);

            const subscribeHeaders = {
                'Authorization': `Bearer ${token}`,
                'contextId': chatroomId,
                'context': 'ROOM'
            };

            stompClient.subscribe(`/topic/chatroom/${chatroomId}`, function(message) {
                console.log("Message received:", message);
            }, subscribeHeaders);

            enableButtonsAfterConnect(token, chatroomId);
        }, function(error) {
            console.error('Connection error:', error);
            alert('Connection failed. Please check your token and chat room ID.');
        });
    }