Merge pull request #26 from SafeNet-2024/feature/add-security

[feat] 웹소켓 연결 및 메시지 전송시 JWT 토큰 검증 과정 추가

src/main/java/com/SafeNet/Backend/domain/message/controller/MessageController.java

@@ -3,6 +3,7 @@
 import com.SafeNet.Backend.domain.message.dto.MessageDto;
 import com.SafeNet.Backend.domain.message.service.MessageService;
 import com.SafeNet.Backend.domain.messageroom.service.MessageRoomService;
+import com.SafeNet.Backend.global.auth.JwtTokenProvider;
 import com.SafeNet.Backend.global.pubsub.RedisPublisher;
 import io.swagger.v3.oas.annotations.Operation;
 import io.swagger.v3.oas.annotations.Parameter;
@@ -13,58 +14,42 @@
 import io.swagger.v3.oas.annotations.tags.Tag;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;  // 로깅을 위해 추가
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.http.ResponseEntity;
 import org.springframework.messaging.handler.annotation.MessageMapping;
+import org.springframework.security.access.AccessDeniedException;
 import org.springframework.stereotype.Controller;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.PathVariable;
 import org.springframework.web.bind.annotation.RequestHeader;
 
-import java.time.LocalDateTime;
 import java.util.List;
+
 @Slf4j
 @RequiredArgsConstructor
 @Controller
 @Tag(name = "Message", description = "Message API")
 public class MessageController {
+    @Autowired
+    private JwtTokenProvider jwtTokenProvider;
+
     private final RedisPublisher redisPublisher;
     private final MessageRoomService messageRoomService;
     private final MessageService messageService;
 
-    // websocket "/pub/chat/message"로 들어오는 메시징을 처리한다.
 
+    // websocket "/pub/chat/message"로 들어오는 메시징을 처리한다.
     @MessageMapping("/chat/message")
-    @Operation(summary = "메시지 발송", description = "WebSocket을 통해 메시지를 발송한다")
     public void message(@RequestHeader(name = "ACCESS_TOKEN", required = false) String accessToken,
-                        @RequestHeader(name = "REFRESH_TOKEN", required = false) String refreshToken,
                         MessageDto messageDto) {
 
-        // 로깅 추가
-        log.info("Received message:");
-        log.info(" - sender: {}", messageDto.getSender());
-        log.info(" - roomId: {}", messageDto.getRoomId());
-        log.info(" - message: {}", messageDto.getMessage());
-        log.info(" - sentTime: {}", messageDto.getSentTime());
-        log.info("ACCESS_TOKEN: {}", accessToken);
-        log.info("REFRESH_TOKEN: {}", refreshToken);
-
-        // 클라이언트 채팅방(topic) 입장, 대화를 위해 리스너와 연동
-        messageRoomService.enterMessageRoom(messageDto.getRoomId());
-
-        // MessageDto 객체를 빌더를 통해 생성하고 현재 시간을 설정
-        MessageDto messageWithTime = MessageDto.builder()
-                .sender(messageDto.getSender())
-                .roomId(messageDto.getRoomId())
-                .message(messageDto.getMessage())
-                .sentTime(LocalDateTime.now().toString())
-                .build();
-
-        // Websocket에 발행된 메시지를 redis로 발행한다(publish)
-        // 해당 쪽지방을 구독(subscribe)한 클라이언트에게 메시지가 실시간 전송
-        redisPublisher.publish(messageRoomService.getTopic(messageDto.getRoomId()), messageWithTime);
+        // Access Token 검증
+        if (accessToken == null || !jwtTokenProvider.validateToken(accessToken)) { // 메시지 전송 전 유효한 토큰인지 검증
+            throw new AccessDeniedException("Invalid or expired token");
+        }
 
-        // DB와 Redis에 메시지 저장
-        messageService.saveMessage(messageDto);
+        // 메시지 전송 로직 호출
+        messageRoomService.handleMessage(messageDto.getRoomId(), messageDto.getSender(), messageDto);
     }
 
     // 대화 내역 조회

src/main/java/com/SafeNet/Backend/domain/message/service/MessageService.java

@@ -34,6 +34,12 @@ public MessageService(
     public void saveMessage(MessageDto messageDto) {
         MessageRoom messageRoom = messageRoomRepository.findByRoomId(messageDto.getRoomId())
                 .orElseThrow(() -> new IllegalArgumentException("해당 쪽지방이 존재하지 않습니다."));
+
+        if (!messageRoom.isFirstMessageSent()) {
+            messageRoom.setFirstMessageSent(true);
+            messageRoomRepository.save(messageRoom); // 첫 메시지 여부를 true로 업데이트
+        }
+
         Message message = Message.builder()
                 .sender(messageDto.getSender())
                 .messageRoom(messageRoom)

src/main/java/com/SafeNet/Backend/domain/messageRoom/entity/MessageRoom.java

@@ -36,6 +36,9 @@ public class MessageRoom {
     @Column(nullable = false)
     private String receiver; // 채팅방 수신자
 
+    @Column(nullable = false)
+    private boolean firstMessageSent = false; // 보낸 메시지가 존재하는지
+
     @CreationTimestamp
     @Column(name = "created_at", nullable = false, updatable = false)
     private LocalDateTime createdAt; // 현재 시간 자동 할당
@@ -50,4 +53,8 @@ public class MessageRoom {
     @ManyToOne(fetch = LAZY)
     @JoinColumn(name = "post_id", nullable = false)
     private Post post;
+
+    public void setFirstMessageSent(boolean firstMessageSent) {
+        this.firstMessageSent = firstMessageSent;
+    }
 }

src/main/java/com/SafeNet/Backend/domain/messageRoom/service/MessageRoomService.java

@@ -2,16 +2,19 @@
 
 import com.SafeNet.Backend.domain.member.entity.Member;
 import com.SafeNet.Backend.domain.member.repository.MemberRepository;
+import com.SafeNet.Backend.domain.message.dto.MessageDto;
 import com.SafeNet.Backend.domain.message.entity.Message;
 import com.SafeNet.Backend.domain.message.dto.MessageResponseDto;
 import com.SafeNet.Backend.domain.message.repository.MessageRepository;
+import com.SafeNet.Backend.domain.message.service.MessageService;
 import com.SafeNet.Backend.domain.messageroom.entity.MessageRoom;
 import com.SafeNet.Backend.domain.messageroom.dto.MessageRoomDto;
 import com.SafeNet.Backend.domain.messageroom.repository.MessageRoomRepository;
 import com.SafeNet.Backend.domain.post.entity.Post;
 import com.SafeNet.Backend.domain.post.exception.PostException;
 import com.SafeNet.Backend.domain.post.repository.PostRepository;
 import com.SafeNet.Backend.global.exception.CustomException;
+import com.SafeNet.Backend.global.pubsub.RedisPublisher;
 import com.SafeNet.Backend.global.pubsub.RedisSubscriber;
 import jakarta.annotation.PostConstruct;
 import lombok.extern.slf4j.Slf4j;
@@ -23,6 +26,7 @@
 import org.springframework.http.HttpStatus;
 import org.springframework.stereotype.Service;
 
+import java.time.LocalDateTime;
 import java.util.*;
 
 @Slf4j
@@ -38,6 +42,9 @@ public class MessageRoomService {
 
     private HashOperations<String, String, MessageRoomDto> opsHashMessageRoom;
     private Map<String, ChannelTopic> topics;
+    private final RedisPublisher redisPublisher; // RedisPublisher 주입
+
+    private final MessageService messageService;
 
     private static final String Message_Rooms = "MESSAGE_ROOM"; // Redis에 채팅방 데이터를 저장하기 위한 해시맵의 키
 
@@ -48,14 +55,18 @@ public MessageRoomService(
             MemberRepository memberRepository,
             RedisMessageListenerContainer redisMessageListener,
             RedisSubscriber redisSubscriber,
-            @Qualifier("customRedisTemplate") RedisTemplate<String, Object> redisTemplate) {
+            @Qualifier("customRedisTemplate") RedisTemplate<String, Object> redisTemplate,
+            RedisPublisher redisPublisher,
+            MessageService messageService) {
         this.messageRoomRepository = messageRoomRepository;
         this.messageRepository = messageRepository;
         this.postRepository = postRepository;
         this.memberRepository = memberRepository;
         this.redisMessageListener = redisMessageListener;
         this.redisSubscriber = redisSubscriber;
         this.redisTemplate = redisTemplate;
+        this.redisPublisher = redisPublisher;
+        this.messageService = messageService;
     }
 
     @PostConstruct
@@ -127,15 +138,14 @@ public List<MessageResponseDto> findAllRoomByUser(String email) {
         List<MessageResponseDto> messageRoomDtos = new ArrayList<>();
 
         for (MessageRoom messageRoom : messageRooms) {
-            //  member가 sender인 경우
-            if (member.getName().equals(messageRoom.getSender())) {
+            if (member.getName().equals(messageRoom.getSender())) { //  member가 sender인 경우
                 // 가장 최신 메시지 & 생성 시간 조회
                 // TimeStamped 클래스에서 설정해둔 보낸 시간(sentTime)을 통해 각 채팅방에서 가장 최근 메시지와 그 메시지가 보내진 시간을 꺼낸다.
                 Message latestMessage = messageRepository.findTopByMessageRoom_RoomIdOrderBySentTimeDesc(messageRoom.getRoomId());
                 MessageResponseDto messageRoomDto = setLatestMessage(messageRoom, latestMessage, messageRoom.getReceiver());
                 messageRoomDtos.add(messageRoomDto);
 
-            } else {  // user가 receiver인 경우
+            } else if (member.getName().equals(messageRoom.getReceiver()) || messageRoom.isFirstMessageSent()) {  // user가 receiver이면서 첫번째 메시지가 있는 경우
                 // 가장 최신 메시지 & 생성 시간 조회
                 Message latestMessage = messageRepository.findTopByMessageRoom_RoomIdOrderBySentTimeDesc(messageRoom.getRoomId());
                 MessageResponseDto messageRoomDto = setLatestMessage(messageRoom, latestMessage, messageRoom.getSender());
@@ -267,18 +277,47 @@ private MessageRoomDto convertToDto(MessageRoom messageRoom) {
     }
 
     // 쪽지방 입장
-    public void enterMessageRoom(String roomId) {
+    public void enterMessageRoom(String roomId, String username) {
         ChannelTopic topic = topics.get(roomId);
 
         if (topic == null) {
             topic = new ChannelTopic(roomId);
             redisMessageListener.addMessageListener(redisSubscriber, topic);
             topics.put(roomId, topic);
+
+            // 입장 메시지 추가 (채팅방에 처음 입장하는 경우)
+            MessageDto enterMessage = MessageDto.builder()
+                    .sender(username) // 메시지 보낸 사람으로 username 설정
+                    .roomId(roomId)
+                    .message(username + "이(가) 채팅방에 입장했습니다")
+                    .sentTime(LocalDateTime.now().toString())
+                    .build();
+            redisPublisher.publish(topic, enterMessage);
         }
     }
 
     // redis 채널에서 쪽지방 조회
     public ChannelTopic getTopic(String roomId) {
         return topics.get(roomId);
     }
+
+    public void handleMessage(String roomId, String username, MessageDto messageDto) {
+        // 클라이언트 채팅방(topic) 입장, 대화를 위해 리스너와 연동
+        enterMessageRoom(roomId, username);
+
+        // 메시지 전송 로직
+        MessageDto messageWithTime = MessageDto.builder()
+                .sender(messageDto.getSender())
+                .roomId(messageDto.getRoomId())
+                .message(messageDto.getMessage())
+                .sentTime(LocalDateTime.now().toString())
+                .build();
+
+        // Websocket에 발행된 메시지를 redis로 발행한다(publish)
+        // 해당 쪽지방을 구독(subscribe)한 클라이언트에게 메시지가 실시간 전송
+        redisPublisher.publish(getTopic(messageDto.getRoomId()), messageWithTime);
+
+        // DB와 Redis에 메시지 저장
+        messageService.saveMessage(messageDto);
+    }
 }

src/main/java/com/SafeNet/Backend/global/config/AuthChannelInterceptor.java

@@ -0,0 +1,46 @@
+package com.SafeNet.Backend.global.config;
+
+import com.SafeNet.Backend.global.auth.JwtTokenProvider;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageChannel;
+import org.springframework.messaging.simp.stomp.StompCommand;
+import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
+import org.springframework.messaging.support.ChannelInterceptor;
+import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
+import org.springframework.stereotype.Component;
+
+import java.util.Collections;
+
+// WebSocket 메시지의 헤더에서 ACCESS_TOKEN을 추출하고 검증
+// 유효한 토큰이 있는 경우 사용자 인증 정보를 설정하고, 유효하지 않은 경우 연결을 차단
+@Component
+public class AuthChannelInterceptor implements ChannelInterceptor {
+
+    private final JwtTokenProvider jwtTokenProvider;
+
+    @Autowired
+    public AuthChannelInterceptor(JwtTokenProvider jwtTokenProvider) {
+        this.jwtTokenProvider = jwtTokenProvider;
+    }
+
+    @Override
+    public Message<?> preSend(Message<?> message, MessageChannel channel) {
+        StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
+        if (StompCommand.CONNECT.equals(accessor.getCommand())) { // CONNECT 프레임은 서버에 대한 인증 및 기타 설정과 관련된 정보를 전송하기 위해 사용
+            String token = accessor.getFirstNativeHeader("ACCESS_TOKEN");
+            if (token != null && token.startsWith("Bearer ")) {
+                token = token.substring(7);
+                if (jwtTokenProvider.validateToken(token)) {
+                    String username = jwtTokenProvider.getAuthentication(token).getName();
+                    accessor.setUser(new UsernamePasswordAuthenticationToken(username, null, Collections.emptyList()));
+                } else {
+                    throw new IllegalArgumentException("Invalid or expired token");
+                }
+            } else {
+                throw new IllegalArgumentException("Missing or invalid ACCESS_TOKEN header");
+            }
+        }
+        return message;
+    }
+}

src/main/java/com/SafeNet/Backend/global/config/WebSockConfig.java

@@ -1,6 +1,7 @@
 package com.SafeNet.Backend.global.config;
 
 import org.springframework.context.annotation.Configuration;
+import org.springframework.messaging.simp.config.ChannelRegistration;
 import org.springframework.messaging.simp.config.MessageBrokerRegistry;
 import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
 import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
@@ -9,6 +10,16 @@
 @Configuration
 @EnableWebSocketMessageBroker
 public class WebSockConfig implements WebSocketMessageBrokerConfigurer {
+    private final AuthChannelInterceptor authChannelInterceptor;
+
+    public WebSockConfig(AuthChannelInterceptor authChannelInterceptor) {
+        this.authChannelInterceptor = authChannelInterceptor;
+    }
+
+    @Override
+    public void configureClientInboundChannel(ChannelRegistration registration) {
+        registration.interceptors(authChannelInterceptor); // 웹소켓 검증을 위한 인터셉터 등록
+    }
 
     @Override
     public void configureMessageBroker(MessageBrokerRegistry config) {