mirror of
https://github.com/serty2005/rmser.git
synced 2026-02-05 03:12:34 -06:00
130 lines
3.2 KiB
Go
130 lines
3.2 KiB
Go
package ws
|
||
|
||
import (
|
||
"net/http"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/gorilla/websocket"
|
||
"go.uber.org/zap"
|
||
"rmser/internal/domain/account"
|
||
"rmser/pkg/logger"
|
||
)
|
||
|
||
type Server struct {
|
||
clients map[string]*websocket.Conn // session_id -> conn
|
||
register chan *clientParams
|
||
unregister chan string
|
||
mu sync.RWMutex
|
||
upgrader websocket.Upgrader
|
||
}
|
||
|
||
type clientParams struct {
|
||
sessionID string
|
||
conn *websocket.Conn
|
||
}
|
||
|
||
func NewServer() *Server {
|
||
return &Server{
|
||
clients: make(map[string]*websocket.Conn),
|
||
register: make(chan *clientParams),
|
||
unregister: make(chan string),
|
||
upgrader: websocket.Upgrader{
|
||
CheckOrigin: func(r *http.Request) bool { return true }, // Allow All CORS
|
||
},
|
||
}
|
||
}
|
||
|
||
func (s *Server) Run() {
|
||
for {
|
||
select {
|
||
case params := <-s.register:
|
||
s.mu.Lock()
|
||
// Если уже есть соединение с таким ID, закрываем старое
|
||
if old, ok := s.clients[params.sessionID]; ok {
|
||
old.Close()
|
||
}
|
||
s.clients[params.sessionID] = params.conn
|
||
s.mu.Unlock()
|
||
logger.Log.Info("WS Client connected", zap.String("session_id", params.sessionID))
|
||
|
||
case sessionID := <-s.unregister:
|
||
s.mu.Lock()
|
||
if conn, ok := s.clients[sessionID]; ok {
|
||
conn.Close()
|
||
delete(s.clients, sessionID)
|
||
}
|
||
s.mu.Unlock()
|
||
logger.Log.Info("WS Client disconnected", zap.String("session_id", sessionID))
|
||
}
|
||
}
|
||
}
|
||
|
||
// HandleConnections обрабатывает входящие WS запросы
|
||
func (s *Server) HandleConnections(c *gin.Context) {
|
||
sessionID := c.Query("session_id")
|
||
if sessionID == "" {
|
||
c.Status(http.StatusBadRequest)
|
||
return
|
||
}
|
||
|
||
conn, err := s.upgrader.Upgrade(c.Writer, c.Request, nil)
|
||
if err != nil {
|
||
logger.Log.Error("WS Upgrade error", zap.Error(err))
|
||
return
|
||
}
|
||
|
||
s.register <- &clientParams{sessionID: sessionID, conn: conn}
|
||
|
||
// Читаем сообщения, чтобы держать соединение открытым и ловить Disconnect
|
||
go func() {
|
||
defer func() {
|
||
s.unregister <- sessionID
|
||
}()
|
||
for {
|
||
if _, _, err := conn.NextReader(); err != nil {
|
||
break
|
||
}
|
||
}
|
||
}()
|
||
}
|
||
|
||
// SendAuthSuccess отправляет токен и закрывает соединение (так как задача выполнена)
|
||
func (s *Server) SendAuthSuccess(sessionID string, token string, user account.User) {
|
||
s.mu.RLock()
|
||
conn, ok := s.clients[sessionID]
|
||
s.mu.RUnlock()
|
||
|
||
if !ok {
|
||
logger.Log.Warn("WS Client not found for auth", zap.String("session_id", sessionID))
|
||
return
|
||
}
|
||
|
||
resp := map[string]interface{}{
|
||
"event": "auth_success",
|
||
"data": map[string]interface{}{
|
||
"token": token,
|
||
"user": map[string]interface{}{
|
||
"id": user.ID,
|
||
"telegram_id": user.TelegramID,
|
||
"username": user.Username,
|
||
"first_name": user.FirstName,
|
||
"last_name": user.LastName,
|
||
"photo_url": user.PhotoURL,
|
||
"is_system_admin": user.IsSystemAdmin,
|
||
},
|
||
},
|
||
}
|
||
|
||
if err := conn.WriteJSON(resp); err != nil {
|
||
logger.Log.Error("WS Write error", zap.Error(err))
|
||
} else {
|
||
logger.Log.Info("WS Auth sent successfully", zap.String("session_id", sessionID))
|
||
}
|
||
|
||
// Даем время на доставку и закрываем
|
||
time.Sleep(500 * time.Millisecond)
|
||
s.unregister <- sessionID
|
||
}
|