Spring

Spring STOMP로 채팅, 알림 구현 + Spring Security, JWT

계양 꿀주먹 2024. 8. 19. 22:24

스터디카페 예약 서비스 프로젝트를 진행하면서 STOMP를 사용해 채팅과 알림 기능을 구현한 내용을 정리한 글입니다.

틀린 부분이 있을 수 있으니 댓글로 의견을 주시면 감사합니다.

 

1. build.gradle 의존성 추가

// websocket
implementation 'org.springframework.boot:spring-boot-starter-websocket'
// STOMP
implementation 'org.webjars:stomp-websocket:2.3.3'

 

2. StompWebSocketConfig 

@EnableWebSocketMessageBroker
@Configuration
@RequiredArgsConstructor
public class StompWebSocketConfig implements WebSocketMessageBrokerConfigurer {

    private final StompHandler stompHandler;

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {

        registry.addEndpoint("/ws")
                .setAllowedOriginPatterns("*")
                .withSockJS();
    }

    @Override
    public void configureMessageBroker(MessageBrokerRegistry registry) {

        registry.setApplicationDestinationPrefixes("/pub");     // 사용자 -> 서버
        registry.enableSimpleBroker("/sub");    // 서버 -> 사용자
    }

    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {
        registration.interceptors(stompHandler);
    }
}

 

Spring 에서 WebSocket과 STOMP를 사용하기 위한 Config 입니다.

  • registerStompEndpoints : 클라이언트가 STOMP를 사용하기 위해 서버에 연결할 수 있는 EndPoint를 등록합니다.
    • addEndpoint("/ws"): 클라이언트가 /ws 엔드포인트를 통해 WebSocket에 연결할 수 있도록 설정합니다.
    • setAllowedOriginPatterns("*"): 도메인에서의 요청을 허용합니다. 이는 CORS 설정을 위해 사용됩니다.
    • withSockJS(): WebSocket을 지원하지 않는 브라우저를 위해 SockJS를 사용할 수 있도록 설정합니다. SockJS는 WebSocket과 유사한 기능을 제공하는 폴백 옵션입니다.
  • configureMessageBroker : 메시지 브로커를 설정하는 메소드이며, 메시지 발행(pub)과 구독(sub) 경로를 설정합니다.
    • setApplicationDestinationPrefixes : 발행 경로를 설정
    • enableSimpleBroker : 구독 경로를 설정
  • configureClientInboundChannel : 클라이언트에게 들어오는 메시지를 처리하기 위한 채널을 설정하는 메소드입니다.
    WebSocket은 기존 HTTP의 Spring Security와 독립적이기 때문에 StompHandler에서 Security 관련 설정을 합니다.

    ** 독립적인 이유 : HTTP는 각 요청마다 SecurityFilterChain에서 HTTP의 요청을 가로채 처리하지만, WebSocket은 핸드셰이크 과정 후 별도의 HTTP 요청이 없기 때문에 HTTP와 같은 인층 체계를 쓸 수 없습니다.

 

3. StompHandler

@Component
@RequiredArgsConstructor
public class StompHandler implements ChannelInterceptor {

    private final JwtUtil jwtUtil;

    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
        String accessToken = "";

        // 연결, 메시지 구독, 메시지 전송 요청에 대해 실행
        if((accessor.getCommand() == StompCommand.CONNECT) || (accessor.getCommand() == StompCommand.SEND)
                || (accessor.getCommand() == StompCommand.SUBSCRIBE)) {

            accessToken = accessor.getFirstNativeHeader("Authorization");

            if (StringUtils.hasText(accessToken) && accessToken.startsWith("Bearer ")) {
                accessToken = accessToken.substring(7);
            }

            try {
                jwtUtil.validateAccessToken(accessToken);
            } catch (ExpiredJwtException e) {
                throw new CatchStudyException(ErrorCode.EXPIRED_ACCESS_TOKEN);
            }

            Authentication authentication = jwtUtil.getAuthentication(accessToken);
            SecurityContextHolder.getContext().setAuthentication(authentication);
            accessor.setUser(authentication);
        }

        return message;
    }
}

 

모든 메시지에 JWT Token을 검증하면 자원 낭비가 많이 일어나게 되어 연결 요청, 구독 요청, 메시지 전송에 대해서만 JWT 검사를 하도록 했습니다.

 

4. Message, ChatRoom, ChatNotification Entity

@Entity
@Getter
@NoArgsConstructor
@Table(name = "message")
public class Message {
    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    @Column(name = "message_id")
    private Long messageId;

