Refactor WebSocket client handling and token validation

Centralize WebSocket client initialization in the main server setup and remove redundant code. Introduce middleware token validation in WebSocket connections to verify user authenticity and enhance logging for client connections and disconnections.

cmd/api/main.go

@@ -42,6 +42,9 @@ func main() {
 	// Connect NATS
 	natServer, err := nats.NewNATS()
 
+	// WebSocket Clients
+	wsClients := make(map[string]*websocket.Conn)
+
 	s := handler.NewServer(db)
 
 	// Initialize Validator
@@ -71,7 +74,6 @@ func main() {
 	delv.NewDeliveryHandler(s, "/delivery", deliveryService, middlewares, validate)
 
 	// Notification
-	wsClients := make(map[string]*websocket.Conn)
 	notifyService := notification.NewNotificationService(db, env, natServer)
 	err = notifyService.SubscribeNewOrders(wsClients)
 	if err != nil {

cmd/api/middleware/middleware.go

@@ -16,7 +16,7 @@ type UserClaims struct {
 	jwt.RegisteredClaims
 }
 
-func validateToken(token string) (bool, int64) {
+func ValidateToken(token string) (bool, int64) {
 
 	tokenInfo, err := jwt.ParseWithClaims(token, &UserClaims{}, func(token *jwt.Token) (interface{}, error) {
 		return []byte(os.Getenv("JWT_SECRET_KEY")), nil
@@ -50,7 +50,7 @@ func AuthMiddleware() gin.HandlerFunc {
 		}
 		token := tokenParts[1]
 
-		tokenValidation, userID := validateToken(token)
+		tokenValidation, userID := ValidateToken(token)
 		if !tokenValidation {
 			c.JSON(http.StatusUnauthorized, gin.H{"error": "Invalid Token"})
 			c.Abort()

pkg/handler/notification/notify.go

@@ -1,8 +1,10 @@
 package notification
 
 import (
+	"Go_Food_Delivery/cmd/api/middleware"
 	"github.com/gin-gonic/gin"
 	"log"
+	"strconv"
 )
 
 func (s *NotifyHandler) notifyOrders(c *gin.Context) {
@@ -19,15 +21,23 @@ func (s *NotifyHandler) notifyOrders(c *gin.Context) {
 		conn.Close()
 		return
 	}
-	userId := "1"
-	s.clients[userId] = conn
-	log.Printf("New client connected::%s", userId)
+	valid, userIdInt := middleware.ValidateToken(token)
+	if !valid {
+		log.Println("Invalid Token!")
+		conn.Close()
+		return
+	}
+
+	userID := strconv.FormatInt(userIdInt, 10)
+
+	s.clients[userID] = conn
+	log.Printf("New client connected::%s", userID)
 
 	for {
 		_, _, err := conn.ReadMessage()
 		if err != nil {
-			log.Printf("Client disconnected: %s::%v", userId, err)
-			delete(s.clients, userId)
+			log.Printf("Client disconnected: %s::%v", userID, err)
+			delete(s.clients, userID)
 			break
 		}
 	}

pkg/nats/nats_server.go

@@ -1,7 +1,6 @@
 package nats
 
 import (
-	"fmt"
 	"github.com/gorilla/websocket"
 	"github.com/nats-io/nats.go"
 	"log"
@@ -34,10 +33,8 @@ func (n *NATS) Sub(topic string, clients map[string]*websocket.Conn) error {
 	_, err := n.Conn.Subscribe(topic, func(msg *nats.Msg) {
 		message := string(msg.Data)
 		slog.Info("MESSAGE_REPLY_FROM_NATS", "RECEIVED_MESSAGE", message)
-		fmt.Println("CLIENTS::", clients)
 		userId, messageData := n.formatMessage(message)
 		if conn, ok := clients[userId]; ok {
-			fmt.Println("SENDING_MESSAGE_TO_CLIENT")
 			err := conn.WriteMessage(websocket.TextMessage, []byte(messageData))
 			if err != nil {
 				log.Println("Error sending message to client:", err)