[Spring] @EnableWebSocket으로 채팅 구현

수정중

---

 

 

1. WebSocketConfig 생성 

WebSocket에 연결을 하기 위해 먼저 WebSocketConfig를 만들었습니다.

package com.jam.client.chat.webSocket;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;

import com.jam.client.chat.service.ChatService;
import com.jam.client.member.service.MemberService;

import lombok.RequiredArgsConstructor;

@Configuration
@EnableWebSocket // 웹소켓 서버 사용
//@EnableWebSocketMessageBroker // STOMP 사용
@RequiredArgsConstructor
public class WebSocketConfig implements WebSocketConfigurer  {
	
    private final ChatService chatService;
    private final CustomWebSocketInterceptor customWebSocketInterceptor;

    @Bean
    public WebSocketHandler webSocketHandler() {
        return new WebSocketHandler(chatService);
    }

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(webSocketHandler(), "/ws")
        	.addInterceptors(customWebSocketInterceptor)
                .setAllowedOrigins("*");
    }
}

 

 

✔ @EnableWebSocket vs @EnableWebSocketMessageBroker

/	핸들러 기반 WebSocket	STOMP 방식
설정 어노테이션	@EnableWebSocket	@EnableWebSocketMessageBroker
핸들러 방식	TextWebSocketHandler 직접 구현	@MessageMapping + SimpMessagingTemplate
클라이언트 연결 방식	new WebSocket()	SockJS + Stomp.js
채팅방 라우팅	직접 구현	prefix로 자동 분기 (/pub, /sub)
확장성/관리	커스터마이징 쉬움	기능 많고 대규모에 적합

핸들러 방식은 클라이언트에게 메시지를 보내기 위해 WebSocketSession 객체를 직접 저장하고 관리해야 하며,
chatRoomId -> Set<Session> 구조처럼 매핑을 구현해야 합니다.

STOMP 방식은 구독 주소(/sub/xxx)에 따라 브로커가 자동으로 처리합니다.
예를 들어 클라이언트가 /sub/chat/room/123에 구독하고 있다면, 서버가 convertAndSend("/sub/chat/room/123", message)를 호출할 때 브로커가 자동으로 처리합니다.

저의 프로젝트는 Redis와 JWT를 활용하여 커스터마이징된 구조로 구성되어 있으며, 1:1 채팅 기능만을 지원합니다.
STOMP 방식의 강점인 구독 기반 메시지 처리, 자동 라우팅, 브로커를 통한 알림 기능 등이 필요하지 않기 때문에,
고수준 STOMP 방식 대신 로우레벨 WebSocket 방식(@EnableWebSocket)을 사용하여 직접 세션 및 메시지를 제어하도록 설정했습니다.

 

💡코드 설명

@Bean
public WebSocketHandler webSocketHandler() {
    return new WebSocketHandler(chatService);
}

• WebSocketHandler를 Spring Bean으로 등록합니다.
  이 핸들러는 클라이언트의 WebSocket 연결과 메시지를 처리하는 핵심 클래스입니다.

@Override
public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
    registry.addHandler(webSocketHandler(), "/ws")
        .addInterceptors(customWebSocketInterceptor)
        .setAllowedOrigins("*");
}

 

• WebSocket 엔드포인트를 등록하기 위해 WebSocketHandlerRegistry를 오버라이드합니다.
• /ws 경로로 들어오는 WebSocket 연결을 webSocketHandler()로 처리하겠다는 의미입니다.
  클라이언트는 ws://서버주소/ws로 WebSocket 연결을 시도하게 됩니다.
• 저는 JWT토큰으로 사용자 인증을 처리하기 위해 WebSocket 핸드셋 요청 시 실행되는 인터셉터를 추가했습니다.
• CORS 설정으로, 어떤 Origin에서든 WebSocket 연결을 허용하겠다는 의미입니다.
  운영 환경에서는 "*" 대신 허용된 Origin을 지정하는 게 더 안전하지만 현재 개발중이기 때문에 *로 설정했습니다.
  
  

2. WebSocketInterceptor

package com.jam.client.chat.webSocket;

import java.util.HashMap;
import java.util.Map;

import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.http.server.ServletServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import com.jam.global.jwt.JwtService;

import lombok.extern.log4j.Log4j;

@Log4j
@Component
public class CustomWebSocketInterceptor implements HandshakeInterceptor {

	private final JwtService jwtService;

	public CustomWebSocketInterceptor(JwtService jwtService) {
		this.jwtService = jwtService;
	}
	
