Spring+Kotlin) 웹소켓에 Jwt Filter 적용 및 Security 이용

2024.06.25

원래는 WebSocket의 HandshakeInterceptor를 정의하여 함수 beforeHandshake 에서 header에 담긴 jwt를 확인-검증하였다. 하지만 채팅 메시지를 보내는 과정에서 userId를 앱이나 웹에 노출시키고 싶지 않기 때문에 웹소켓 통신에도 UserDetails를 추가하여 SecurityContext에 올리기로 하였다.

override fun beforeHandshake(
        request: ServerHttpRequest,
        response: ServerHttpResponse,
        wsHandler: WebSocketHandler,
        attributes: MutableMap<String, Any>
    ): Boolean {
        val accessToken = (request as ServletServerHttpRequest)
            .headers["AccessToken"]
            ?.firstOrNull()

        if(accessToken.isNullOrEmpty()) {
            logger.info { "HandShake Interceptor) Token Not Found" }
            return false
        } else if(!tokenManager.validateToken(accessToken)) {
            logger.info { "HandShake Interceptor) Invalid Token" }
            return false
        }
       
        return true
    }

여기에 UserId를 체크하고 UserDetails를 Context에 추가하는 코드를 더해준다.

override fun beforeHandshake(
        request: ServerHttpRequest,
        response: ServerHttpResponse,
        wsHandler: WebSocketHandler,
        attributes: MutableMap<String, Any>
    ): Boolean {
        val accessToken = (request as ServletServerHttpRequest)
            .headers["AccessToken"]
            ?.firstOrNull()

        if(accessToken.isNullOrEmpty()) {
            logger.info { "HandShake Interceptor) Token Not Found" }
            return false
        } else if(!tokenManager.validateToken(accessToken)) {
            logger.info { "HandShake Interceptor) Invalid Token" }
            return false
        }

        val userId = tokenManager.getUserIdFromToken(accessToken).toLong()
        val user = userService.findUser(userId)
        if(user != null) {
            val userDetails = UserDetailsImpl(user)
            val authentication = UsernamePasswordAuthenticationToken(userDetails, null, userDetails.authorities)
            SecurityContextHolder.getContext().authentication = authentication

            return true
        } else {
            logger.info { "HandShake Interceptor) Invalid User Id" }
            return false
        }
    }

2024.06.28

위에서 의도했던 코드들이 작동하지 않는다. 그 이유는 다음과 같았다.

@AuthenticationPrincipal is specifically for REST/MVC methods.

UserDetails는 Rest와 Mvc 메소드에서 접근가능하다는 것이다. 따라서 위 코드는 작동하지 않았다. 그래서 25일부터 오늘까지 계속 다양한 방법을 시도했었다.
그 시도들을 알아보기 전에 요구사항을 살펴보자.

1. UserId를 client에게 노출시키고 싶지 않다.
2. UserId를 이용하여 ChatMessage를 저장한다.

-> UserId 프로퍼티를 어딘가에 숨기고 websocket 통신 중 메시지를 서버에서 받았을때 가져와야한다.

방법 1) Principal 이용

websocket session에 참여 중인 사용자들을 Principal을 이용하여 구분할 수 있다. Map<String, User>와 같은 형태로 이용할 수 있지만 결국 특정 가능한 string을 유저에게 보내줘야 한다. 이는 userId를 암호화하여 해결할 수 있겠다.

방법 2) Client가 보낼때 Header에 추가

사용자가 채팅방을 불러올때 부르는 API에 UserId, UserName을 암호화한 값을 json body에 넣어 보내주고 이를 client는 header에 추가하여 메시지를 전송한다.

방법 3) Client 인터셉터에서 access token 헤더 추가

방법2에서 RetrofitClient 클래스에서 정의한 interceptor가 작동하지 않아 Stomp Library의 코드를 뜯어보았더니, 새로운 request를 생성 후 통신을 하는 걸 알게 되었다. 따라서 새로운 request를 가져오지 않고 기존 header가 추가된 request를 이용하게 해볼 수 있을 것 같다.

