From 2e401a544fc44fa8f355b67a0738a1c77e941604 Mon Sep 17 00:00:00 2001
From: Harkamal Randhawa <harkamalsr@gmail.com>
Date: Tue, 10 Mar 2026 20:04:26 -0600
Subject: [PATCH] Fix chat auth

---
 .../WebSocketAuthChannelInterceptor.java      | 12 +++++++
 .../controller/ChatWebSocketController.java   | 35 ++++++++++++++++---
 2 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/src/main/java/com/petshop/backend/config/WebSocketAuthChannelInterceptor.java b/src/main/java/com/petshop/backend/config/WebSocketAuthChannelInterceptor.java
index 2b467025..6f01c384 100644
--- a/src/main/java/com/petshop/backend/config/WebSocketAuthChannelInterceptor.java
+++ b/src/main/java/com/petshop/backend/config/WebSocketAuthChannelInterceptor.java
@@ -59,10 +59,22 @@ public class WebSocketAuthChannelInterceptor implements ChannelInterceptor {
                     Collections.singletonList(new SimpleGrantedAuthority("ROLE_" + user.getRole().name()))
             );
             accessor.setUser(authentication);
+            accessor.getSessionAttributes().put("user", authentication);
+            return message;
+        }
+
+        if (StompCommand.DISCONNECT.equals(command) || StompCommand.UNSUBSCRIBE.equals(command)) {
             return message;
         }
 
         Principal principal = accessor.getUser();
+        if (principal == null && accessor.getSessionAttributes() != null) {
+            Object sessionUser = accessor.getSessionAttributes().get("user");
+            if (sessionUser instanceof Principal storedPrincipal) {
+                accessor.setUser(storedPrincipal);
+                principal = storedPrincipal;
+            }
+        }
         if (principal == null) {
             throw new IllegalArgumentException("Unauthenticated websocket session");
         }
diff --git a/src/main/java/com/petshop/backend/controller/ChatWebSocketController.java b/src/main/java/com/petshop/backend/controller/ChatWebSocketController.java
index 011268e6..bb30bd44 100644
--- a/src/main/java/com/petshop/backend/controller/ChatWebSocketController.java
+++ b/src/main/java/com/petshop/backend/controller/ChatWebSocketController.java
@@ -4,36 +4,61 @@ import com.petshop.backend.dto.chat.MessageRequest;
 import com.petshop.backend.dto.chat.MessageResponse;
 import com.petshop.backend.entity.User;
 import com.petshop.backend.repository.UserRepository;
+import com.petshop.backend.security.JwtUtil;
 import com.petshop.backend.service.ChatRealtimeService;
 import com.petshop.backend.service.ChatService;
 import jakarta.validation.Valid;
 import org.springframework.messaging.handler.annotation.DestinationVariable;
 import org.springframework.messaging.handler.annotation.MessageMapping;
 import org.springframework.messaging.handler.annotation.Payload;
+import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
 import org.springframework.messaging.simp.annotation.SendToUser;
-import org.springframework.security.core.Authentication;
 import org.springframework.stereotype.Controller;
 
+import java.security.Principal;
+
 @Controller
 public class ChatWebSocketController {
 
     private final ChatService chatService;
     private final ChatRealtimeService chatRealtimeService;
     private final UserRepository userRepository;
+    private final JwtUtil jwtUtil;
 
-    public ChatWebSocketController(ChatService chatService, ChatRealtimeService chatRealtimeService, UserRepository userRepository) {
+    public ChatWebSocketController(ChatService chatService, ChatRealtimeService chatRealtimeService, UserRepository userRepository, JwtUtil jwtUtil) {
         this.chatService = chatService;
         this.chatRealtimeService = chatRealtimeService;
         this.userRepository = userRepository;
+        this.jwtUtil = jwtUtil;
     }
 
     @MessageMapping("/chat/conversations/{id}/messages")
     @SendToUser("/queue/chat/errors")
-    public void sendMessage(@DestinationVariable Long id, @Valid @Payload MessageRequest request, Authentication authentication) {
-        User user = userRepository.findByUsername(authentication.getName())
-                .orElseThrow(() -> new IllegalArgumentException("User not found"));
+    public void sendMessage(@DestinationVariable Long id, @Valid @Payload MessageRequest request, SimpMessageHeaderAccessor headerAccessor) {
+        User user = resolveUser(headerAccessor);
         MessageResponse message = chatService.sendMessage(id, user.getId(), user.getRole(), request);
         chatRealtimeService.publishMessage(id, message);
         chatRealtimeService.publishConversationUpdate(id);
     }
+
+    private User resolveUser(SimpMessageHeaderAccessor headerAccessor) {
+        Principal principal = headerAccessor.getUser();
+        if (principal != null) {
+            return userRepository.findByUsername(principal.getName())
+                    .orElseThrow(() -> new IllegalArgumentException("User not found"));
+        }
+
+        String tokenHeader = headerAccessor.getFirstNativeHeader("Authorization");
+        if (tokenHeader == null || tokenHeader.isBlank()) {
+            tokenHeader = headerAccessor.getFirstNativeHeader("token");
+        }
+        if (tokenHeader == null || tokenHeader.isBlank()) {
+            throw new IllegalArgumentException("User not authenticated");
+        }
+
+        String token = tokenHeader.startsWith("Bearer ") ? tokenHeader.substring(7) : tokenHeader;
+        String username = jwtUtil.extractUsername(token);
+        return userRepository.findByUsername(username)
+                .orElseThrow(() -> new IllegalArgumentException("User not found"));
+    }
 }