STOMP, JWT 인증을 사용한 채팅 구현

우선 일단 RDB 에다가 저장하는 채팅부터 구현 후, 리팩토링을 통해서 Redis, Message Queue를 사용하여 개발할 예정이다. 추후 진행할 때, 왜 이렇게 구현해야하는지 작성할 예정이고 지금은 Spring Boot 코드를 작성할 것이다. 

 

나는 React + Next.js 가 프론트이기 때문에 프론트 코드는 작성하지 않을 예정이다. 인터넷을 보면 많이 있으니 해당 코드를 보면서 참고하여 구현하면 되겠다. 

 

 

😜 디렉토리 구조

 

message
│
├── controller
│   ├── request
│   │   ├── CreateChatMessageRequest
│   │   └── CreateChatRoomRequest
│   ├── response
│   │   ├── CreateChatRoomResponse
│   │   ├── FindChatRoomResponse
│   │   ├── ChatRoomController
│   │   └── MessageController
│
├── domain
│   ├── chat
│   └── message
│       └── Message
│
├── dto
│   ├── ChatRoomDetailDto
│   ├── ChatRoomListDto
│   └── CreateFindChatRoom
│
├── entity
│   ├── ChatRoomEntity
│   └── MessageEntity
│
├── repository
│   ├── chat
│   │   ├── ChatRoomRepository
│   │   ├── CustomChatRepository
│   │   └── CustomChatRepositoryImpl
│   └── message
│       └── MessageRepository
│
└── service
    ├── chat
    │   ├── ChatRoomCommandServiceImpl
    │   ├── ChatRoomCommandService
    │   ├── ChatRoomQueryService
    │   └── ChatRoomQueryServiceImpl
    └── message
        ├── MessageCommandService
        └── MessageCommandServiceImpl

 

 

build.gradle

 

    // stomp 설정 추가
    implementation 'org.springframework.boot:spring-boot-starter-websocket'

 

 

 

🤗 ChatRoomEntity

 

package com.realworld.feature.message.entity;


import com.realworld.feature.message.domain.chat.CreateChatRoom;
import com.realworld.feature.product.entity.ProductJpaEntity;
import jakarta.persistence.*;
import lombok.*;
import org.springframework.data.annotation.CreatedDate;
import org.springframework.data.jpa.domain.support.AuditingEntityListener;
import java.time.LocalDateTime;
import java.util.UUID;


@ToString
@Builder
@Entity
@Getter
@Table(name="chat_rooms")
@EntityListeners(AuditingEntityListener.class)
@AllArgsConstructor
@NoArgsConstructor
public class ChatRoomEntity {

    @Id
    @Column(name="chat_room_id")
    private UUID roomId;

    @CreatedDate
    @Column(name = "create_at")
    private LocalDateTime createdAt;

    @ManyToOne(fetch = FetchType.LAZY)
    @JoinColumn(name = "product_seq")
    private ProductJpaEntity product;

    @Column(name = "room_maker_id")
    private String roomMaker;

    @Column(name = "guest_id")
    private String guest;


    public CreateChatRoom createToDomain(){
        return CreateChatRoom.builder()
                .roomId(this.roomId.toString())
                .lastMessage(this.roomId.toString())
                .roomMaker(this.roomMaker)
                .guest(this.guest)
                .productSeq(this.product.getProductSeq())
                .createdAt(String.valueOf(this.createdAt))
                .build();
    }


}

 

나의 채팅의 경우 상품에 물려서 Chatting방이 생성되어 유저간 통신이 되어야 하기 때문에 ManyToOne으로 단방향 매핑을 맺어주었다. 또한 ProductEntity에서는 해당 데이터를 가지고 있을 필요가 없기 때문에 컬럼을 저장하지 않았다. 

이런 방식은 JPA를 어느정도 공부하거나 프로젝트를 진행하면 왜 이런 방식으로 진행하는지 알 수 있다. 

 

 

🧐 MessageEntity

 

package com.realworld.feature.message.entity;

import com.realworld.feature.message.domain.message.Message;
import jakarta.persistence.*;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Getter;
import lombok.NoArgsConstructor;
import org.springframework.data.annotation.CreatedDate;
import org.springframework.data.jpa.domain.support.AuditingEntityListener;
import java.time.LocalDateTime;

