package handlers import ( "context" "log" "net/http" "tailly_back_v2/internal/domain" "tailly_back_v2/internal/service" "tailly_back_v2/internal/ws" "tailly_back_v2/pkg/auth" "time" "github.com/gorilla/websocket" ) var upgrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, CheckOrigin: func(r *http.Request) bool { return true }, Subprotocols: []string{"graphql-transport-ws"}, } type ChatHandler struct { chatService service.ChatService hub *ws.Hub tokenAuth *auth.TokenAuth } func NewChatHandler(chatService service.ChatService, hub *ws.Hub, tokenAuth *auth.TokenAuth) *ChatHandler { return &ChatHandler{ chatService: chatService, hub: hub, tokenAuth: tokenAuth, } } func (h *ChatHandler) HandleWebSocket(w http.ResponseWriter, r *http.Request) { log.Printf("Incoming WebSocket headers: %+v", r.Header) log.Printf("Cookies: %+v", r.Cookies()) // Убираем все проверки токена здесь, так как будем проверять его в первом сообщении conn, err := upgrader.Upgrade(w, r, nil) if err != nil { log.Printf("WebSocket upgrade error: %v", err) http.Error(w, "Could not upgrade to WebSocket", http.StatusBadRequest) return } client := &ws.Client{ UserID: 0, // Пока не авторизован Send: make(chan *domain.Message, 256), } h.hub.RegisterClient(client) ctx, cancel := context.WithCancel(context.Background()) defer cancel() go h.readPump(ctx, conn, client) go h.writePump(ctx, conn, client) <-ctx.Done() } func (h *ChatHandler) readPump(ctx context.Context, conn *websocket.Conn, client *ws.Client) { ticker := time.NewTicker(30 * time.Second) defer ticker.Stop() defer conn.Close() // 1. Ожидаем первое сообщение с токеном var authMsg struct { Type string `json:"type"` Token string `json:"token"` } if err := conn.ReadJSON(&authMsg); err != nil { log.Printf("Failed to read auth message: %v", err) return } if authMsg.Type != "auth" || authMsg.Token == "" { log.Println("First message must be auth with token") conn.WriteJSON(map[string]string{"type": "error", "message": "Auth required"}) return } // 2. Валидируем токен userID, err := h.tokenAuth.ValidateAccessToken(authMsg.Token) if err != nil { log.Printf("Token validation error: %v", err) conn.WriteJSON(map[string]string{"type": "error", "message": "Invalid token"}) return } client.UserID = userID log.Printf("WebSocket authenticated, userID=%d", userID) // 3. Отправляем подтверждение авторизации if err := conn.WriteJSON(map[string]string{"type": "auth_success"}); err != nil { log.Printf("Failed to send auth confirmation: %v", err) return } // 4. Основной цикл обработки сообщений for { select { case <-ticker.C: if err := conn.WriteJSON(map[string]string{"type": "ping"}); err != nil { log.Printf("Ping error: %v", err) return } case <-ctx.Done(): return default: var msg struct { Type string `json:"type"` Payload struct { ReceiverID int `json:"receiverId"` Content string `json:"content"` } `json:"payload"` } if err := conn.ReadJSON(&msg); err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) { log.Printf("WebSocket error: %v", err) } return } switch msg.Type { case "pong": continue case "ping": conn.WriteJSON(map[string]string{"type": "pong"}) continue case "message": if client.UserID == 0 { conn.WriteJSON(map[string]string{"type": "error", "message": "Not authenticated"}) continue } if msg.Payload.ReceiverID == 0 { conn.WriteJSON(map[string]interface{}{ "type": "error", "message": "Invalid receiver ID", }) continue } chat, err := h.chatService.GetOrCreateChat(ctx, client.UserID, msg.Payload.ReceiverID) if err != nil { log.Printf("Chat error: %v", err) conn.WriteJSON(map[string]interface{}{ "type": "error", "message": "Chat error", "details": err.Error(), }) continue } message, err := h.chatService.SendMessage( ctx, client.UserID, chat.ID, msg.Payload.Content, ) if err != nil { log.Printf("Message send error: %v", err) continue } h.hub.Broadcast(message) default: conn.WriteJSON(map[string]string{"type": "error", "message": "Unknown message type"}) } } } } func (h *ChatHandler) writePump(ctx context.Context, conn *websocket.Conn, client *ws.Client) { defer conn.Close() for { select { case <-ctx.Done(): return case message, ok := <-client.Send: if !ok { // Канал закрыт conn.WriteMessage(websocket.CloseMessage, []byte{}) return } if err := conn.WriteJSON(message); err != nil { log.Printf("WebSocket write error: %v", err) return } } } }