From 402af01806442d02ab3359f8b94c6b3192704437 Mon Sep 17 00:00:00 2001 From: yeseong0412 Date: Tue, 12 Nov 2024 12:13:45 +0900 Subject: [PATCH 1/6] =?UTF-8?q?Refactor=20::=20Stomp=20User=20=EC=A0=95?= =?UTF-8?q?=EB=B3=B4=20=EC=A0=84=EB=8B=AC=20=EB=B0=A9=EB=B2=95=20=EB=B3=80?= =?UTF-8?q?=EA=B2=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../websocket/config/StompWebSocketConfig.kt | 43 +++++++------------ .../controller/StompRabbitMQController.kt | 11 +++-- .../websocket/util/SecurityUtils.kt | 13 ++++++ 3 files changed, 33 insertions(+), 34 deletions(-) create mode 100644 src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/util/SecurityUtils.kt diff --git a/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/config/StompWebSocketConfig.kt b/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/config/StompWebSocketConfig.kt index bc25ff7e1..f7d9283be 100644 --- a/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/config/StompWebSocketConfig.kt +++ b/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/config/StompWebSocketConfig.kt @@ -1,8 +1,8 @@ package com.seugi.api.domain.chat.presentation.websocket.config import com.seugi.api.domain.chat.application.service.chat.room.ChatRoomService +import com.seugi.api.domain.chat.presentation.websocket.util.SecurityUtils import com.seugi.api.domain.member.application.exception.MemberErrorCode -import com.seugi.api.global.auth.jwt.JwtUserDetails import com.seugi.api.global.auth.jwt.JwtUtils import com.seugi.api.global.exception.CustomException import org.springframework.beans.factory.annotation.Value @@ -15,7 +15,6 @@ import org.springframework.messaging.simp.config.ChannelRegistration import org.springframework.messaging.simp.config.MessageBrokerRegistry import org.springframework.messaging.simp.stomp.StompHeaderAccessor import org.springframework.messaging.support.ChannelInterceptor -import org.springframework.messaging.support.MessageBuilder import org.springframework.messaging.support.MessageHeaderAccessor import org.springframework.util.AntPathMatcher import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker @@ -50,9 +49,9 @@ class StompWebSocketConfig( val accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor::class.java)!! when (accessor.messageType) { - SimpMessageType.CONNECT -> handleConnect(message, accessor) + SimpMessageType.CONNECT -> handleConnect(accessor) SimpMessageType.SUBSCRIBE -> handleSubscribe(accessor) - SimpMessageType.UNSUBSCRIBE, SimpMessageType.DISCONNECT -> handleUnsubscribeOrDisconnect() + SimpMessageType.UNSUBSCRIBE, SimpMessageType.DISCONNECT -> handleUnsubscribeOrDisconnect(accessor) else -> {} } return message @@ -60,20 +59,13 @@ class StompWebSocketConfig( }) } - private fun handleConnect(message: Message<*>, accessor: StompHeaderAccessor) { + private fun handleConnect(accessor: StompHeaderAccessor) { val authToken = accessor.getNativeHeader("Authorization")?.firstOrNull() if (authToken != null && authToken.startsWith("Bearer ")) { val auth = jwtUtils.getAuthentication(authToken) - val userDetails = auth.principal as? JwtUserDetails - val userId: String? = userDetails?.id?.value?.toString() - - if (userId != null) { - val simpAttributes = SimpAttributesContextHolder.currentAttributes() - simpAttributes.setAttribute("user-id", userId) - MessageBuilder.createMessage(message.payload, accessor.messageHeaders) - } else { - throw CustomException(MemberErrorCode.MEMBER_NOT_FOUND) - } + accessor.user = auth + } else { + throw CustomException(MemberErrorCode.MEMBER_NOT_FOUND) } } @@ -81,25 +73,20 @@ class StompWebSocketConfig( accessor.destination?.let { val simpAttributes = SimpAttributesContextHolder.currentAttributes() simpAttributes.setAttribute("sub", it.substringAfterLast(".")) - val userId = simpAttributes.getAttribute("user-id") as String chatRoomService.sub( - userId = userId.toLong(), + userId = SecurityUtils.getUserId(accessor.user), roomId = it.substringAfterLast(".") ) } } - private fun handleUnsubscribeOrDisconnect() { - val simpAttributes = SimpAttributesContextHolder.currentAttributes() - val userId = simpAttributes.getAttribute("user-id") as String? - val roomId = simpAttributes.getAttribute("sub") as String? - userId?.let { - roomId?.let { - chatRoomService.unSub( - userId = userId.toLong(), - roomId = it - ) - } + private fun handleUnsubscribeOrDisconnect(accessor: StompHeaderAccessor) { + accessor.destination?.let { + val simpAttributes = SimpAttributesContextHolder.currentAttributes() + chatRoomService.unSub( + userId = SecurityUtils.getUserId(accessor.user), + roomId = simpAttributes.getAttribute("sub").toString() + ) } } } \ No newline at end of file diff --git a/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/controller/StompRabbitMQController.kt b/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/controller/StompRabbitMQController.kt index 83a36cabc..9730f004b 100644 --- a/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/controller/StompRabbitMQController.kt +++ b/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/controller/StompRabbitMQController.kt @@ -2,21 +2,20 @@ package com.seugi.api.domain.chat.presentation.websocket.controller import com.seugi.api.domain.chat.application.service.message.MessageService import com.seugi.api.domain.chat.presentation.websocket.dto.ChatMessageDto +import com.seugi.api.domain.chat.presentation.websocket.util.SecurityUtils import org.springframework.messaging.handler.annotation.MessageMapping -import org.springframework.messaging.simp.SimpAttributesContextHolder import org.springframework.stereotype.Controller +import java.security.Principal @Controller class StompRabbitMQController( - private val messageService: MessageService + private val messageService: MessageService, ) { @MessageMapping("chat.message") - fun send(chat: ChatMessageDto) { - val simpAttributes = SimpAttributesContextHolder.currentAttributes() - val userId = simpAttributes.getAttribute("user-id") as String? - messageService.sendAndSaveMessage(chat, userId!!.toLong()) + fun send(chat: ChatMessageDto, principal: Principal) { + messageService.sendAndSaveMessage(chat, SecurityUtils.getUserId(principal)) } } \ No newline at end of file diff --git a/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/util/SecurityUtils.kt b/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/util/SecurityUtils.kt new file mode 100644 index 000000000..ed81b2875 --- /dev/null +++ b/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/util/SecurityUtils.kt @@ -0,0 +1,13 @@ +package com.seugi.api.domain.chat.presentation.websocket.util + +import com.seugi.api.global.auth.jwt.JwtUserDetails +import org.springframework.security.authentication.UsernamePasswordAuthenticationToken +import java.security.Principal + +object SecurityUtils { + + fun getUserId(principal: Principal?): Long { + return (principal as? UsernamePasswordAuthenticationToken)?.principal.let { it as? JwtUserDetails }?.member?.id?.value + ?: -1 + } +} \ No newline at end of file From 0c041fb47fc12bfb0c18eeecec69ad0a566fb7c2 Mon Sep 17 00:00:00 2001 From: yeseong0412 Date: Tue, 19 Nov 2024 17:04:48 +0900 Subject: [PATCH 2/6] =?UTF-8?q?Refactor=20::=20Response=20=EA=B5=AC?= =?UTF-8?q?=EC=A1=B0=20=EC=B6=94=EC=83=81=ED=99=94?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Response 베이스 클래스와 인터페이스를 추가하고, ErrorResponse, BaseResponse 가 상속받도록 변경함 --- .../com/seugi/api/global/response/BaseResponse.kt | 12 ++++++------ .../com/seugi/api/global/response/ErrorResponse.kt | 11 +++++++++++ .../seugi/api/global/response/ResponseInterface.kt | 8 ++++++++ 3 files changed, 25 insertions(+), 6 deletions(-) create mode 100644 src/main/kotlin/com/seugi/api/global/response/ErrorResponse.kt create mode 100644 src/main/kotlin/com/seugi/api/global/response/ResponseInterface.kt diff --git a/src/main/kotlin/com/seugi/api/global/response/BaseResponse.kt b/src/main/kotlin/com/seugi/api/global/response/BaseResponse.kt index ffbe3c4d2..d204b3508 100644 --- a/src/main/kotlin/com/seugi/api/global/response/BaseResponse.kt +++ b/src/main/kotlin/com/seugi/api/global/response/BaseResponse.kt @@ -7,13 +7,13 @@ import org.springframework.http.HttpStatus @JsonInclude(JsonInclude.Include.NON_NULL) data class BaseResponse( - val status: Int = HttpStatus.OK.value(), - val success: Boolean = true, - val state: String? = "OK", - val message: String, - val data: T? = null + override val status: Int = HttpStatus.OK.value(), + override val success: Boolean = true, + override val state: String = "OK", + override val message: String, + val data: T? = null, -) { + ) : ResponseInterface { // errorResponse constructor constructor(code: CustomErrorCode) : this( diff --git a/src/main/kotlin/com/seugi/api/global/response/ErrorResponse.kt b/src/main/kotlin/com/seugi/api/global/response/ErrorResponse.kt new file mode 100644 index 000000000..acbfac30d --- /dev/null +++ b/src/main/kotlin/com/seugi/api/global/response/ErrorResponse.kt @@ -0,0 +1,11 @@ +package com.seugi.api.global.response + +import com.fasterxml.jackson.annotation.JsonInclude + +@JsonInclude(JsonInclude.Include.NON_NULL) +data class ErrorResponse( + override val status: Int = 1000, + override val success: Boolean = false, + override val state: String = "Error", + override val message: String, +) : ResponseInterface diff --git a/src/main/kotlin/com/seugi/api/global/response/ResponseInterface.kt b/src/main/kotlin/com/seugi/api/global/response/ResponseInterface.kt new file mode 100644 index 000000000..1bb82a47f --- /dev/null +++ b/src/main/kotlin/com/seugi/api/global/response/ResponseInterface.kt @@ -0,0 +1,8 @@ +package com.seugi.api.global.response + +interface ResponseInterface { + val status: Int + val success: Boolean + val state: String + val message: String +} \ No newline at end of file From 4e086a1292b5761c3df9066c53344d5e455ba4c1 Mon Sep 17 00:00:00 2001 From: yeseong0412 Date: Wed, 20 Nov 2024 13:43:13 +0900 Subject: [PATCH 3/6] =?UTF-8?q?Refactor=20::=20=ED=8C=A8=ED=82=A4=EC=A7=80?= =?UTF-8?q?=20=EC=9C=84=EC=B9=98=20=EB=B3=80=EA=B2=BD=20=EB=B0=8F=20?= =?UTF-8?q?=EA=B8=B0=EB=B3=B8=EA=B0=92=20=EB=B3=80=EA=B2=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit RabbitMQConfig 패키지 위치 변경 ErrorResponse Status 기본값 변경 --- .../presentation/websocket => global}/config/RabbitMQConfig.kt | 2 +- src/main/kotlin/com/seugi/api/global/response/ErrorResponse.kt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) rename src/main/kotlin/com/seugi/api/{domain/chat/presentation/websocket => global}/config/RabbitMQConfig.kt (96%) diff --git a/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/config/RabbitMQConfig.kt b/src/main/kotlin/com/seugi/api/global/config/RabbitMQConfig.kt similarity index 96% rename from src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/config/RabbitMQConfig.kt rename to src/main/kotlin/com/seugi/api/global/config/RabbitMQConfig.kt index c0add3b53..29708a5f3 100644 --- a/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/config/RabbitMQConfig.kt +++ b/src/main/kotlin/com/seugi/api/global/config/RabbitMQConfig.kt @@ -1,4 +1,4 @@ -package com.seugi.api.domain.chat.presentation.websocket.config +package com.seugi.api.global.config import org.springframework.amqp.core.BindingBuilder import org.springframework.amqp.core.Queue diff --git a/src/main/kotlin/com/seugi/api/global/response/ErrorResponse.kt b/src/main/kotlin/com/seugi/api/global/response/ErrorResponse.kt index acbfac30d..1dc9942fc 100644 --- a/src/main/kotlin/com/seugi/api/global/response/ErrorResponse.kt +++ b/src/main/kotlin/com/seugi/api/global/response/ErrorResponse.kt @@ -4,7 +4,7 @@ import com.fasterxml.jackson.annotation.JsonInclude @JsonInclude(JsonInclude.Include.NON_NULL) data class ErrorResponse( - override val status: Int = 1000, + override val status: Int = 4500, override val success: Boolean = false, override val state: String = "Error", override val message: String, From eaeeef5033cd8c95fbb0775e8cfe66ab4843fae4 Mon Sep 17 00:00:00 2001 From: yeseong0412 Date: Thu, 21 Nov 2024 19:56:13 +0900 Subject: [PATCH 4/6] =?UTF-8?q?Feat=20::=20=EC=86=8C=EC=BC=93=20=EB=B9=84?= =?UTF-8?q?=EC=A6=88=EB=8B=88=EC=8A=A4=20=ED=95=B8=EB=93=A4=EB=A7=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 임시 --- .../exception/CustomSocketExceptionHandler.kt | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) create mode 100644 src/main/kotlin/com/seugi/api/global/exception/CustomSocketExceptionHandler.kt diff --git a/src/main/kotlin/com/seugi/api/global/exception/CustomSocketExceptionHandler.kt b/src/main/kotlin/com/seugi/api/global/exception/CustomSocketExceptionHandler.kt new file mode 100644 index 000000000..adcff91a4 --- /dev/null +++ b/src/main/kotlin/com/seugi/api/global/exception/CustomSocketExceptionHandler.kt @@ -0,0 +1,46 @@ +package com.seugi.api.global.exception + +import com.seugi.api.global.response.ErrorResponse +import org.springframework.messaging.handler.annotation.MessageExceptionHandler +import org.springframework.messaging.simp.SimpAttributesContextHolder +import org.springframework.messaging.simp.SimpMessagingTemplate +import org.springframework.web.bind.annotation.ControllerAdvice +import java.net.SocketException +import java.security.Principal + +@ControllerAdvice +class CustomSocketExceptionHandler( +// private val template: RabbitMessagingTemplate + private val template: SimpMessagingTemplate, +) { + + @MessageExceptionHandler(SocketException::class) + fun handleSocketException(principal: Principal, ex: Exception) { + template.convertAndSendToUser( + principal.name, + "/error", + ErrorResponse(status = 4500, message = ex.cause?.message ?: "Socket Error") + ) + } + + @MessageExceptionHandler(RuntimeException::class) + fun handleRuntimeException(principal: Principal, ex: Exception) { + val simpAttributes = SimpAttributesContextHolder.currentAttributes() + println( + principal.name + " | " + simpAttributes.getAttribute("sub").toString() + " | " + simpAttributes.sessionId + ) + template.convertAndSendToUser( + principal.name, + "/queue/errors", + ErrorResponse(status = 4500, message = ex.cause?.message ?: "Socket Error") + ) + } + +// @Throws(IOException::class) +// private fun removeSession(message: Message<*>) { +// val stompHeaderAccessor = StompHeaderAccessor.wrap(message) +// val sessionId = stompHeaderAccessor.sessionId +// } + + +} \ No newline at end of file From f9a7cad3822a8127e91580bfa170eca546591ab6 Mon Sep 17 00:00:00 2001 From: yeseong0412 Date: Thu, 21 Nov 2024 20:43:55 +0900 Subject: [PATCH 5/6] =?UTF-8?q?Fix=20::=20=EC=86=8C=EC=BC=93=20=EC=98=88?= =?UTF-8?q?=EC=99=B8=20=ED=95=B8=EB=93=A4=EB=9F=AC=20=EC=88=98=EC=A0=95=20?= =?UTF-8?q?=EB=B0=8F=20=EC=84=B8=EC=85=98=20=EC=A0=9C=EA=B1=B0=20=EA=B8=B0?= =?UTF-8?q?=EB=8A=A5=20=EC=B6=94=EA=B0=80?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SocketException 핸들러에 세션 제거 기능을 추가했습니다. 또한, runtime 예외 메시지를 표시하는 방식을 개선했습니다. --- .../exception/CustomSocketExceptionHandler.kt | 33 ++++++++++--------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/src/main/kotlin/com/seugi/api/global/exception/CustomSocketExceptionHandler.kt b/src/main/kotlin/com/seugi/api/global/exception/CustomSocketExceptionHandler.kt index adcff91a4..9a3e215e6 100644 --- a/src/main/kotlin/com/seugi/api/global/exception/CustomSocketExceptionHandler.kt +++ b/src/main/kotlin/com/seugi/api/global/exception/CustomSocketExceptionHandler.kt @@ -2,45 +2,46 @@ package com.seugi.api.global.exception import com.seugi.api.global.response.ErrorResponse import org.springframework.messaging.handler.annotation.MessageExceptionHandler -import org.springframework.messaging.simp.SimpAttributesContextHolder +import org.springframework.messaging.Message import org.springframework.messaging.simp.SimpMessagingTemplate +import org.springframework.messaging.simp.stomp.StompHeaderAccessor import org.springframework.web.bind.annotation.ControllerAdvice +import java.io.IOException import java.net.SocketException import java.security.Principal @ControllerAdvice class CustomSocketExceptionHandler( -// private val template: RabbitMessagingTemplate - private val template: SimpMessagingTemplate, + private val template: SimpMessagingTemplate ) { + private val bindingUrl = "/queue/errors" + @MessageExceptionHandler(SocketException::class) - fun handleSocketException(principal: Principal, ex: Exception) { + fun handleSocketException(message: Message<*>, principal: Principal, ex: Exception) { + removeSession(message) template.convertAndSendToUser( principal.name, - "/error", + bindingUrl, ErrorResponse(status = 4500, message = ex.cause?.message ?: "Socket Error") ) } @MessageExceptionHandler(RuntimeException::class) fun handleRuntimeException(principal: Principal, ex: Exception) { - val simpAttributes = SimpAttributesContextHolder.currentAttributes() - println( - principal.name + " | " + simpAttributes.getAttribute("sub").toString() + " | " + simpAttributes.sessionId - ) template.convertAndSendToUser( principal.name, - "/queue/errors", - ErrorResponse(status = 4500, message = ex.cause?.message ?: "Socket Error") + bindingUrl, + ErrorResponse(status = 4500, message = ex.cause?.stackTraceToString() ?: "Socket Error") ) } -// @Throws(IOException::class) -// private fun removeSession(message: Message<*>) { -// val stompHeaderAccessor = StompHeaderAccessor.wrap(message) -// val sessionId = stompHeaderAccessor.sessionId -// } + @Throws(IOException::class) + private fun removeSession(message: Message<*>) { + val stompHeaderAccessor = StompHeaderAccessor.wrap(message) + val sessionId = stompHeaderAccessor.sessionId + stompHeaderAccessor.sessionAttributes?.remove(sessionId) + } } \ No newline at end of file From 75ef6057853b16bbab6eef7d5b1ba6b36f46845b Mon Sep 17 00:00:00 2001 From: yeseong0412 Date: Thu, 21 Nov 2024 20:44:25 +0900 Subject: [PATCH 6/6] =?UTF-8?q?Feat=20::=20STOMP=20=EC=98=A4=EB=A5=98=20?= =?UTF-8?q?=ED=95=B8=EB=93=A4=EB=9F=AC=20=EC=B6=94=EA=B0=80?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit STOMP 웹소켓 오류 처리를 위한 StompErrorHandler 추가. StompWebSocketConfig에 에러 핸들러 설정 및 사용자 목적지 접두사 설정. --- .../websocket/config/StompWebSocketConfig.kt | 14 +++-- .../websocket/handler/StompErrorHandler.kt | 61 +++++++++++++++++++ 2 files changed, 71 insertions(+), 4 deletions(-) create mode 100644 src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/handler/StompErrorHandler.kt diff --git a/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/config/StompWebSocketConfig.kt b/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/config/StompWebSocketConfig.kt index f7d9283be..afd3d78d4 100644 --- a/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/config/StompWebSocketConfig.kt +++ b/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/config/StompWebSocketConfig.kt @@ -1,6 +1,7 @@ package com.seugi.api.domain.chat.presentation.websocket.config import com.seugi.api.domain.chat.application.service.chat.room.ChatRoomService +import com.seugi.api.domain.chat.presentation.websocket.handler.StompErrorHandler import com.seugi.api.domain.chat.presentation.websocket.util.SecurityUtils import com.seugi.api.domain.member.application.exception.MemberErrorCode import com.seugi.api.global.auth.jwt.JwtUtils @@ -28,11 +29,13 @@ class StompWebSocketConfig( private val jwtUtils: JwtUtils, private val chatRoomService: ChatRoomService, @Value("\${spring.rabbitmq.host}") private val rabbitmqHost: String, + private val stompErrorHandler: StompErrorHandler ) : WebSocketMessageBrokerConfigurer { override fun registerStompEndpoints(registry: StompEndpointRegistry) { registry.addEndpoint("/stomp/chat") .setAllowedOrigins("*") + registry.setErrorHandler(stompErrorHandler) } override fun configureMessageBroker(registry: MessageBrokerRegistry) { @@ -41,6 +44,7 @@ class StompWebSocketConfig( registry.enableStompBrokerRelay("/queue", "/topic", "/exchange", "/amq/queue") .setRelayHost(rabbitmqHost) .setVirtualHost("/") + registry.setUserDestinationPrefix("/user") } override fun configureClientInboundChannel(registration: ChannelRegistration) { @@ -73,10 +77,12 @@ class StompWebSocketConfig( accessor.destination?.let { val simpAttributes = SimpAttributesContextHolder.currentAttributes() simpAttributes.setAttribute("sub", it.substringAfterLast(".")) - chatRoomService.sub( - userId = SecurityUtils.getUserId(accessor.user), - roomId = it.substringAfterLast(".") - ) + if (it.contains(".")) { + chatRoomService.sub( + userId = SecurityUtils.getUserId(accessor.user), + roomId = it.substringAfterLast(".") + ) + } } } diff --git a/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/handler/StompErrorHandler.kt b/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/handler/StompErrorHandler.kt new file mode 100644 index 000000000..b7bc2750d --- /dev/null +++ b/src/main/kotlin/com/seugi/api/domain/chat/presentation/websocket/handler/StompErrorHandler.kt @@ -0,0 +1,61 @@ +package com.seugi.api.domain.chat.presentation.websocket.handler + +import com.fasterxml.jackson.core.JsonProcessingException +import com.fasterxml.jackson.databind.ObjectMapper +import com.seugi.api.global.response.ErrorResponse +import io.jsonwebtoken.ExpiredJwtException +import io.jsonwebtoken.MalformedJwtException +import io.jsonwebtoken.UnsupportedJwtException +import org.springframework.context.annotation.Configuration +import org.springframework.messaging.Message +import org.springframework.messaging.MessageDeliveryException +import org.springframework.messaging.simp.stomp.StompCommand +import org.springframework.messaging.simp.stomp.StompHeaderAccessor +import org.springframework.messaging.support.MessageBuilder +import org.springframework.security.access.AccessDeniedException +import org.springframework.web.socket.messaging.StompSubProtocolErrorHandler +import java.nio.charset.StandardCharsets +import java.security.SignatureException + +@Configuration +class StompErrorHandler(private val objectMapper: ObjectMapper) : StompSubProtocolErrorHandler() { + + override fun handleClientMessageProcessingError(clientMessage: Message?, ex: Throwable): Message? { + + return when (ex) { + is MessageDeliveryException -> { + when (val cause = ex.cause) { + is AccessDeniedException -> { + sendErrorMessage(ErrorResponse(status = 4403, message = "Access denied")) + } + else -> { + if (isJwtException(cause)) { + sendErrorMessage(ErrorResponse(status = 4403, message = cause?.message ?: "JWT Exception")) + } else { + sendErrorMessage(ErrorResponse(status = 4403, message = cause?.stackTraceToString() ?: "Unhandled exception")) + } + } + } + } + else -> { + sendErrorMessage(ErrorResponse(status = 4400, message = ex.message ?: "Unhandled root exception")) + } + } + } + + private fun isJwtException(ex: Throwable?): Boolean { + return ex is SignatureException || ex is ExpiredJwtException || ex is MalformedJwtException || ex is UnsupportedJwtException || ex is IllegalArgumentException + } + + private fun sendErrorMessage(errorResponse: ErrorResponse): Message { + val headers = StompHeaderAccessor.create(StompCommand.ERROR).apply { + message = errorResponse.message + } + return try { + val json = objectMapper.writeValueAsString(errorResponse) + MessageBuilder.createMessage(json.toByteArray(StandardCharsets.UTF_8), headers.messageHeaders) + } catch (e: JsonProcessingException) { + MessageBuilder.createMessage(errorResponse.message.toByteArray(StandardCharsets.UTF_8), headers.messageHeaders) + } + } +} \ No newline at end of file