@Builder
@Entity
@Getter
@AllArgsConstructor
@NoArgsConstructor
@Table(name="messages")
@EntityListeners(AuditingEntityListener.class)
public class MessageEntity {

    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    @Column(name = "message_id")
    private Long messageId;

    @JoinColumn(name = "chat_room_id", insertable = false, updatable = false) // 단순히 값만 필요하기에 해당 처리
    private String roomId;

    @JoinColumn(name="user_id", insertable = false, updatable = false)
    private String writer;


    private String message;

    @CreatedDate
    @Column(name="create_at")
    private LocalDateTime createAt;


    public Message toDomain(MessageEntity entity) {
        return Message.builder()
                .roomId(this.roomId)
                .message(this.message)
                .writer(this.writer)
                .createAt(this.createAt)
                .build();
    }
}

 

해당 코드는 MessageEntity 코드이다. 

메시지를 단순히 RDB에 저장하고, 이후에 채팅방을 불러올 때, 데이터를 가져와주기 위하여 만들어주었다. 

 

 

😎 ChatRoomController

 

package com.realworld.feature.message.controller;

import com.realworld.feature.message.controller.request.CreateChatRoomRequest;
import com.realworld.feature.message.controller.resonse.CreateChatRoomResponse;
import com.realworld.feature.message.controller.resonse.FindChatRoomResponse;
import com.realworld.feature.message.dto.ChatRoomListDto;
import com.realworld.feature.message.service.chat.ChatRoomCommandService;
import com.realworld.feature.message.service.chat.ChatRoomQueryService;
import com.realworld.global.code.SuccessCode;
import com.realworld.global.response.ApiResponse;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.ResponseEntity;
import org.springframework.security.core.annotation.AuthenticationPrincipal;
import org.springframework.security.core.userdetails.User;
import org.springframework.web.bind.annotation.*;

import java.util.List;

@Slf4j
@RestController
@RequiredArgsConstructor
@RequestMapping("/api/v1/chat")
public class ChatRoomController {

    private final ChatRoomCommandService chatRoomCommandService;
    private final ChatRoomQueryService chatRoomQueryService;

    @PostMapping("/personal")
    public ResponseEntity<ApiResponse<CreateChatRoomResponse>> createChatRoomForPersonal(@AuthenticationPrincipal User user, @RequestBody CreateChatRoomRequest request) {

        // 유저의 정보를 이용해 1:1 채팅방을 생성하는 서비스 호출
        CreateChatRoomResponse response =chatRoomCommandService.createChatRoomForPersonal(user.getUsername(), request);

        ApiResponse<CreateChatRoomResponse> apiResponse = null;


        // 응답 데이터와 성공 상태, 메시지를 포함한 ApiResponse 객체 생성
        if(response.isExist()) apiResponse = new ApiResponse<>(response, SuccessCode.SELECT_SUCCESS.getStatus(), SuccessCode.SELECT_SUCCESS.getMessage());
        else apiResponse = new ApiResponse<>(response, SuccessCode.INSERT_SUCCESS.getStatus(), SuccessCode.INSERT_SUCCESS.getMessage());

        // HTTP 200 OK 응답과 함께 API 응답 데이터 반환
        return ResponseEntity.ok(apiResponse);
    }

    @GetMapping("")
    public ResponseEntity<ApiResponse<FindChatRoomResponse>> findChatRoomList(@AuthenticationPrincipal User user){

        List<ChatRoomListDto> list = chatRoomQueryService.findByChatRoomList(user.getUsername());

        FindChatRoomResponse response  = FindChatRoomResponse.builder()
                .list(list)
                .build();

        ApiResponse<FindChatRoomResponse> apiResponse = new ApiResponse<>(response, SuccessCode.SELECT_SUCCESS.getStatus(), SuccessCode.SELECT_SUCCESS.getMessage());


        return ResponseEntity.ok(apiResponse);
    }
}

 

Dto를 통해서 QueryDSL로 받기 때문에 Dto로 받아서 처리해주었다. 