방법1과 방법2를 수행하였으며 다음 주에 방법3을 시도해보아야겠다.

2024.07.19

고민하다가 다른 방법이 하나 또 생각났다.

방법 4) attribute 이용

  방법 1에서 이용한 것과 비슷하다. beforeHandShake에서 access_token을 확인한다. 인증된 유저라면 session Id와 임의의 겹치지 않는 정수를 담은 암호화된 데이터를 generate한다. 이를 cookie에 넣고 또한 Map<"sessionId:랜덤 정수", User>의 형태로 attribute에 저장한다.
  유저가 메시지를 보내고 서버에서 받을때 쿠키 값을 확인하여 암호화된 데이터를 얻는다. 이 데이터를 부호화하여 roomId와 랜덤 정수를 얻는다. attribute에 "sessionId:랜덤 정수" 값을 key로 하여 value인 User를 얻는다.
  이후 Disconnect할때 해당 유저의 쿠키 값과 attribute 값을 지워준다.

step1) Jwt 확인 및 인증 후 attribute에 추가

beforeHandShake 중 인증된 이후

val sessionId = request.servletRequest.session.id
val random = Random.nextInt(0, Int.MAX_VALUE)
val encryption = sessionUserEncryption.encrypt(sessionId, random)
(response as ServletServerHttpResponse).servletResponse.addCookie(
	Cookie("SessionToUser", encryption.toString())
)
val user = userService.findUser(tokenManager.getUserIdFromToken(accessToken).toLong())!!
attributes[encryption] = SessionUser(user.id!!, user.nickname, user.email!!)

SessionUser.kt

data class SessionUser(
    val userId: Long,
    val userName: String,
    val userEmail: String
)

step2) 메시지 수신 후 room id 및 user id 가져오기
하던 도중 stomp의 end point에서 쿠키에 접근할 수 없음을 알았다....

해결

CustomInterceptor.kt

class CustomInterceptor: DefaultHandshakeHandler() {

    @Autowired
    private lateinit var sessionUserEncryption: SessionUserEncryption

    override fun determineUser(
        request: ServerHttpRequest,
        wsHandler: WebSocketHandler,
        attributes: MutableMap<String, Any>
    ): Principal {
        println("It's a determineUser!")
        var accessToken = ""
        (request as ServletServerHttpRequest).servletRequest.cookies?.let {
            for(cookie in it) {
                if(cookie.name == "access_token") accessToken = cookie.value
            }
        }

        val sessionId = request.servletRequest.session.id
        val random = Random.nextInt(0, Int.MAX_VALUE)

        println("sessionId: $sessionId, random: $random")

        val user = userService.findUser(tokenManager.getUserIdFromToken(accessToken).toLong())!!
        return SessionUser(userId = user.id!!, userName = user.nickname, userEmail = user.email!!)
    }
}

SessinUser.kt

data class SessionUser(
    val userId: Long,
    val userName: String,
    val userEmail: String
): Principal {
    override fun getName() = userName
}

HandShake가 성공한 이후 DefaultHandShakeHandler를 상속한 CustomeInterceptor의 determineUser 함수가 호출된다. 이 determineUser가 반환한 SessionUser: Principal 객체는 메시지의 end point로 전달되며 이를 파라미터로 받을 수 있다.

End Point

@MessageMapping("/chat") //여기로 전송되면 메서드 호출 -> WebSocketConfig prefixes 에서 적용한건 앞에 생략
    fun chat(rawMessage: RawMessage, principal: Principal): ChatResponseDto {

        val sessionUser = principal as SessionUser

        val dto = chatService.addChat(rawMessage, sessionUser.userId)
        val roomId = dto.roomId
        firebaseManager.sendToTopic(
            rawMessage.topicUrl, dto
        )

        simpleMessage.convertAndSend("/sub/chat/${roomId}", Json.encodeToString(dto))

        return dto
    }

이로써 user id를 노출하지 않고 유저를 특정하여 message를 저장할 수 있다.