    @Column(columnDefinition = "TEXT")
    private String chat;

    @Column(name = "create_date")
    private LocalDateTime createDate;


    @ManyToOne(fetch = FetchType.LAZY)
    @JoinColumn(name = "chat_room_id")
    private ChatRoom chatRoom;

    @ManyToOne(fetch = FetchType.LAZY)
    @JoinColumn(name = "user_id")
    private Users user;

    public Message(MessageRequestDto messageRequestDto, Users user, ChatRoom chatRoom) {
        this.chat = messageRequestDto.getChat();
        this.createDate = LocalDateTime.now();
        this.chatRoom = chatRoom;
        this.user = user;
    }
}
@Entity
@Getter
@NoArgsConstructor
@Table(name = "chat_room")
public class ChatRoom {
    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    @Column(name = "chat_room_id")
    private Long chatRoomId;

    @ManyToOne(fetch = FetchType.LAZY)
    @JoinColumn(name = "user_id")
    private Users user;

    @ManyToOne(fetch = FetchType.LAZY)
    @JoinColumn(name = "cafe_id")
    private StudyCafe studyCafe;

    public ChatRoom(Users user, StudyCafe studyCafe) {
        this.user = user;
        this.studyCafe = studyCafe;
    }
}
@Entity
@Getter
@NoArgsConstructor
@Table(name = "chat_notification")
public class ChatNotification {

    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    @Column
    private long chatNotificationId;

    @ManyToOne(fetch = FetchType.LAZY)
    @JoinColumn(name = "chat_room_id")
    private ChatRoom chatRoom;

    @ManyToOne(fetch = FetchType.LAZY)
    @JoinColumn(name = "user_id")
    private Users user;

    @Column
    private boolean status;

    public ChatNotification(ChatRoom chatRoom, Users user) {
        this.chatRoom = chatRoom;
        this.user = user;
        this.status = false;
    }

    public void readNotification() {
        this.status = true;
    }
}

 

5. MessageRequestDto, MessageResponseDto

@Getter
@JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
@AllArgsConstructor
@NoArgsConstructor
public class MessageRequestDto {

    private long chatRoomId;
    private String chat;
}
@Getter
@JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
public class MessageResponseDto {

    private final long userId;
    private final long messageId;
    private final String chat;
    private final LocalDateTime createDate;

    public MessageResponseDto(Message message) {
        this.userId = message.getUser().getUserId();
        this.messageId = message.getMessageId();
        this.chat = message.getChat();
        this.createDate = message.getCreateDate();
    }
}

 

6. MessageConroller

@RestController
@RequiredArgsConstructor
public class MessageController {

    private final ChatService chatService;

    @MessageMapping("/{chatRoomId}/chat")
    @SendTo("/sub/{chatRoomId}/chat")
    public MessageResponseDto createMessage(@DestinationVariable long chatRoomId, MessageRequestDto messageRequestDto,
                              @Header("Authorization") String token) {
       return chatService.createMessage(chatRoomId, messageRequestDto, token);
    }
}

 

STOMP를 연결 후 채팅 메시지를 처리하는 Controller 입니다. STOMP의 명령 COMMAND가 SEND일 때 실행됩니다.

  • @MessageMapping : /pub/1/chat 으로 요청이 들어오면 메소드가 호출되며, @DestinationVariable을 통해 chatRoomId가 1로 매핑됩니다.
  • @SendTo : /sub/1/chat 을 구독한 모든 클라이언트들에게 메소드의 처리 결과가 반환됩니다.
  • @Header : 클라이언트가 보낸 메시지에 포함된 JWT Token을 가져옵니다.

7. ChatService 

@Service
@Transactional
@RequiredArgsConstructor
public class ChatService {

private final ChatRoomRepository chatRoomRepository;
    private final MessageRepository messageRepository;
    private final UsersRepository usersRepository;
    private final StudyCafeRepository studyCafeRepository;
    private final UsersService usersService;
    private final ChatNotificationRepository chatNotificationRepository;
    private final JwtUtil jwtUtil;
    private final Map<Long, Map<String, String>> chatRoomMap = new ConcurrentHashMap<>(); // chatRoomId, <session id, 접속한 유저 email>
    private final Map<String, Long> sessionToChatRoom = new ConcurrentHashMap<>();
    