/personal 컨트롤러 경우에 각각의 경우에 응답 결과가 다르게 처리되어야하기 때문에 분기문을 통해서 응답 값을 보내주었다. 

 

 

 

😎 ChatRoomCommandServiceImpl

 

package com.realworld.feature.message.service.chat;

import com.realworld.feature.member.repository.MemberRepository;
import com.realworld.feature.member.service.MemberQueryService;
import com.realworld.feature.message.controller.request.CreateChatRoomRequest;
import com.realworld.feature.message.controller.resonse.CreateChatRoomResponse;
import com.realworld.feature.message.domain.chat.CreateChatRoom;
import com.realworld.feature.message.dto.ChatRoomDetailDto;
import com.realworld.feature.message.entity.ChatRoomEntity;
import com.realworld.feature.message.repository.chat.ChatRoomRepository;
import com.realworld.feature.product.entity.ProductJpaEntity;
import com.realworld.feature.product.repository.ProductRepository;
import com.realworld.global.code.ErrorCode;
import com.realworld.global.config.exception.CustomChatExceptionHandler;
import com.realworld.global.config.exception.CustomProductExceptionHandler;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.util.Optional;

import static com.realworld.feature.message.controller.resonse.CreateChatRoomResponse.convertToResponse;
import static com.realworld.feature.message.domain.chat.CreateChatRoom.initialCreateChatRoom;
import static com.realworld.feature.message.domain.chat.FindChatRoom.toCreateModel;

@Slf4j
@Service
@RequiredArgsConstructor
public class ChatRoomCommandCommandServiceImpl implements ChatRoomCommandService {

    private final MemberQueryService memberQueryService;
    private final ChatRoomRepository chatRoomRepository;
    private final ProductRepository productRepository;
    private final ChatRoomQueryService chatRoomQueryService;
    private final MemberRepository memberRepository;

    @Transactional
    @Override
    public  CreateChatRoomResponse createChatRoomForPersonal(String userId, CreateChatRoomRequest request) {

        // 요청한 유저가 방 생성자인지 확인
        if(!userId.equals(request.getRoomMakerId()))
            throw new CustomChatExceptionHandler(ErrorCode.CHAT_USER_NOT_FOUND);

        // 채팅방이 이미 존재하는지 확인
        Optional<ChatRoomDetailDto> chatRoom = chatRoomQueryService.findByChatRoom(toCreateModel(request));


        if(chatRoom.isPresent()){
            return convertToResponse(chatRoom.get(), true);
        }

        // 제품 실제로 존재하는지 확인
        ProductJpaEntity entity = productRepository.findById(request.getProductSeq()).orElseThrow(() -> new CustomProductExceptionHandler(ErrorCode.NOT_EXISTS_PRODUCT));

        // 제품을 생성한 사용자와 RoomMaker와 동일하면 에러
        if(entity.getMember().getUserId().equals(request.getRoomMakerId())) throw new CustomChatExceptionHandler(ErrorCode.CHAT_DUPLICATE_USER);

        // 채팅방 저장
        CreateChatRoom room = initialCreateChatRoom(request.getRoomMakerId(), request.getGuestId());

        room.setProduct(entity);
        ChatRoomEntity roomSave = chatRoomRepository.save(room.toEntity());

        // 채팅방 응답 반환
        return convertToResponse(roomSave, false);
    }


}

 

ChatRoomService를 interface 로 구현 후, 만들어주어야한다. 

 

 

if(!userId.equals(request.getRoomMakerId()))
    throw new CustomChatExceptionHandler(ErrorCode.CHAT_USER_NOT_FOUND);

 

요청한 유저가 방 생성자인지 즉, 로그인한 유저가 방생성자인지 확인 후, 아니라면 에러를 뱉어주는 코드이다. 

 

// 채팅방이 이미 존재하는지 확인
Optional<ChatRoomDetailDto> chatRoom = chatRoomQueryService.findByChatRoom(toCreateModel(request));

채팅방 존재 여부를 확인해주는 코드이다. 아래에서 findByChatRoom 코드를 작성 해 줄 예정이다.

 

if(chatRoom.isPresent()){
    return convertToResponse(chatRoom.get(), true);
}

 

