add login on front.

internal/api/websocket.go

@@ -3,15 +3,30 @@ package api
 import (
 	"log"
 	"net/http"
+	"os"
+	"shagram/internal/auth"
 	"shagram/internal/db"
 	"shagram/internal/websocket"
+	"strings"
 
 	"github.com/gin-gonic/gin"
 	gws "github.com/gorilla/websocket"
 )
 
 var upgrader = gws.Upgrader{
-	CheckOrigin: func(r *http.Request) bool { return true },
+	CheckOrigin: func(r *http.Request) bool {
+		origin := r.Header.Get("Origin")
+		if origin == "" {
+			return false
+		}
+		allowed := strings.Split(os.Getenv("WS_ALLOWED_ORIGINS"), ",")
+		for _, a := range allowed {
+			if strings.TrimSpace(a) == origin {
+				return true
+			}
+		}
+		return false
+	},
 }
 
 func WebSocketHandler(hub *websocket.Hub, database *db.DB) gin.HandlerFunc {
@@ -19,6 +34,14 @@ func WebSocketHandler(hub *websocket.Hub, database *db.DB) gin.HandlerFunc {
 		roomID := c.Param("room")
 		room := hub.GetOrCreateRoom(roomID)
 
+		tokenString := c.Query("token")
+		claims, err := auth.ParseAccessToken(tokenString)
+		if err != nil {
+			c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid token"})
+			return
+		}
+		username := claims.Username
+
 		conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
 		if err != nil {
 			return
@@ -35,25 +58,24 @@ func WebSocketHandler(hub *websocket.Hub, database *db.DB) gin.HandlerFunc {
 
 			for {
 				var msg map[string]string
-				err := client.Conn.ReadJSON(&msg)
-				if err != nil {
+				if err := client.Conn.ReadJSON(&msg); err != nil {
 					break
 				}
-				user := msg["user"]
-				if user == "" {
-					user = "Anonymous"
+
+				text := msg["text"]
+				if text == "" {
+					continue
 				}
+
 				_, err = database.Exec(`
 					INSERT INTO messages (room_id, user, text)
-					VALUES (?, ?, ?)`, roomID, user, msg["text"])
+					VALUES (?, ?, ?)`, roomID, username, text)
 				if err != nil {
 					log.Printf("Save message error: %v", err)
 				}
-				formattedMsg := user + ": " + msg["text"]
-				message := []byte(formattedMsg)
-				room.Broadcast(message)
+
+				room.Broadcast([]byte(username + ": " + text))
 			}
 		}()
 	}
-
 }