	@Override
	public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
			Map<String, Object> attributes) throws Exception {
		
		if (request instanceof ServletServerHttpRequest) {
			
	        ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) request;
	        HttpServletRequest httpRequest = servletRequest.getServletRequest();

	        ServletServerHttpResponse servletResponse = (ServletServerHttpResponse) response;
	        HttpServletResponse httpResponse = servletResponse.getServletResponse();
	        
	        Map<String, String> userMap = new HashMap<>();
	        
	        Cookie[] cookies = httpRequest.getCookies();
	        
	        if (cookies != null) {
	        	userMap = jwtService.getUserInfo(cookies, httpRequest, httpResponse); 
	        }
	        
	        if(userMap != null) {
		        attributes.put("userId", userMap.get("userId"));
		        attributes.put("auth", userMap.get("auth"));
	        }
	    }
	    return true;
	}
}

• WebSocket 연결 전에 JWT 토큰을 통해 사용자가 인증된 사용자임을 확인하고, WebSocketSession의 attributes에 사용자 ID와 권한을 저장합니다
• 저장한 정보는 WebSocketHandler에서 session.getAttributes().get("userId") 형태로 접근 가능합니다.

 

3. WebSocketHandler 구현

package com.jam.client.chat.webSocket;

import java.io.IOException;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.jam.client.chat.service.ChatService;
import com.jam.client.chat.vo.ChatVO;
import com.jam.global.util.JsonUtils;

import lombok.RequiredArgsConstructor;
import lombok.extern.log4j.Log4j;

@Log4j
@RequiredArgsConstructor
public class WebSocketHandler extends TextWebSocketHandler  {

    private final ChatService chatService;
    
    // 현재 연결된 모든 WebSocket 세션을 관리
    private final Set<WebSocketSession> sessions = java.util.concurrent.ConcurrentHashMap.newKeySet();

    /* chatRoomId: {session1, session2}
     * 채팅에 입장하면 세션 추가됨. 채팅 나가면 세션 삭제
     * 특정 채팅방에 참여한 세션들을 관리. 
     * 키: 채팅방 ID, 값: 해당 채팅방에 연결된 세션들 
     * 특정 채팅방에 속한 클라이언트들에게만 메시지를 보냄.*/
    private final Map<String,Set<WebSocketSession>> chatRoomSession = new ConcurrentHashMap<>();