존재하면 바로 Response로 Controller에 던져준다. 

 

// 제품 실제로 존재하는지 확인
        ProductJpaEntity entity = productRepository.findById(request.getProductSeq()).orElseThrow(() -> new CustomProductExceptionHandler(ErrorCode.NOT_EXISTS_PRODUCT));

 

제품 실제로 존재하는지 확인해주는 코드이다. 이를 통해서 실제 존재 여부를 판단해준다. 

 

 // 제품을 생성한 사용자와 RoomMaker와 동일하면 에러
        if(entity.getMember().getUserId().equals(request.getRoomMakerId())) throw new CustomChatExceptionHandler(ErrorCode.CHAT_DUPLICATE_USER);

 

제품 생성한 사용자가 자신과 채팅할 수 없게 하기 위한 코드이다. 

 

위에 CustomExcpetion코드는 ControllerAdvice를 통해서 구현해주면 된다. 

 

 

 

아래는 레포지토리 구조이다. 

 

📄 ChatRoomRepository

 

public interface ChatRoomRepository extends JpaRepository<ChatRoomEntity, UUID>, CustomChatRepository {

}

 

 

🚴🏻‍♂️ ChatRoomQueryServiceImpl 

 

package com.realworld.feature.message.service.chat;

import com.realworld.feature.message.domain.chat.FindChatRoom;
import com.realworld.feature.message.dto.ChatRoomDetailDto;
import com.realworld.feature.message.dto.ChatRoomListDto;
import com.realworld.feature.message.entity.ChatRoomEntity;
import com.realworld.feature.message.repository.chat.ChatRoomRepository;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Optional;

@Slf4j
@Service
@RequiredArgsConstructor
public class ChatRoomQueryServiceImpl implements ChatRoomQueryService{

    private final ChatRoomRepository repository;


    @Override
    public Optional<ChatRoomDetailDto> findByChatRoom(FindChatRoom room) {
        return repository.findChatRoom(room);
    }

    @Override
    public List<ChatRoomListDto> findByChatRoomList(String userId) {

        return repository.findByChatRoomList(userId, userId);
    }
}

 

 

 

 

📖 CustomChatRepositoryImpl

 

package com.realworld.feature.message.repository.chat;

import com.querydsl.core.BooleanBuilder;
import com.querydsl.core.types.Predicate;
import com.querydsl.core.types.Projections;
import com.querydsl.core.types.dsl.BooleanExpression;
import com.querydsl.jpa.impl.JPAQueryFactory;
import com.realworld.feature.member.entity.QMemberJpaEntity;
import com.realworld.feature.message.domain.chat.FindChatRoom;
import com.realworld.feature.message.dto.ChatRoomDetailDto;
import com.realworld.feature.message.dto.ChatRoomListDto;
import com.realworld.feature.message.entity.QChatRoomEntity;
import lombok.RequiredArgsConstructor;

import java.util.List;
import java.util.Optional;


@RequiredArgsConstructor
public class CustomChatRepositoryImpl implements CustomChatRepository{

    private final JPAQueryFactory queryFactory;

    private final QChatRoomEntity qChatRoom = QChatRoomEntity.chatRoomEntity;
    private final QMemberJpaEntity qMemberJpa = QMemberJpaEntity.memberJpaEntity;

    @Override
    public Optional<ChatRoomDetailDto> findChatRoom(FindChatRoom findChatRoom) {

        return Optional.ofNullable(queryFactory
                .select(Projections.constructor(ChatRoomDetailDto.class, qChatRoom.roomId, qChatRoom.product.price, qChatRoom.product.title, qChatRoom.product.thumbnailUrl,  qChatRoom.product.productSeq, qMemberJpa.userId.as("userId"), qMemberJpa.profileImage.as("profileImage")))
                .from(qChatRoom)
                .join(qMemberJpa).on(qChatRoom.guest.eq(qMemberJpa.userId))
                .where(
                        qChatRoom.product.productSeq.eq(findChatRoom.getProductSeq())
                                .and(qChatRoom.guest.eq(findChatRoom.getGuest()))
                )
                .fetchOne()
        );
    }


