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

Gradle에 필요한 설정 import

  // WebSocket
  implementation("org.springframework.boot:spring-boot-starter-websocket")
  implementation("com.fasterxml.jackson.core:jackson-databind")
  implementation("com.fasterxml.jackson.module:jackson-module-kotlin")
  implementation("org.jetbrains.kotlin:kotlin-reflect")
  implementation("org.jetbrains.kotlin:kotlin-stdlib-jdk8")

  // Web Security
  implementation("org.springframework.boot:spring-boot-starter-security")

WebSocketConfig

에 STOMP 관련 설정을 해준다

package kpring.chat.global.config

import org.springframework.beans.factory.annotation.Value
import org.springframework.context.annotation.Bean
import org.springframework.context.annotation.Configuration
import org.springframework.messaging.simp.config.MessageBrokerRegistry
import org.springframework.web.cors.CorsConfiguration
import org.springframework.web.cors.UrlBasedCorsConfigurationSource
import org.springframework.web.filter.CorsFilter
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker
import org.springframework.web.socket.config.annotation.StompEndpointRegistry
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer
import org.springframework.web.socket.config.annotation.WebSocketTransportRegistration

@Configuration
@EnableWebSocketMessageBroker
class WebSocketConfig : 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 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)
  }
}

SecurityConfig

를 만들어서 WebSocket경로에 대해서 Spring Security 설정을 걸러주었다

package kpring.chat.global.config

import org.springframework.context.annotation.Bean
import org.springframework.context.annotation.Configuration
import org.springframework.security.config.annotation.web.builders.HttpSecurity
import org.springframework.security.config.annotation.web.configuration.EnableWebSecurity
import org.springframework.security.web.SecurityFilterChain

@Configuration
@EnableWebSecurity
class SecurityConfig {
  @Bean
  fun securityFilterChain(http: HttpSecurity): SecurityFilterChain {
    http
      .authorizeHttpRequests { auth ->
        auth
          .requestMatchers("/ws/**").permitAll() // WebSocket 경로를 인증 없이 허용
          .anyRequest().authenticated() // 나머지 요청은 인증 필요
      }
      .csrf { csrf -> csrf.disable() } // CSRF 보호 비활성화
      .cors { cors -> cors.disable() } // CORS 비활성화
      .httpBasic { basic -> basic.disable() } // 기본 인증 비활성화
    return http.build()
  }
}

근데 저렇게 하면 보안상에 이슈가 있지 않나 싶다

그 채팅방에서 탈퇴당한 사람도 채팅방 id만 있으면 subscribe 될 것 같기도 하고…

이 부분은 추후에 더 생각하고 개선해 봐야겠다

일단 생각나는 것은 시간이 지나면 바뀌는 채팅방 id당 코드를 만들어서 해당 채팅방에 소속된 유저인지 검증한 후 그 코드를 subscribe하게 하는 것인데 너무 비효율적인 것 같기도 하다

Controller

package kpring.chat.chat.api.v1

import kpring.chat.chat.service.ChatService
import kpring.chat.global.exception.ErrorCode
import kpring.chat.global.exception.GlobalException
import kpring.core.auth.client.AuthClient
import kpring.core.chat.chat.dto.request.CreateChatRequest
import kpring.core.chat.chat.dto.request.DeleteChatRequest
import kpring.core.chat.chat.dto.request.GetChatsRequest
import kpring.core.chat.chat.dto.request.UpdateChatRequest
import kpring.core.chat.model.ChatType
import kpring.core.server.client.ServerClient
import kpring.core.server.dto.request.GetServerCondition
import lombok.RequiredArgsConstructor
import org.slf4j.Logger
import org.slf4j.LoggerFactory
import org.springframework.messaging.handler.annotation.MessageMapping
import org.springframework.messaging.handler.annotation.Payload
import org.springframework.messaging.simp.SimpMessageHeaderAccessor
import org.springframework.messaging.simp.SimpMessagingTemplate
import org.springframework.stereotype.Controller
import org.springframework.validation.annotation.Validated