    /* session: chatRoomId
     * 특정 세션이 참여한 채팅방 관리, 채팅방 나갈 때 채팅방 id 찾기 위함  */
    private final Map<WebSocketSession, String> sessionToChatRoom = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        log.info("소켓 연결됨" + session.getId());
        sessions.add(session);
    }

    /*  상태코드: 메시지
		401: 인증 안 됨(로그인/토큰 없음)
		403: 인증은 됐는데 권한 없음 (방 멤버 아님)
		404: 방ID나 상대방ID가 없음
		400: 요청 값 이상(없거나 형식 오류)
		500: 서버 오류
	*/
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
    	try {
        	String payload = message.getPayload();
        	
            // 페이로드(JSON 형식의 메시지 데이터) -> ChatVO로 변환
            ChatVO chatVO = convertToChatMessageVO(payload);

            String chatRoomId = chatVO.getChatRoomId();
            
    	    Map<String, Object> attributes = session.getAttributes();
    	    String userId = (String) attributes.get("userId");
            
    	    if (userId == null) {
    	    	sendError(session, 401, "로그인이 필요한 서비스 입니다. 로그인 하시겠습니까?");
    	    	session.close(CloseStatus.POLICY_VIOLATION);
    	        return;
    	    }
    	    
            // 채팅방 세션 없으면 만듦
            if (!chatRoomSession.containsKey(chatRoomId)) {
                chatRoomSession.put(chatRoomId, ConcurrentHashMap.newKeySet());
            }
            
            // 특정 채팅방에 연결된 세션 집합
            Set<WebSocketSession> sessionSet = chatRoomSession.get(chatRoomId);

            // ENTER(채팅방 입장)
            if (chatVO.getType().equals(ChatVO.Type.ENTER)) {
            	
            	sessionSet.add(session);
            	sessionToChatRoom.put(session, chatRoomId);
            	
        	    Map<String, String> info = chatService.getChatPartner(chatVO.getChatRoomId(), userId);
        	    
        	    if (info == null || info.get("chatPartnerId") == null) {
        	    	sessionSet.remove(session);
        	    	if (sessionSet.isEmpty()) chatRoomSession.remove(chatRoomId);
        	    	sessionToChatRoom.remove(session);
        	    	sendMessage(session, "ERROR", Map.of("code", 404, "message", "채팅방을 찾을 수 없습니다."));
        	    	return;
        		}
        	    
        	    session.getAttributes().put("partnerId", info.get("chatPartnerId"));
        	    session.getAttributes().put("partnerName", info.get("chatPartnerName"));
        	    
        	    sendMessage(session, "PARTNER_INFO", Map.of("partnerName", info.get("chatPartnerName"))); 
        	    
        	    log.info("사용자 " + session.getId() + "가 채팅방 " + chatRoomId + "에 입장했습니다.");
            } else if (chatVO.getType().equals(ChatVO.Type.LEAVE)) {
            	
                // 1. chatRoomSession에서 제거
                if (chatRoomSession.containsKey(chatRoomId)) {
                    sessionSet.remove(session);

                    // 채팅방에 세션이 없으면 맵에서 키 제거
                    if (sessionSet.isEmpty()) {
                        chatRoomSession.remove(chatRoomId);
                    }
                }

                sessionToChatRoom.remove(session);

                log.info("사용자 " + session.getId() + "가 채팅방 " + chatRoomId + "에서 퇴장했습니다.");
            
            } else if (chatVO.getType().equals(ChatVO.Type.MESSAGE)) {
            	try {

            		// payload 검증
            		if (chatRoomId == null || chatVO.getMessage() == null || chatVO.getMessage().isBlank()) {
            			sendMessage(session, "ERROR", Map.of("code", 400, "message", "잘못된 요청 입니다. 다시 시도해 주세요."));
            			return;
            		}
            		// 방 세션 검증
            		if (sessionSet == null || !sessionSet.contains(session)) {
            			sendError(session, 403, "잘못된 접근 입니다.");
            			session.close(CloseStatus.POLICY_VIOLATION);
            			return;
            		}
            		
            		// redis에 채팅방 있는지
            		if (!chatService.roomExists(chatRoomId)) {
            			sendMessage(session, "ERROR", Map.of("code", 404, "message", "채팅방을 찾을 수 없습니다."));
            			return;
            		}
            		
            		// Redis에서 멤버십 확인
            		if (!chatService.isMemberOfRoom(userId, chatRoomId)) {
            			sendMessage(session, "ERROR", Map.of("code", 403, "message", "잘못된 접근 입니다."));
            			return;
            		}
            		
            		String partnerId = (String) attributes.get("partnerId");
            		
            		if (partnerId == null) {
            			sendMessage(session, "ERROR", Map.of("code", 404, "message", "채팅방을 찾을 수 없습니다."));
            			return;
            		}
            		
            		chatVO.setReceiverId(partnerId);
            	    chatVO.setSenderId(userId);

            	    // 첫 메시지인지 확인
            	    boolean isFirst = chatService.ensureRoomOnFirstMessage(chatRoomId, userId, partnerId);

            		if (isFirst) { 
            	    	String userName = chatService.getUserNameFromRedis(userId);
            			String partnerName = (String) attributes.get("partnerName");
            			
            			if (chatVO.getChatDate() == null) chatVO.setChatDate(LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm")));
            			
            			if (userName != null && partnerName != null) {
                			broadcastRoomCreated(chatRoomId, userId, userName, partnerId, partnerName, chatVO.getMessage(), chatVO.getChatDate());
                		}
            		}
            		
            	    // 메시지 저장 후 해당 방에만 브로드캐스트
            	    chatService.saveChat(chatVO);
            	    sendMessageToChatRoom(chatVO, "MESSAGE", sessionSet);
                    
        	    } catch (Exception e) {
        	        throw new RuntimeException("Failed to save chat message", e);
        	    }
            }
        } catch (Exception e) {
            log.error("Exception: " + e.getClass().getName() + ": " + e.getMessage());
            e.printStackTrace();
            
            sendError(session, 500,  "서버에 문제가 발생했습니다. 잠시 후 다시 시도해 주세요.");
        }
    }

    // 클라이언트가 연결 끊음
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        
    	log.info("연결 끊김" + session.getId());
    	
        sessions.remove(session);
        
        // 세션이 속한 채팅방 ID 찾기
        //String chatRoomId = sessionToChatRoom.get(session);
        String chatRoomId = sessionToChatRoom.remove(session);
        
        if (chatRoomId != null) {
            // 해당 채팅방의 세션 집합에서 제거
            Set<WebSocketSession> chatRoomSessions = chatRoomSession.get(chatRoomId);
            if (chatRoomSessions != null) {
                chatRoomSessions.remove(session);
            }
        }

        chatRoomSession.values().forEach(this::removeClosedSession);
    }

    // 닫힌 세션을 chatRoomSession에서 제거
    private void removeClosedSession(Set<WebSocketSession> chatRoomSession) {
        chatRoomSession.removeIf(sess -> !sessions.contains(sess));
    }
    
    private void sendMessage(WebSocketSession session, String type, Object data) {
        if (session == null || !session.isOpen()) {
            log.warn("Skipping message. Session is null or closed.");
            return;
        }
        
        try {
            Map<String, Object> payload = new HashMap<>();
            payload.put("type", type);
            payload.put("data", data);

            String json = JsonUtils.toJson(payload);
            session.sendMessage(new TextMessage(json));
        } catch (IOException e) {
            log.error("Error sending WebSocket message: " + e.getMessage(), e);
        }
    }


    // 특정 채팅방의 모든 클라이언트에게 메시지를 전송
    private void sendMessageToChatRoom(ChatVO chatVO, String type, Set<WebSocketSession> sessionSet) {
    	sessionSet.stream().forEach(sess -> sendMessage(sess, type, chatVO));
    }
    
    // 현재 연결된 세션들 중 userId가 일치하는 모든 세션에 전송 (PC+모바일에서 동시 사용해도 새 채팅방 생성되도록)
    private void sendMessageToUser(String userId, String type, Object data) {
    	for (WebSocketSession s : sessions) {
    		if (!s.isOpen()) continue;
    		Object uid = s.getAttributes().get("userId");
    		if (userId != null && userId.equals(uid)) {
    			sendMessage(s, type, data);
    		}
    	}
    }
    
    private void broadcastRoomCreated(
    		String chatRoomId, 
    		String senderId, String senderName, 
    		String partnerId, String partnerName, 
    		String firstMessage, String chatDate) {

    	// 송신자에게는 partner = 상대 이름
    	Map<String, Object> toSender = Map.of(
    		"chatRoomId", chatRoomId,
    		"partnerId", partnerId,
    		"partner", partnerName,
    		"message", firstMessage,
    		"chatDate", chatDate
    	);
    	sendMessageToUser(senderId, "ROOM_CREATED", toSender);

    	// 수신자에게는 partner = 보내는 사람 이름
    	Map<String, Object> toPartner = Map.of(
    		"chatRoomId", chatRoomId,
    		"partnerId", senderId,
    		"partner", senderName,
    		"message", firstMessage,
    		"chatDate", chatDate
    	);
    	sendMessageToUser(partnerId, "ROOM_CREATED", toPartner);
    }

    private void sendError(WebSocketSession session, int code, String msg) {
    	sendMessage(session, "ERROR", Map.of("code", code, "message", msg));
    }

    public ChatVO convertToChatMessageVO(String payload) {
		ObjectMapper objectMapper = new ObjectMapper();
        try {
            return objectMapper.readValue(payload, ChatVO.class);
        } catch (IOException e) {
            throw new RuntimeException("Failed to convert payload to VO", e);
        }
	}
}

 

 