    @Override
    public List<ChatRoomListDto> findByChatRoomList(String roomMaker, String guest) {

        return queryFactory.selectDistinct(
                        Projections.constructor(ChatRoomListDto.class, qChatRoom.roomId, qChatRoom.product.title)
                )
                .from(qChatRoom)
                .where(findChatRoomListCriteria(roomMaker, guest))
                .fetch();

    }


    private Predicate findChatRoomListCriteria(String roomMaker, String guest) {
        BooleanBuilder builder = new BooleanBuilder();
        BooleanExpression roomMakerCondition = qChatRoom.roomMaker.eq(roomMaker);
        BooleanExpression guestCondition = qChatRoom.guest.eq(guest);
        builder.or(roomMakerCondition).or(guestCondition);

        return builder;
    }


}

 

 

아래에서 해당 코드를 설명하겠다. 

🛞 findChatRoom 메서드

 

@Override
public Optional<ChatRoomDetailDto> findChatRoom(FindChatRoom findChatRoom) {
    return Optional.ofNullable(queryFactory
            .select(Projections.constructor(ChatRoomDetailDto.class, qChatRoom.roomId, qChatRoom.product.price, qChatRoom.product.title, qChatRoom.product.thumbnailUrl, qChatRoom.product.productSeq, qMemberJpa.userId.as("userId"), qMemberJpa.profileImage.as("profileImage")))
            .from(qChatRoom)
            .join(qMemberJpa).on(qChatRoom.guest.eq(qMemberJpa.userId))
            .where(
                    qChatRoom.product.productSeq.eq(findChatRoom.getProductSeq())
                            .and(qChatRoom.guest.eq(findChatRoom.getGuest()))
            )
            .fetchOne()
    );
}

 

findChatRoom : 채팅방을 찾기 위한 조건을 담고 있는 객체입니다
Optional<ChatRoomDetailDto> : 채팅방의 상세 정보 DTO를 담고 있는 Optional 객체

select : select 절에서 ChatRoomDetailDto 객체를 생성하기 위해 필요한 필드들을 지정
from : from 절에서 qChatRoom (채팅방 테이블)과 qMemberJpa (멤버 테이블)을 조인
join : join 절에서 Member와 조인한다. 

where : productSeq, guest 기준으로 필터링한다. 

이렇게 해준 이유는 유저에서 ImageUrl을 프로필을 가져와주어야한다. 

여기서 내가 힘들었던 부분은 원래는 공통 테이블에 프로필 이미지가 있었는데 진짜 JPA는 조인을 하고 양방향 매핑을 많이 할수록 힘든 구조이기 때문에 불필요한 조인이아닌, 유저 테이블 자체에 사실상 프로필 이미지가 있는게 구조적으로 맞기에 해당 방식으로 변경해주었다. 

 

⚙️ findByChatRoomList(String roomMaker, String guest) 메서드

 

@Override
public List<ChatRoomListDto> findByChatRoomList(String roomMaker, String guest) {
    return queryFactory.selectDistinct(
                    Projections.constructor(ChatRoomListDto.class, qChatRoom.roomId, qChatRoom.product.title)
            )
            .from(qChatRoom)
            .where(findChatRoomListCriteria(roomMaker, guest))
            .fetch();
}

 

 

roomMaker: 채팅방을 생성한 사용자의 ID
guest: 채팅방에 참여한 사용자의 ID

List<ChatRoomListDto> : 채팅방 목록을 담고 있는 DTO 리스트입니다.

select Distinct 절에서 중복되지 않은 채팅방 정보를 조회하기 위해 ChatRoomListDto 객체를 생성합니다.
from 절에서 qChatRoom 테이블을 지정합니다.
where 절에서 findChatRoomListCriteria 메서드를 호출하여 필터 조건을 지정합니다.

 

😆 findChatRoomListCriteria(String roomMaker, String guest) 메서드

private Predicate findChatRoomListCriteria(String roomMaker, String guest) {
    BooleanBuilder builder = new BooleanBuilder();
    BooleanExpression roomMakerCondition = qChatRoom.roomMaker.eq(roomMaker);
    BooleanExpression guestCondition = qChatRoom.guest.eq(guest);
    builder.or(roomMakerCondition).or(guestCondition);

    return builder;
}

 