@Controller
@RequiredArgsConstructor
class WebSocketChatController(
  private val chatService: ChatService,
  private val authClient: AuthClient,
  private val serverClient: ServerClient,
  private val simpMessagingTemplate: SimpMessagingTemplate,
) {
  private val logger: Logger = LoggerFactory.getLogger(WebSocketChatController::class.java)

  @MessageMapping("/chat/create")
  fun createChat(
    @Payload @Validated request: CreateChatRequest,
    headerAccessor: SimpMessageHeaderAccessor,
  ) {
    val token = headerAccessor.getFirstNativeHeader("Authorization") ?: throw GlobalException(ErrorCode.MISSING_TOKEN)
    val userId = authClient.getTokenInfo(token).data!!.userId
    val contextId = request.contextId

    val result =
      when (request.type) {
        ChatType.ROOM -> chatService.createRoomChat(request, userId)
        ChatType.SERVER ->
          chatService.createServerChat(
            request,
            userId,
            serverClient.getServerList(token, GetServerCondition()).body!!.data!!,
          )
      }
    simpMessagingTemplate.convertAndSend("/topic/chatroom/$contextId", result)
  }

  @MessageMapping("/chat/update")
  fun updateChat(
    @Payload @Validated request: UpdateChatRequest,
    headerAccessor: SimpMessageHeaderAccessor,
  ) {
    val token = headerAccessor.getFirstNativeHeader("Authorization") ?: throw GlobalException(ErrorCode.MISSING_TOKEN)
    val userId = authClient.getTokenInfo(token).data!!.userId
    val contextId = request.contextId

    val result = chatService.updateChat(request, userId)
    simpMessagingTemplate.convertAndSend("/topic/chatroom/$contextId", result)
  }

  @MessageMapping("/chat/delete")
  fun deleteChat(
    @Payload @Validated request: DeleteChatRequest,
    headerAccessor: SimpMessageHeaderAccessor,
  ) {
    val token = headerAccessor.getFirstNativeHeader("Authorization") ?: throw GlobalException(ErrorCode.MISSING_TOKEN)
    val userId = authClient.getTokenInfo(token).data!!.userId
    val chatId = request.id
    val contextId = request.contextId

    val result = chatService.deleteChat(chatId, userId)
    simpMessagingTemplate.convertAndSend("/topic/chatroom/$contextId", result)
  }

  @MessageMapping("/chat")
  fun getChats(
    @Payload @Validated request: GetChatsRequest,
    headerAccessor: SimpMessageHeaderAccessor,
  ) {
    val token = headerAccessor.getFirstNativeHeader("Authorization") ?: throw GlobalException(ErrorCode.MISSING_TOKEN)
    val userId = authClient.getTokenInfo(token).data!!.userId
    val type = request.type
    val contextId = request.contextId
    val page = request.page
    val size = request.size

    val result =
      when (type) {
        ChatType.ROOM -> chatService.getRoomChats(contextId, userId, page, size)
        ChatType.SERVER ->
          chatService.getServerChats(
            contextId,
            userId,
            page,
            size,
            serverClient.getServerList(token, GetServerCondition()).body!!.data!!,
          )
      }
    logger.info("Chat result: $result")
    simpMessagingTemplate.convertAndSend("/topic/chatroom/$contextId", result)
  }
}

Front 연결 예시 코드

html안에 모든 것을 다 넣어버렸다

WebSocket Chat Application

Authorization Token:

 

 Send Message

Create Chat

RoomServerCreate Chat

Get Chats

RoomServer Get Chats

Update Chat

  Update Chat

Delete Chat

