𝒘𝒆𝒃𝒔𝒐𝒄𝒌𝒆𝒕에서 사용자 정보 가져오기 오류

❎ 문제발생

클라이언트에서 메시지를 보낼때 오류가 발생하였다. 

fun processMessage(message: MessageDto, user: UserPrincipal): Message {

    val chatRoom = chatRoomRepository.findById(message.chatRoomId)
        .orElseThrow { IllegalArgumentException("Chat room not found") }

    val member = oAuthRepository.findById(user.memberId)
        .orElseThrow { IllegalArgumentException("Member not found") }

    return messageRepository.save(
        Message(
            content = message.content,
            userId = member,
            chatRoom = chatRoom
        )
    )
}

해당 코드에서 user.memberId에 해당하는 유저가 존재하지 않다고 말해준다.

 

🆘 해결 과정

디버깅을 해보니 UserPrincipal에 제대로된 사용자 정보가 담기지 않는다는 것을 확인할 수 있었다.

 

@MutationMapping
fun createChatRoom(
    @AuthenticationPrincipal user: UserPrincipal,
    @Argument roomName: String,
    @Argument roomType: String
): ChatRoom {
    return chatService.createChatRoom(user.memberId, roomName, roomType)
}

 

반면 채팅방 생성에서는 문제없이 사용자 정보를 잘 가져오는것을 확인할 수 있다.

 

따라서 http통신에서는 @AuthenticationPrincipal을 통해 사용자 정보를 가져올수 있지만

websocket통신에서는 다른방법을 사용해야한다.

 

현재 웹소켓을 통해 토큰을 파라미터로 전달하고 있기 때문에 해당 파라미터를 가로채서 토큰을 추출하고 사용자 정보를 담아야한다.

WebSocket 메세지 메소드(@MessageMapping)은 Http요청과 별개로 동작한다.

package hjp.hjchat.infra.security.jwt.websocket

import hjp.hjchat.infra.security.jwt.JwtAuthenticationToken
import hjp.hjchat.infra.security.jwt.JwtTokenManager
import hjp.hjchat.infra.security.jwt.UserPrincipal
import org.springframework.http.server.ServerHttpRequest
import org.springframework.http.server.ServerHttpResponse
import org.springframework.security.core.context.SecurityContextHolder
import org.springframework.web.socket.WebSocketHandler
import org.springframework.web.socket.server.HandshakeInterceptor
import org.springframework.stereotype.Component
import java.nio.file.AccessDeniedException


@Component
class JwtWebSocketHandshakeInterceptor(
    private val jwtTokenManager: JwtTokenManager,
) : HandshakeInterceptor {

    override fun beforeHandshake(
        request: ServerHttpRequest,
        response: ServerHttpResponse,
        wsHandler: WebSocketHandler,
        attributes: MutableMap<String, Any>
    ): Boolean {
        val queryParams = request.uri.query
        val token = queryParams?.substringAfter("token=")?.substringBefore("&")

        if (token == null || token.isBlank()) {
            throw AccessDeniedException("JWT Token is missing in query parameters")
        }

        jwtTokenManager.validateToken(token).onSuccess { claims ->
            val memberId = claims.body.subject.toLong()
            val memberRole = claims.body.get("memberRole", String::class.java)

            val userPrincipal = UserPrincipal(
                memberId = memberId,
                memberRole = setOf(memberRole)
            )


            val authentication = JwtAuthenticationToken(
                userPrincipal = userPrincipal,
                details = null
            )
            
            SecurityContextHolder.getContext().authentication = authentication
            attributes["userPrincipal"] = userPrincipal
        }.onFailure {
            throw AccessDeniedException("Invalid JWT Token")
        }

        return true
    }


    override fun afterHandshake(
        request: ServerHttpRequest,
        response: ServerHttpResponse,
        wsHandler: WebSocketHandler,
        exception: Exception?
    ) {

    }
}

 

✔ 해결

@AuthenticationPrincipal user: UserPrincipal,

밑에 코드로 변경

headerAccessor: SimpMessageHeaderAccessor

HandShakeIntercepter를 통해 사용자 정보를 attributes에 직접 저장을 하고  저장한 값을 직접적으로 가져오는 식으로 문제를 해결하였다.

 