- BooleanBuilder를 사용하여 필터 조건을 동적으로 생성합니다.
- or 연산을 사용하여 roomMaker 또는 guest 조건 중 하나라도 일치하면 해당 채팅방을 조회하도록 조건을 설정합니다.

roomMaker 또는 guest와 일치하는 채팅방 목록을 반환
왜냐면 둘 중하나라도 있을 때, 채팅방을 가져와야하기 때문이다. 
Product  포토카드를 생성한 사람이라면 guest로 채팅방이 존재 할 것이기에 해당 방식으로 가져와주어야 양쪽의 채팅방 다 가져와 줄 수 있다. 

마지막은 STOMP 관련 설정이다.  Controller, Config, Handler 까지 3개를 만들었다. 나는 여기서 JWT 토큰을 가지고 인증된 사용자만 ws를 통해서 소켓을 열게 만들었다. 

 

 

👻 MessageController

 

 

 

package com.realworld.feature.message.controller;

import com.realworld.feature.message.controller.request.CreateChatMessageRequest;
import com.realworld.feature.message.service.message.MessageCommandService;
import com.realworld.global.code.SuccessCode;
import com.realworld.global.response.ApiResponse;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.ResponseEntity;
import org.springframework.messaging.handler.annotation.DestinationVariable;
import org.springframework.messaging.handler.annotation.MessageMapping;
import org.springframework.messaging.simp.SimpMessagingTemplate;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

@Slf4j
@RestController
@RequestMapping("/api/v1")
@RequiredArgsConstructor
public class MessageController {

    private final MessageCommandService messageCommandService;
    private final SimpMessagingTemplate messagingTemplate;

    @MessageMapping(value = "/{roomId}/messages")
    public ResponseEntity<ApiResponse<?>> enterMessage(@DestinationVariable("roomId") String roomId, @RequestBody CreateChatMessageRequest request) {
        log.info("enter :: {}", request.getRoomId());

        messageCommandService.saveMessage(request);

        messagingTemplate.convertAndSend("/sub/"+ roomId, request);

        ApiResponse<?> apiResponse = new ApiResponse<>(null, SuccessCode.INSERT_SUCCESS.getStatus(), SuccessCode.INSERT_SUCCESS.getMessage());

        return ResponseEntity.ok(apiResponse);
    }


}

 

 

 

 

MessageCommandService: 메시지를 저장하는 서비스 클래스.

SimpMessagingTemplate: 메시지를 특정 목적지로 보내기 위한 Spring의 메시징 템플릿.

@MessageMapping: STOMP 메시지의 경로를 매핑하는 데 사용됩니다. /app/{roomId}/messages 경로로 들어오는 웹소켓 메시지가 이 메서드로 라우팅됩니다.
여기서 {roomId}는 채팅방 ID를 나타내며, 이 경로 변수는 @DestinationVariable("roomId")로 추출됩니다.

메시지 저장: messageCommandService.saveMessage(request)를 호출하여 메시지를 데이터베이스에 저장합니다.
메시지 전송: messagingTemplate.convertAndSend("/sub/"+ roomId, request)를 사용해 해당 채팅방에 구독된 클라이언트에게 메시지를 전송합니다. 여기서 /sub/{roomId}는 구독 경로입니다.
응답 생성: ApiResponse<?> apiResponse = new ApiResponse<>(null, SuccessCode.INSERT_SUCCESS.getStatus(), SuccessCode.INSERT_SUCCESS.getMessage())를 통해 성공 응답을 생성합니다.
응답 반환: ResponseEntity.ok(apiResponse)를 통해 HTTP 200 OK 상태로 응답을 반환합니다.

 

 

 

 

 

 

📪 WebSocketConfig 

 

package com.realworld.global.config;

