[트러블슈팅] WebSocket + STOMP 적용하여 알림 구현하기

[간단 요약]

StompHeaderAccessor.wrap(message)은 변경이 실제 메시지에 반영 안 되는 경우가 있다. 따라서 CONNECT 때는 담겨 있던 인증 객체 정보가 SUBSCRIBE 때 사라지는 경우가 발생한다.
이를 해결하기 위해선 accessor를 StompHeaderAccessor.wrap(message) 대신 MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class)로 선언해준다. 그리고
CONNECT 직후 accessor.setLeaveMutable(true); 처리를 해주고, preSend() 마지막의 메시지도 MessageBuilder.createMessage(...)로 다시 만들어 준다.
Tip: 트러블슈팅은 GPT가 잘한다.

[Github Commit]

https://github.com/all-in-market/alarm-server/pull/1

Feat/restock notification by ginsengcandy · Pull Request #1 · all-in-market/alarm-server

[문제 상황]

1. 요구사항

: 다음 요구사항에 따라 실시간 알림 수신을 WebSocket + STOMP로 구현하려고 했다.

 

1. 고객이 품절 상품 재입고 알림 신청
2. 판매자가 재고 증가
3. notification DB 저장
4. 실시간 websocket push
5. 고객은 실시간 수신 + API 조회 가능

2. WebSocket/STOMP 설계

: 재입고 알림을 신청한 고객 개인에게 알림이 전송된다.

- 개인 알림 subscribe path

/user/queue/notifications

- 서버에서 전송 시

messagingTemplate.convertAndSendToUser(
	userId.toString(),
    "/queue/notifications",
    payload
);

 

이 조합은 다음과 같은 장점이 있다. 첫째, user destination routing을 자동 지원하며, 둘째, session별 private delivery가 가능하다.

3. Broker 설정

@Override
public void configureMessageBroker(MessageBrokerRegistry registry) {
    registry.enableSimpleBroker("/queue", "/topic");
    registry.setApplicationDestinationPrefixes("/app");
    registry.setUserDestinationPrefix("/user");
}

• /queue: 개인 알림
• /topic: broadcast reserve (미사용)
• /app: client -> server (미사용)
• /user: private routing

4. JWT 인증 + Principal Binding

: WebSocket CONNECT 시 JWT 인증 후 session에 Principal을 저장해야 했다.

: JwtChannelInterceptor로 하여금 CONNECT 요청만 인터셉트하도록 설정했다.

if (StompCommand.CONNECT.equals(accessor.getCommand())) {
    authenticate(accessor);

    return MessageBuilder.createMessage(
            message.getPayload(),
            accessor.getMessageHeaders()
    );
}

이후 JWT에서 토큰을 추출하여 검증하고, 사용자 아이디와 권한을 추출하여 Principal 객체와 Authentication 객체를 생성했다.

UserPrincipal principal = new UserPrincipal(userId, userRole);

Authentication authentication =
    new UsernamePasswordAuthenticationToken(
        principal,
        null,
        Collections.emptyList()
    );

특히 Principal에서 getName()이 userId를 반환하게 하여 2단계에서 개별 사용자 전송용 경로를 만들 때 항상 userId를 포함시키도록 했다.

public record UserPrincipal(
    Long userId,
    UserRole userRole
) implements Principal {

    @Override
    public String getName() {
        return String.valueOf(userId);
    }
}

문제는 CONNECT 때 존재하던 user가 SUBSCRIBE 때 유실됐다는 것이다(user = null). 참고로 두 프레임 모두에서 세션 ID 는 동일했다. 즉 JWT나 reconnect 문제는 아니었다. 유일한 경우의 수는 Message mutation persistence 문제였다.

기존에 사용했던 코드 StompHeaderAccessor.wrap(message)는 새 wrapper accessor를 만들 수 있는데, 그렇게 되면 실제 message-bound accessor state(user/session)과 분리될 수 있다.

그 결과 CONNECT에서 setUser 했지만 SUBSCRIBE에서 user가 유실된다.

[해결 과정]

StompHeaderAccessor.wrap(message)를 다음과 같이 수정했다.

이는 message에 attach된 실제 accessor를 직접 조회하는 형식으로 바꾼 것이다. 이래야 실제 session-bound principal state가 유지된다.

StompHeaderAccessor accessor =
    MessageHeaderAccessor.getAccessor(
        message,
        StompHeaderAccessor.class
    );

이 변경 후 user propagation이 정상화되었다.

두 번째 문제는 CONNECT 후 accessor.setUser(authentication);만 하고 return message;를 통해 원본 message만 반환했다는 것이다. 이 경우 mutated headers downstream에 반영이 안 될 수 있다는 문제가 있었다.

그래서 CONNECT branch에서 rebuilt message를 반환했다. CONNECT에서 setUser()만으로 끝내지 말고, mutated headers 포함한 새 message를 반환한 것이다. 

return MessageBuilder.createMessage(
    message.getPayload(),
    accessor.getMessageHeaders()
);

최종 코드는 다음과 같다.

authenticate(accessor);
accessor.setUser(authentication);
accessor.setLeaveMutable(true);

return MessageBuilder.createMessage(
    message.getPayload(),
    accessor.getMessageHeaders()
);

 

[교훈]

로깅과 디버깅이 중요하다.

트러블슈팅은 chatGPT가 잘한다.

JWT 인증 객체 정보를 STOMP SUBSCRIBER까지 유지하려면 wrap을 사용하면 안 된다.