    @EventListener
    public void handleSessionConnect(SessionConnectEvent event) {
        StompHeaderAccessor accessor = StompHeaderAccessor.wrap(event.getMessage());

        long chatRoomId = Long.parseLong((String) ((List) accessor.getNativeHeader("chatRoomId")).get(0));
        String sessionId = accessor.getSessionId();
        String accessToken = ((String) ((List) accessor.getNativeHeader("Authorization")).get(0));
        String email = jwtUtil.getEmailFromJwtToken(accessToken);

        Map<String, String> userMap = chatRoomMap.getOrDefault(chatRoomId, new HashMap<>());
        userMap.put(sessionId, email);

        chatRoomMap.put(chatRoomId, userMap);
        sessionToChatRoom.put(sessionId, chatRoomId);
    }

    @EventListener
    public void handleSessionDisconnect(SessionDisconnectEvent event) {
        StompHeaderAccessor accessor = StompHeaderAccessor.wrap(event.getMessage());
        String sessionId = accessor.getSessionId();
        Long chatRoomId = sessionToChatRoom.get(sessionId);

        Map<String, String> userMap = chatRoomMap.get(chatRoomId);
        userMap.remove(sessionId);
        sessionToChatRoom.remove(sessionId);

        if (userMap.isEmpty()) chatRoomMap.remove(chatRoomId);
        else chatRoomMap.put(chatRoomId, userMap);
    }
    
     @Transactional
    public MessageResponseDto createMessage(long chatRoomId, MessageRequestDto messageRequestDto, String token) {

        String accessToken = token.substring(7);
        String email = jwtUtil.getEmailFromJwtToken(accessToken);
        Users user = usersRepository.findByEmail(email);
        ChatRoom chatRoom = chatRoomRepository.findById(chatRoomId).
                orElseThrow(() -> new CatchStudyException(ErrorCode.CHATROOM_NOT_FOUND));

        int userCount = chatRoomMap.get(chatRoomId).size(); // 참여중인 채탕방의 유저 수 1 or 2

        if(userCount == 1) {
            if(user.getAuthor().equals(Author.roleUser)) {
                Users studyCafeUser = usersRepository.findByUserId(chatRoom.getStudyCafe().getUser().getUserId()).
                        orElseThrow(() -> new CatchStudyException(ErrorCode.USER_NOT_FOUND));
                chatNotificationRepository.save(new ChatNotification(chatRoom, studyCafeUser));
            } else {
                Users client =  chatRoom.getUser();
                chatNotificationRepository.save(new ChatNotification(chatRoom, client));
            }
        }

        return new MessageResponseDto(messageRepository.save(new Message(messageRequestDto, user, chatRoom)));
    }
  • chatRoomMap : 채팅방에 접속한 유저를 알기 위한 Map 입니다.
    chatRoomId를 key로 <sessionId, email>을 가진 userMap을 value로 갖습니다.
  • userMap : sessionId를 key로 유저의 email을 저장한 Map입니다.
  • sessionToChatRoom : sessionId를 통해 chatRoom을 찾기 위한 Map입니다. 
  • handleSessionConnect : STOMP가 연결되었을 때 실행되는 메소드입니다. 저희 서비스에서는 채팅방에 입장할 때, 연결합니다.
    연결할 때 JWT Token에서 email을 가져오고 chatRoomId를 key값으로 <sessionId, email> 형태의 Map을 담아 저장해 관리하며 연결이 해제됐을 때, chatRoomId를 쉽게 찾을 수 있도록 sessionToChatRoom Map에 값을 저장합니다.
  • handleSessionDisConnect : STOMP가 연결 해제를 감지했을 때 실행되는 메소드입니다.
    연결이 해제될 때는 header에 값이 없으므로, 해제될 때의 sessionId값을 통해 sesstionToChatRoom을 통해 쉽게 chatRoomId를 가져와서 채팅방의 유저 목록을 불러와 sessionId 값으로 목록에서 제거 후 유저 목록이 비어있는 경우 채팅방을 제거하고, 남아있다면 다시 값을 저장해 업데이트합니다.
  • createMessage : 채팅이 전송되었을 때 실행되는 메소드입니다.
    스터디카페 사장님과 사용자의 1 : 1 대화이기 때문에 채팅이 전송되었을 때는 현재 참여중인 사용자의 수가 1명 혹은 2명입니다.
    chatRoomId를 통해 채팅방에 접속 중인 사용자의 수가 2명이라면 알림을 생성할 필요가 없으니 1명일 때만 알림을 저장하여 채팅방 목록을 확인했을 때 알 수 있도록 합니다.

 

 


참고 

https://docs.spring.io/spring-framework/reference/web/websocket/stomp.html

 

전체코드

https://github.com/fourix4/Back-End