import com.realworld.global.config.handler.StompHandler;
import lombok.RequiredArgsConstructor;
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;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;

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

    private final StompHandler stompHandler;

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/ws")
                .setAllowedOriginPatterns("*")
                .withSockJS();

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

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

    @Override
    public void configureMessageBroker(MessageBrokerRegistry registry) {
        // 메모리 기반 메시지 브로커가 해당 api를 구독하고 있는 클라이언트에게 메시지를 전달한다.
        // to subscriber
        registry.setApplicationDestinationPrefixes("/pub");

        // 클라이언트부터 메시지를 받을 api의 prefix를 설정한다.
        // publish
        registry.enableSimpleBroker("/sub");
    }
}

 

아래에서 해당 코드를 설명하겠습니다. 

 

 

@Override
public void registerStompEndpoints(StompEndpointRegistry registry) {
    registry.addEndpoint("/ws")
            .setAllowedOriginPatterns("*")
            .withSockJS();

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

registry.addEndpoint("/ws"): /ws 엔드포인트를 통해 WebSocket 연결을 수립할 수 있습니다.
setAllowedOriginPatterns("*"): 모든 도메인에서 오는 요청을 허용합니다.
withSockJS(): SockJS를 지원하도록 설정합니다. SockJS는 WebSocket을 지원하지 않는 브라우저에서 폴백 메커니즘을 제공하는 라이브러리입니다.

 

 

@Override
public void configureMessageBroker(MessageBrokerRegistry registry) {
    registry.setApplicationDestinationPrefixes("/pub");

    registry.enableSimpleBroker("/sub");
}

 

setApplicationDestinationPrefixes("/pub"): 클라이언트가 서버로 메시지를 보낼 때 사용하는 경로의 접두사를 설정합니다.

클라이언트가 메시지를 /pub/hello로 보낸다면, 서버에서는 해당 메시지를 처리할 수 있는 메서드가 필요합니다.

enableSimpleBroker("/sub"): 메모리 기반의 간단한 메시지 브로커를 활성화하고, 해당 브로커가 구독자들에게 메시지를 전달할 때 사용하는 접두사를 설정합니다. 클라이언트는 /sub로 시작하는 경로를 구독하여 서버에서 보낸 메시지를 받을 수 있습니다.

 

 

 

🙌🏻  StompHandler

 

package com.realworld.global.config.handler;

import com.realworld.global.code.ErrorCode;
import com.realworld.global.config.exception.CustomJwtExceptionHandler;
import com.realworld.global.config.jwt.JwtTokenProvider;
import io.jsonwebtoken.ExpiredJwtException;
import io.jsonwebtoken.MalformedJwtException;
import io.jsonwebtoken.UnsupportedJwtException;
import io.jsonwebtoken.security.SecurityException;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpHeaders;
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.core.Authentication;
import org.springframework.security.core.userdetails.User;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

@Slf4j
@Component
@RequiredArgsConstructor
public class StompHandler implements ChannelInterceptor {

    private final JwtTokenProvider jwtTokenProvider;

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

        String id = null;
        if(accessor.getCommand() == StompCommand.CONNECT) {
            String token = resolveToken(accessor);

            try {
                if(token == null) throw new CustomJwtExceptionHandler(ErrorCode.JWT_UNKNOWN_ERROR);

                if (jwtTokenProvider.validateToken(token)) {
                    Authentication authentication = jwtTokenProvider.getAuthentication(token);
                    User user = (User) authentication.getPrincipal();
                    id = user.getUsername();
                }
            } catch (SecurityException | MalformedJwtException | IllegalArgumentException e) {
                throw new CustomJwtExceptionHandler(ErrorCode.JWT_WRONG_TYPE_TOKEN_ERROR);
            } catch (ExpiredJwtException e){
                throw new CustomJwtExceptionHandler(ErrorCode.JWT_TOKEN_EXPIRED_ERROR);
            } catch(UnsupportedJwtException e){
                throw new CustomJwtExceptionHandler(ErrorCode.UNSUPPORTED_TOKEN_ERROR);
            } catch (Exception e){
                throw new CustomJwtExceptionHandler(ErrorCode.JWT_UNKNOWN_ERROR);
            }

            accessor.addNativeHeader("senderUserId", id);
        }

        return message;
    }

    private String resolveToken(StompHeaderAccessor accessor) {
        String bearerToken = accessor.getFirstNativeHeader(HttpHeaders.AUTHORIZATION);

        if(StringUtils.hasText(bearerToken) && bearerToken.startsWith("Bearer ")){
            return bearerToken.substring(7);
        }
        return null;
    }
}

 

 