4. 클라이언트 측 

STOMP 방식이라면 라이브러리 추가해야 함
<script src="https://cdn.jsdelivr.net/npm/sockjs-client@1/dist/sockjs.min.js"></script>
<script src="https://cdn.jsdelivr.net/npm/stompjs@2.3.3/lib/stomp.min.js"></script>

로우레벨 WebSocket 방식은 추가해야 할 스트립트 없습니당

function initWebSocket(chatRoomId) {
	try {
		// 웹소켓 중복 연결 방지
		if (socket) {
			if (socket.readyState === WebSocket.OPEN) {
				socket.send(JSON.stringify({ type: "ENTER", chatRoomId: roomId }));
				return;
			}
			if (socket.readyState === WebSocket.CONNECTING) {
				return; // 연결 중이면 대기
			}
		}
		
		socket = new WebSocket("ws://localhost:8080/ws");	
		
		let retried = false;
		socket.onopen = function() { 
		  retried = false; 
		  socket.send(JSON.stringify({ type: "ENTER", chatRoomId: roomId })); 
		};
        
        // LEAVE 메시지를 전송하고 싶다면
        // socket.send(JSON.stringify({ type: "LEAVE", chatRoomId: roomId }));

		socket.onclose = function() {
		  console.log("WebSocket is closed now.");
		  if (currentRoomId && !retried) {
		    retried = true;
		    setTimeout(() => initWebSocket(currentRoomId), 2000);
		  }
		};
		
		socket.onerror = function(error) {
			console.log("WebSocket error: " + error);
		};

		socket.onmessage = function(event) {
			const { type, data } = JSON.parse(event.data);
			if (type === "MESSAGE") {
				// 메시지 수신
			} else if (type === "PARTNER_INFO") {
				// 상대방 정보 갱신
			} else if (type === 'ROOM_CREATED'){
				// 처음 나눈 대화라면 채팅방 목록에 채팅 추가
			} else if (type === "ERROR") {
				sessionStorage.removeItem("chatRoomId");

				if(data.code === 401){
					// 로그인 되지 않음
				}else if ([403, 404].includes(data.code)) {
					// 채팅방에 접근 권한 없음
				} else { // 400, 500
					// 채팅방ID 또는 상대방 정보가 없거나 서버 오류
				}
			}
		};
	} catch (err) {
		console.error("채팅 초기화 중 오류 발생:", err);
	}
}