69 lines
1.4 KiB
Go
69 lines
1.4 KiB
Go
// Filename: internal/handlers/websocket_handler.go
|
|
package handlers
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/gorilla/websocket"
|
|
"github.com/sirupsen/logrus"
|
|
)
|
|
|
|
type connWrapper struct {
|
|
conn *websocket.Conn
|
|
mu sync.Mutex
|
|
}
|
|
|
|
type WebSocketHandler struct {
|
|
logger *logrus.Logger
|
|
clients sync.Map
|
|
upgrader websocket.Upgrader
|
|
}
|
|
|
|
func NewWebSocketHandler(logger *logrus.Logger) *WebSocketHandler {
|
|
return &WebSocketHandler{
|
|
logger: logger,
|
|
upgrader: websocket.Upgrader{
|
|
CheckOrigin: func(r *http.Request) bool { return true },
|
|
},
|
|
}
|
|
}
|
|
|
|
func (h *WebSocketHandler) HandleSystemLogs(c *gin.Context) {
|
|
conn, err := h.upgrader.Upgrade(c.Writer, c.Request, nil)
|
|
if err != nil {
|
|
h.logger.WithError(err).Error("Failed to upgrade websocket")
|
|
return
|
|
}
|
|
defer conn.Close()
|
|
|
|
clientID := time.Now().UnixNano()
|
|
h.clients.Store(clientID, &connWrapper{conn: conn})
|
|
defer h.clients.Delete(clientID)
|
|
|
|
for {
|
|
if _, _, err := conn.ReadMessage(); err != nil {
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
func (h *WebSocketHandler) BroadcastLog(entry *logrus.Entry) {
|
|
msg := map[string]interface{}{
|
|
"timestamp": entry.Time.Format(time.RFC3339),
|
|
"level": entry.Level.String(),
|
|
"message": entry.Message,
|
|
"fields": entry.Data,
|
|
}
|
|
|
|
h.clients.Range(func(key, value interface{}) bool {
|
|
wrapper := value.(*connWrapper)
|
|
wrapper.mu.Lock()
|
|
wrapper.conn.WriteJSON(msg)
|
|
wrapper.mu.Unlock()
|
|
return true
|
|
})
|
|
}
|