RoomServerDelete Chat
    let stompClient = null;

    const getToken = () => document.getElementById('authToken').value;
    const getChatRoomId = () => document.getElementById('chatroomId').value;

    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);

        stompClient.connect({ 'Authorization': `Bearer ${token}` }, function (frame) {
            console.log('Connected: ' + frame);

            // Send a message to verify access to the chat room
            stompClient.send(`/app/connect/${chatroomId}`, { 'Authorization': token }, {});

            // Subscribe to the chat room's topic if access is granted
            stompClient.subscribe(`/topic/chatroom/${chatroomId}`, {});
            enableButtonsAfterConnect(token, chatroomId);

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

    function enableButtonsAfterConnect(token, chatroomId) {
        document.getElementById("sendBtn").disabled = false;
        document.getElementById("createChatBtn").disabled = false;
        document.getElementById("getChatsBtn").disabled = false;
        document.getElementById("updateChatBtn").disabled = false;
        document.getElementById("deleteChatBtn").disabled = false;

        // Sending a message to the chat room
        document.getElementById("sendBtn").addEventListener("click", function() {
            const message = document.getElementById("msg").value;
            if (message.trim() !== "") {
                stompClient.send(`/app/chat/create`, { 'Authorization': token },
                    JSON.stringify({ content: message, type: "ROOM", chatRoomId: chatroomId })
                );
                document.getElementById("msg").value = "";  // Clear the input field after sending
                console.log('Message sent successfully!');
            }
        });

        // Creating a new chat message
        document.getElementById("createChatBtn").addEventListener("click", function() {
            const content = document.getElementById("chatContent").value;
            const type = document.getElementById("chatType").value;
            if (content.trim() !== "") {
                stompClient.send(`/app/chat/create`, { 'Authorization': token },
                    JSON.stringify({ content: content, type: type, contextId: chatroomId })
                );
                document.getElementById("chatContent").value = "";  // Clear input field after creation
                console.log('Chat created successfully!');
            }
        });

        // Fetching chat messages
        document.getElementById("getChatsBtn").addEventListener("click", function() {
            const id = document.getElementById("getId").value;
            const type = document.getElementById("getChatType").value;
            const page = document.getElementById("getPage").value;
            const size = 2;
            if (id && page) {
                stompClient.send(`/app/chat`, { 'Authorization': token },
                    JSON.stringify({ contextId: id, type: type, page: parseInt(page), size: parseInt(size) })
                );
            }
        });

        // Updating a chat
        document.getElementById("updateChatBtn").addEventListener("click", function() {
            const chatId = document.getElementById("updateChatId").value;
            const content = document.getElementById("updateContent").value;
            if (chatId && content.trim() !== "") {
                stompClient.send(`/app/chat/update`, { 'Authorization': token },
                    JSON.stringify({ contextId: chatroomId, id: chatId, content: content, type: "ROOM" })
                );
                document.getElementById("updateChatId").value = "";
                document.getElementById("updateContent").value = "";
                console.log('Chat updated successfully!');
            }
        });

        // Deleting a chat
        document.getElementById("deleteChatBtn").addEventListener("click", function() {
            const chatId = document.getElementById("deleteChatId").value;
            const type = document.getElementById("deleteChatType").value;
            if (chatId) {
                stompClient.send(`/app/chat/delete`, { 'Authorization': token },
                    JSON.stringify({ contextId: chatroomId, id: chatId, type: type })
                );
                document.getElementById("deleteChatId").value = "";
                console.log('Chat deleted successfully!');
            }
        });
    }

    document.addEventListener('DOMContentLoaded', function() {
        document.getElementById("sendBtn").disabled = true;
        document.getElementById("createChatBtn").disabled = true;
        document.getElementById("getChatsBtn").disabled = true;
        document.getElementById("updateChatBtn").disabled = true;
        document.getElementById("deleteChatBtn").disabled = true;

        document.getElementById("connectBtn").addEventListener("click", function() {
            connectWebSocket();
        });
    });

이걸 실행시키면 http://localhost:63343/에서 돌아가니까 이걸 WebSecurityConfig에서 걸러주면 된다

원래 RestAPI로 Controller를 구성할 때에는 Create나 Update Delete는 굳이 DTO를 리턴하지 않았는데 웹소켓을 개발해보니 관점이 바뀌어서 신기했다

같은 채팅방 안에 있는 다른 사람이 Create했을때 subscribe로 CreateResponseDto를 받으면 전체 List로 업데이트하지 않고도 프론트에서 실시간 채팅 반영이 가능해서 훨씬 효율적인 것 같다

ChatType

MessageType을 만들어서 Response에 속성으로 넣었다

채팅의 Type에 따라 프론트에서 처리할 수 있게 해서 트래픽 부하를 줄일 수 있을 것으로 기대한다

Update, Delete 등 상태 변경이 일어났을 때에 해당 채팅의 id로 프론트에서 찾아서 그 채팅만 변경하는것이 list를 아예 새로 읽어오는 것보다 트래픽 부하가 적고 속도가 빠를 것으로 예상한다

package kpring.core.chat.model

enum class MessageType(val type: String) {
  ENTER("ENTER"),
  INVITATION("INVITATION"),
  CHAT("CHAT"),
  UPDATE("UPDATE"),
  DELETE("DELETE"),
  OUT("OUT"),
}

ENTER("ENTER") : 누군가가 채팅방에 입장했을 때 수신되는 ChatResponse

INVITATION("INVITATION") : 누군가가 채팅방에 초대링크를 공유했을 때 수신되는 ChatResponse

CHAT("CHAT") : 일반 채팅

UPDATE("UPDATE") : 이 Type의 ChatResponse가 수신되면 해당 채팅의 id로 프론트 페이지에서 찾아서 그 채팅만 변경할 수 있습니다.

DELETE("DELETE") : 이 Type의 ChatResponse가 수신되면 해당 채팅의 id로 프론트 페이지에서 찾아서 그 채팅만 삭제할 수 있습니다.

OUT("OUT") : 누군가가 채팅방을 나갔을 때 가는 ChatResponse

ChatResponse 에 sender와 messageType 추가

data class ChatResponse(
  val id: String,
  val sender: String,
  val messageType: MessageType,
  val isEdited: Boolean,
  val sentAt: String,
  val content: String,

 

이후 아무나 채팅방 id만 알면 SUBSCRIBE 요청을 날려서 도청이 가능할 것 같아 보안을 다소 개선해 주었다

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