mirror of
https://github.com/h44z/wg-portal.git
synced 2026-01-29 06:36:24 +00:00
101 lines
2.5 KiB
Go
101 lines
2.5 KiB
Go
package handlers
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
|
|
"github.com/go-pkgz/routegroup"
|
|
"github.com/gorilla/websocket"
|
|
|
|
"github.com/h44z/wg-portal/internal/app"
|
|
"github.com/h44z/wg-portal/internal/config"
|
|
"github.com/h44z/wg-portal/internal/domain"
|
|
)
|
|
|
|
type WebsocketEventBus interface {
|
|
Subscribe(topic string, fn any) error
|
|
Unsubscribe(topic string, fn any) error
|
|
}
|
|
|
|
type WebsocketEndpoint struct {
|
|
authenticator Authenticator
|
|
bus WebsocketEventBus
|
|
|
|
upgrader websocket.Upgrader
|
|
}
|
|
|
|
func NewWebsocketEndpoint(cfg *config.Config, auth Authenticator, bus WebsocketEventBus) *WebsocketEndpoint {
|
|
return &WebsocketEndpoint{
|
|
authenticator: auth,
|
|
bus: bus,
|
|
upgrader: websocket.Upgrader{
|
|
ReadBufferSize: 1024,
|
|
WriteBufferSize: 1024,
|
|
CheckOrigin: func(r *http.Request) bool {
|
|
origin := r.Header.Get("Origin")
|
|
return strings.HasPrefix(origin, cfg.Web.ExternalUrl)
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
func (e WebsocketEndpoint) GetName() string {
|
|
return "WebsocketEndpoint"
|
|
}
|
|
|
|
func (e WebsocketEndpoint) RegisterRoutes(g *routegroup.Bundle) {
|
|
g.With(e.authenticator.LoggedIn()).HandleFunc("GET /ws", e.handleWebsocket())
|
|
}
|
|
|
|
// wsMessage represents a message sent over websocket to the frontend
|
|
type wsMessage struct {
|
|
Type string `json:"type"` // either "peer_stats" or "interface_stats"
|
|
Data any `json:"data"` // domain.TrafficDelta
|
|
}
|
|
|
|
func (e WebsocketEndpoint) handleWebsocket() http.HandlerFunc {
|
|
return func(w http.ResponseWriter, r *http.Request) {
|
|
conn, err := e.upgrader.Upgrade(w, r, nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
defer conn.Close()
|
|
|
|
ctx, cancel := context.WithCancel(r.Context())
|
|
defer cancel()
|
|
|
|
writeMutex := sync.Mutex{}
|
|
writeJSON := func(msg wsMessage) error {
|
|
writeMutex.Lock()
|
|
defer writeMutex.Unlock()
|
|
return conn.WriteJSON(msg)
|
|
}
|
|
|
|
peerStatsHandler := func(status domain.TrafficDelta) {
|
|
_ = writeJSON(wsMessage{Type: "peer_stats", Data: status})
|
|
}
|
|
interfaceStatsHandler := func(status domain.TrafficDelta) {
|
|
_ = writeJSON(wsMessage{Type: "interface_stats", Data: status})
|
|
}
|
|
|
|
_ = e.bus.Subscribe(app.TopicPeerStatsUpdated, peerStatsHandler)
|
|
defer e.bus.Unsubscribe(app.TopicPeerStatsUpdated, peerStatsHandler)
|
|
_ = e.bus.Subscribe(app.TopicInterfaceStatsUpdated, interfaceStatsHandler)
|
|
defer e.bus.Unsubscribe(app.TopicInterfaceStatsUpdated, interfaceStatsHandler)
|
|
|
|
// Keep connection open until client disconnects or context is cancelled
|
|
go func() {
|
|
for {
|
|
if _, _, err := conn.ReadMessage(); err != nil {
|
|
cancel()
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
<-ctx.Done()
|
|
}
|
|
}
|