From aad2fb366fccf0981bd90ff3c045c427a0f09393 Mon Sep 17 00:00:00 2001
From: madipo2611 <klimov.aka@yandex.ru>
Date: Tue, 19 Aug 2025 13:50:14 +0300
Subject: [PATCH] =?UTF-8?q?v0.0.23=20=D0=BF=D1=80=D0=B0=D0=B2=D0=BA=D0=B8?=
 =?UTF-8?q?=20=D0=B2=20/ws?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/http/graph/messages_resolvers.go |  8 ++++----
 internal/http/middleware/ws_auth.go       | 25 ++++++++++++++++-------
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/internal/http/graph/messages_resolvers.go b/internal/http/graph/messages_resolvers.go
index 4795245..1bb2763 100644
--- a/internal/http/graph/messages_resolvers.go
+++ b/internal/http/graph/messages_resolvers.go
@@ -148,6 +148,9 @@ func (r *subscriptionResolver) MessageStream(ctx context.Context, userID int) (<
 func (r *subscriptionResolver) runMessageStream(ctx context.Context, userID int, messageChan chan<- *domain.Message) error {
 	log.Printf("Starting new stream for user %d", userID)
 
+	// Создаем отдельный контекст для gRPC стрима
+	grpcCtx := context.Background()
+
 	_, err := r.MessageClient.UpdateMessageStatus(ctx, &proto.UpdateMessageStatusRequest{
 		MessageId: 0, // 0 = все сообщения для пользователя
 		Status:    "DELIVERED",
@@ -157,10 +160,7 @@ func (r *subscriptionResolver) runMessageStream(ctx context.Context, userID int,
 		log.Printf("Failed to mark messages as delivered: %v", err)
 	}
 
-	streamCtx, cancel := context.WithCancel(ctx)
-	defer cancel()
-
-	stream, err := r.MessageClient.StreamMessages(streamCtx, &proto.StreamMessagesRequest{
+	stream, err := r.MessageClient.StreamMessages(grpcCtx, &proto.StreamMessagesRequest{
 		UserId: int32(userID),
 	})
 	if err != nil {
diff --git a/internal/http/middleware/ws_auth.go b/internal/http/middleware/ws_auth.go
index 667da89..ad3add5 100644
--- a/internal/http/middleware/ws_auth.go
+++ b/internal/http/middleware/ws_auth.go
@@ -3,8 +3,8 @@ package middleware
 import (
 	"context"
 	"net/http"
+	"strings"
 	"tailly_back_v2/pkg/auth"
-	"time"
 )
 
 // WSAuthMiddleware проверяет JWT токен для WebSocket соединений
@@ -12,16 +12,14 @@ func WSAuthMiddleware(tokenAuth *auth.TokenAuth) func(http.Handler) http.Handler
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			if r.Header.Get("Upgrade") == "websocket" {
-				ctx, cancel := context.WithTimeout(r.Context(), 30*time.Minute)
-				defer cancel()
-
 				token := extractTokenFromRequest(r)
 				if token != "" {
 					if userID, err := tokenAuth.ValidateAccessToken(token); err == nil {
-						ctx = context.WithValue(ctx, userIDKey, userID)
+						// Создаем контекст без таймаута для WebSocket
+						ctx := context.WithValue(r.Context(), userIDKey, userID)
+						r = r.WithContext(ctx)
 					}
 				}
-				r = r.WithContext(ctx)
 			}
 			next.ServeHTTP(w, r)
 		})
@@ -29,10 +27,23 @@ func WSAuthMiddleware(tokenAuth *auth.TokenAuth) func(http.Handler) http.Handler
 }
 
 func extractTokenFromRequest(r *http.Request) string {
-	// Только проверка кук (как в вашем коде)
+	// Проверяем куки
 	cookie, err := r.Cookie("accessToken")
 	if err == nil {
 		return cookie.Value
 	}
+
+	// Проверяем заголовок Authorization
+	authHeader := r.Header.Get("Authorization")
+	if authHeader != "" && strings.HasPrefix(authHeader, "Bearer ") {
+		return strings.TrimPrefix(authHeader, "Bearer ")
+	}
+
+	// Проверяем query параметры
+	token := r.URL.Query().Get("token")
+	if token != "" {
+		return token
+	}
+
 	return ""
 }