274 lines
7.2 KiB
Go
274 lines
7.2 KiB
Go
package main
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"log"
|
||
"net/http"
|
||
"os"
|
||
"os/signal"
|
||
"strings"
|
||
"sync"
|
||
"syscall"
|
||
"time"
|
||
|
||
"github.com/golang-jwt/jwt/v5"
|
||
"github.com/gorilla/websocket"
|
||
"github.com/rabbitmq/amqp091-go"
|
||
"github.com/redis/go-redis/v9"
|
||
)
|
||
|
||
type Notification struct {
|
||
UserID string `json:"userId"`
|
||
Type string `json:"type"`
|
||
Title string `json:"title"`
|
||
Message string `json:"message"`
|
||
Payload json.RawMessage `json:"payload,omitempty"`
|
||
}
|
||
|
||
type Hub struct {
|
||
mu sync.RWMutex
|
||
connections map[string]map[*websocket.Conn]struct{}
|
||
redis *redis.Client
|
||
}
|
||
|
||
func NewHub(redisClient *redis.Client) *Hub {
|
||
return &Hub{
|
||
connections: make(map[string]map[*websocket.Conn]struct{}),
|
||
redis: redisClient,
|
||
}
|
||
}
|
||
|
||
func (h *Hub) Add(ctx context.Context, userID string, conn *websocket.Conn) {
|
||
h.mu.Lock()
|
||
defer h.mu.Unlock()
|
||
if h.connections[userID] == nil {
|
||
h.connections[userID] = make(map[*websocket.Conn]struct{})
|
||
}
|
||
h.connections[userID][conn] = struct{}{}
|
||
_ = h.redis.SAdd(ctx, "ws:online_users", userID).Err()
|
||
_ = h.redis.Set(ctx, "ws:user:"+userID, "online", 90*time.Second).Err()
|
||
}
|
||
|
||
func (h *Hub) Remove(ctx context.Context, userID string, conn *websocket.Conn) {
|
||
h.mu.Lock()
|
||
defer h.mu.Unlock()
|
||
delete(h.connections[userID], conn)
|
||
if len(h.connections[userID]) == 0 {
|
||
delete(h.connections, userID)
|
||
_ = h.redis.SRem(ctx, "ws:online_users", userID).Err()
|
||
_ = h.redis.Del(ctx, "ws:user:"+userID).Err()
|
||
}
|
||
}
|
||
|
||
func (h *Hub) Send(userID string, notification Notification) {
|
||
h.mu.RLock()
|
||
conns := h.connections[userID]
|
||
h.mu.RUnlock()
|
||
for conn := range conns {
|
||
if err := conn.WriteJSON(notification); err != nil {
|
||
log.Printf("Не удалось отправить уведомление пользователю %s: %v", userID, err)
|
||
}
|
||
}
|
||
}
|
||
|
||
func (h *Hub) Broadcast(notification Notification) {
|
||
h.mu.RLock()
|
||
defer h.mu.RUnlock()
|
||
for userID, conns := range h.connections {
|
||
for conn := range conns {
|
||
if err := conn.WriteJSON(notification); err != nil {
|
||
log.Printf("Не удалось отправить широковещательное уведомление пользователю %s: %v", userID, err)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
func env(key, fallback string) string {
|
||
if value := os.Getenv(key); value != "" {
|
||
return value
|
||
}
|
||
return fallback
|
||
}
|
||
|
||
func resolveUserID(r *http.Request) (string, error) {
|
||
if token := strings.TrimSpace(r.URL.Query().Get("token")); token != "" {
|
||
secret := env("JWT_ACCESS_SECRET", "docker-access-secret")
|
||
parsed, err := jwt.Parse(token, func(token *jwt.Token) (interface{}, error) {
|
||
if token.Method != jwt.SigningMethodHS256 {
|
||
return nil, errors.New("неподдерживаемый алгоритм JWT")
|
||
}
|
||
return []byte(secret), nil
|
||
})
|
||
if err != nil || !parsed.Valid {
|
||
return "", err
|
||
}
|
||
claims, ok := parsed.Claims.(jwt.MapClaims)
|
||
if !ok {
|
||
return "", errors.New("некорректные claims JWT")
|
||
}
|
||
sub, _ := claims["sub"].(string)
|
||
if sub == "" {
|
||
return "", errors.New("JWT не содержит sub")
|
||
}
|
||
return sub, nil
|
||
}
|
||
|
||
userID := strings.TrimSpace(r.URL.Query().Get("user_id"))
|
||
if userID == "" {
|
||
return "", errors.New("не передан token или user_id")
|
||
}
|
||
return userID, nil
|
||
}
|
||
|
||
func main() {
|
||
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
||
defer stop()
|
||
|
||
redisClient := redis.NewClient(&redis.Options{Addr: env("REDIS_ADDR", "localhost:6379")})
|
||
if err := redisClient.Ping(ctx).Err(); err != nil {
|
||
log.Printf("Redis недоступен при старте, сервис продолжит переподключаться: %v", err)
|
||
}
|
||
|
||
hub := NewHub(redisClient)
|
||
upgrader := websocket.Upgrader{
|
||
CheckOrigin: func(r *http.Request) bool {
|
||
return true
|
||
},
|
||
}
|
||
|
||
http.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) {
|
||
w.WriteHeader(http.StatusOK)
|
||
_, _ = w.Write([]byte("ok"))
|
||
})
|
||
|
||
http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
|
||
userID, err := resolveUserID(r)
|
||
if err != nil || userID == "" {
|
||
http.Error(w, "Не удалось авторизовать WebSocket", http.StatusUnauthorized)
|
||
return
|
||
}
|
||
|
||
conn, err := upgrader.Upgrade(w, r, nil)
|
||
if err != nil {
|
||
log.Printf("Не удалось открыть WebSocket: %v", err)
|
||
return
|
||
}
|
||
hub.Add(r.Context(), userID, conn)
|
||
defer func() {
|
||
hub.Remove(context.Background(), userID, conn)
|
||
_ = conn.Close()
|
||
}()
|
||
|
||
for {
|
||
if err := conn.SetReadDeadline(time.Now().Add(2 * time.Minute)); err != nil {
|
||
log.Printf("Не удалось обновить таймаут чтения: %v", err)
|
||
return
|
||
}
|
||
_, _, err := conn.ReadMessage()
|
||
if err != nil {
|
||
return
|
||
}
|
||
_ = redisClient.Set(r.Context(), "ws:user:"+userID, "online", 90*time.Second).Err()
|
||
}
|
||
})
|
||
|
||
go consumeRabbit(ctx, hub)
|
||
|
||
server := &http.Server{
|
||
Addr: ":" + env("PORT", "8085"),
|
||
ReadHeaderTimeout: 10 * time.Second,
|
||
}
|
||
|
||
go func() {
|
||
log.Printf("Media WS запущен на %s", server.Addr)
|
||
if err := server.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) {
|
||
log.Fatalf("Ошибка HTTP сервера: %v", err)
|
||
}
|
||
}()
|
||
|
||
<-ctx.Done()
|
||
shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||
defer cancel()
|
||
if err := server.Shutdown(shutdownCtx); err != nil {
|
||
log.Printf("Не удалось корректно остановить сервер: %v", err)
|
||
}
|
||
_ = redisClient.Close()
|
||
}
|
||
|
||
func consumeRabbit(ctx context.Context, hub *Hub) {
|
||
url := env("RABBITMQ_URL", "amqp://lendry:lendry_password@localhost:5672/")
|
||
for {
|
||
if ctx.Err() != nil {
|
||
return
|
||
}
|
||
|
||
conn, err := amqp091.Dial(url)
|
||
if err != nil {
|
||
log.Printf("RabbitMQ недоступен, повтор через 5 секунд: %v", err)
|
||
time.Sleep(5 * time.Second)
|
||
continue
|
||
}
|
||
|
||
channel, err := conn.Channel()
|
||
if err != nil {
|
||
log.Printf("Не удалось открыть RabbitMQ channel: %v", err)
|
||
_ = conn.Close()
|
||
time.Sleep(5 * time.Second)
|
||
continue
|
||
}
|
||
|
||
queue, err := channel.QueueDeclare("id.notifications", true, false, false, false, nil)
|
||
if err != nil {
|
||
log.Printf("Не удалось объявить очередь уведомлений: %v", err)
|
||
_ = channel.Close()
|
||
_ = conn.Close()
|
||
time.Sleep(5 * time.Second)
|
||
continue
|
||
}
|
||
|
||
deliveries, err := channel.Consume(queue.Name, "media-ws", false, false, false, false, nil)
|
||
if err != nil {
|
||
log.Printf("Не удалось подписаться на уведомления: %v", err)
|
||
_ = channel.Close()
|
||
_ = conn.Close()
|
||
time.Sleep(5 * time.Second)
|
||
continue
|
||
}
|
||
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
_ = channel.Close()
|
||
_ = conn.Close()
|
||
return
|
||
case delivery, ok := <-deliveries:
|
||
if !ok {
|
||
_ = channel.Close()
|
||
_ = conn.Close()
|
||
time.Sleep(2 * time.Second)
|
||
break
|
||
}
|
||
handleDelivery(hub, delivery)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
func handleDelivery(hub *Hub, delivery amqp091.Delivery) {
|
||
var notification Notification
|
||
if err := json.Unmarshal(delivery.Body, ¬ification); err != nil {
|
||
log.Printf("Некорректное уведомление: %v", err)
|
||
_ = delivery.Nack(false, false)
|
||
return
|
||
}
|
||
|
||
if notification.UserID == "" {
|
||
hub.Broadcast(notification)
|
||
} else {
|
||
hub.Send(notification.UserID, notification)
|
||
}
|
||
_ = delivery.Ack(false)
|
||
}
|