ChannelInterceptor 인터페이스: 이 인터페이스는 Spring Messaging에서 채널을 가로채는 역할을 합니다. 이 인터페이스를 구현하면 메시지가 채널로 들어오거나 나가기 전에 이를 가로채어 특정 로직을 추가할 수 있습니다.

 

 

🚍 preSend

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

    String id = null;
    if(accessor.getCommand() == StompCommand.CONNECT) {
        String token = resolveToken(accessor);

        try {
            if(token == null) throw new CustomJwtExceptionHandler(ErrorCode.JWT_UNKNOWN_ERROR);

            if (jwtTokenProvider.validateToken(token)) {
                Authentication authentication = jwtTokenProvider.getAuthentication(token);
                User user = (User) authentication.getPrincipal();
                id = user.getUsername();
            }
        } catch (SecurityException | MalformedJwtException | IllegalArgumentException e) {
            throw new CustomJwtExceptionHandler(ErrorCode.JWT_WRONG_TYPE_TOKEN_ERROR);
        } catch (ExpiredJwtException e){
            throw new CustomJwtExceptionHandler(ErrorCode.JWT_TOKEN_EXPIRED_ERROR);
        } catch(UnsupportedJwtException e){
            throw new CustomJwtExceptionHandler(ErrorCode.UNSUPPORTED_TOKEN_ERROR);
        } catch (Exception e){
            throw new CustomJwtExceptionHandler(ErrorCode.JWT_UNKNOWN_ERROR);
        }

        accessor.addNativeHeader("senderUserId", id);
    }

    return message;
}

 

 

StompHeaderAccessor: STOMP 메시지의 헤더를 다루기 위한 유틸리티 클래스입니다. 이 클래스를 사용하여 STOMP 명령 및 헤더를 쉽게 접근할 수 있습니다.

StompCommand.CONNECT: 클라이언트가 WebSocket 연결을 시도할 때 사용하는 STOMP 명령입니다. 이 명령을 가로채어 연결 시도 시의 메시지를 처리합니다.

토큰 추출: resolveToken 메서드를 통해 HTTP Authorization 헤더에서 JWT 토큰을 추출합니다.

 

토큰 유효성 검증: jwtTokenProvider.validateToken(token) 메서드를 호출하여 토큰의 유효성을 확인합니다. 유효한 토큰인 경우, jwtTokenProvider.getAuthentication(token)을 통해 사용자의 인증 정보를 가져옵니다.

 

예외 처리: 토큰이 잘못되었거나 만료된 경우, 커스텀 예외 CustomJwtExceptionHandler를 발생시켜 처리합니다.

 

추가 작업: 유효한 토큰인 경우, 추출된 사용자 ID를 STOMP 헤더(senderUserId)에 추가합니다. 이렇게 하면 이 ID를 통해 이후 메시지를 처리할 때 사용자를 식별할 수 있습니다.

 

 

📖 resolveToken 

private String resolveToken(StompHeaderAccessor accessor) {
    String bearerToken = accessor.getFirstNativeHeader(HttpHeaders.AUTHORIZATION);

    if(StringUtils.hasText(bearerToken) && bearerToken.startsWith("Bearer ")){
        return bearerToken.substring(7);
    }
    return null;
}

 

 

 

accessor.getFirstNativeHeader(HttpHeaders.AUTHORIZATION) : STOMP 헤더에서 Authorization 헤더를 가져옵니다.

 

StringUtils.hasText(bearerToken) && bearerToken.startsWith("Bearer ") : 토큰이 존재하고, "Bearer "로 시작하는지 확인합니다.

 

bearerToken.substring(7) : "Bearer " 이후의 실제 토큰 부분을 반환합니다.

이제 apic test를 통해서 테스트를 해보자. 

https://apic.app/online/#/tester

apic - The Complete API Solution

 

apic - The Complete API Solution

 

아래 그림을 확인하면 정상적으로 작동하는 모습을 확인 할 수 있다.