@MessageMapping("/send")
@Transactional
fun sendMessage(
    @Payload message: MessageDto,
    headerAccessor: SimpMessageHeaderAccessor
) {
    val userPrincipal = headerAccessor.sessionAttributes?.get("userPrincipal") as? UserPrincipal
        ?: throw IllegalArgumentException("UserPrincipal not found in session attributes")

    val savedMessage = chatService.processMessage(message, userPrincipal)

    messagingTemplate.convertAndSend("/topic/chatroom/${message.chatRoomId}", savedMessage.toResponse())
}

 

테스트

 

 

 

 

전체적인 흐름

 

1. 클라이언트에서 WebSocket에 연결을 요청

const socket = new SockJS(`http://localhost:8080/ws?token=${accessToken}`); // Query Parameter로 토큰 전달
ws = Stomp.over(socket);

쿼리 파라미터에 포함된 accessToken을 통해 JWT 토큰을 서버에 전달

 

2. WebSocket Handshake에서 해당 토큰 인증

package hjp.hjchat.infra.security.jwt.websocket

import hjp.hjchat.infra.security.jwt.JwtAuthenticationToken
import hjp.hjchat.infra.security.jwt.JwtTokenManager
import hjp.hjchat.infra.security.jwt.UserPrincipal
import org.springframework.http.server.ServerHttpRequest
import org.springframework.http.server.ServerHttpResponse
import org.springframework.security.core.context.SecurityContextHolder
import org.springframework.web.socket.WebSocketHandler
import org.springframework.web.socket.server.HandshakeInterceptor
import org.springframework.stereotype.Component
import java.nio.file.AccessDeniedException


@Component
class JwtWebSocketHandshakeInterceptor(
    private val jwtTokenManager: JwtTokenManager,
) : HandshakeInterceptor {

    override fun beforeHandshake(
        request: ServerHttpRequest,
        response: ServerHttpResponse,
        wsHandler: WebSocketHandler,
        attributes: MutableMap<String, Any>
    ): Boolean {
        val queryParams = request.uri.query
        val token = queryParams?.substringAfter("token=")?.substringBefore("&")

        if (token == null || token.isBlank()) {
            throw AccessDeniedException("JWT Token is missing in query parameters")
        }

        jwtTokenManager.validateToken(token).onSuccess { claims ->
            val memberId = claims.body.subject.toLong()
            val memberRole = claims.body.get("memberRole", String::class.java)

            val userPrincipal = UserPrincipal(
                memberId = memberId,
                memberRole = setOf(memberRole)
            )


            val authentication = JwtAuthenticationToken(
                userPrincipal = userPrincipal,
                details = null
            )

            SecurityContextHolder.getContext().authentication = authentication
            attributes["userPrincipal"] = userPrincipal
        }.onFailure {
            throw AccessDeniedException("Invalid JWT Token")
        }

        return true
    }


    override fun afterHandshake(
        request: ServerHttpRequest,
        response: ServerHttpResponse,
        wsHandler: WebSocketHandler,
        exception: Exception?
    ) {

    }
}

 

JWT토큰을 검증하고 생성한 UserPrincipal 객체를 WebSocket Attributes에 저장, WebSocket 연결 이후에도 메시지 핸들러가 사용자 정보를 사용할 수 있도록 SecurityContextHolder에 인증정보 저장

 

3. WebSocket 메시지 기능 수행

@MessageMapping("/send")
@Transactional
fun sendMessage(
    @Payload message: MessageDto,
    headerAccessor: SimpMessageHeaderAccessor
) {
    val userPrincipal = headerAccessor.sessionAttributes?.get("userPrincipal") as? UserPrincipal
        ?: throw IllegalArgumentException("UserPrincipal not found in session attributes")

    val savedMessage = chatService.processMessage(message, userPrincipal)

    messagingTemplate.convertAndSend("/topic/chatroom/${message.chatRoomId}", savedMessage.toResponse())
}

클라이언트가 STOMP 메시지를 서버로 전송하면  @MessageMapping 메서드가 이를 처리함

 

4. 메시지를 저장하고 채팅방을 생성하면 MySQL DB에서 저장하고 관리

5. 실시간 메시지 브로드캐스팅

messagingTemplate.convertAndSend("/topic/chatroom/${message.chatRoomId}", savedMessage.toResponse())

 

메시지가 저장된 후, SimpMessagingTemplate를 사용하여 /topic/chatroom/{chatRoomId}로 메시지를 브로드 캐스트함