mihomo/hub/route/connections.go

92 lines
1.8 KiB
Go

package route
import (
"bytes"
"encoding/json"
"net/http"
"strconv"
"time"
"github.com/Dreamacro/clash/tunnel/statistic"
"github.com/go-chi/chi/v5"
"github.com/go-chi/render"
"github.com/gorilla/websocket"
)
func connectionRouter() http.Handler {
r := chi.NewRouter()
r.Get("/", getConnections)
r.Delete("/", closeAllConnections)
r.Delete("/{id}", closeConnection)
return r
}
func getConnections(w http.ResponseWriter, r *http.Request) {
if !websocket.IsWebSocketUpgrade(r) {
snapshot := statistic.DefaultManager.Snapshot()
render.JSON(w, r, snapshot)
return
}
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
return
}
intervalStr := r.URL.Query().Get("interval")
interval := 1000
if intervalStr != "" {
t, err := strconv.Atoi(intervalStr)
if err != nil {
render.Status(r, http.StatusBadRequest)
render.JSON(w, r, ErrBadRequest)
return
}
interval = t
}
buf := &bytes.Buffer{}
sendSnapshot := func() error {
buf.Reset()
snapshot := statistic.DefaultManager.Snapshot()
if err := json.NewEncoder(buf).Encode(snapshot); err != nil {
return err
}
return conn.WriteMessage(websocket.TextMessage, buf.Bytes())
}
if err := sendSnapshot(); err != nil {
return
}
tick := time.NewTicker(time.Millisecond * time.Duration(interval))
defer tick.Stop()
for range tick.C {
if err := sendSnapshot(); err != nil {
break
}
}
}
func closeConnection(w http.ResponseWriter, r *http.Request) {
id := chi.URLParam(r, "id")
statistic.DefaultManager.ConnectionsRange(func(c statistic.Tracker) bool {
if id == c.ID() {
_ = c.Close()
return false
}
return true
})
render.NoContent(w, r)
}
func closeAllConnections(w http.ResponseWriter, r *http.Request) {
statistic.DefaultManager.ConnectionsRange(func(c statistic.Tracker) bool {
_ = c.Close()
return true
})
render.NoContent(w, r)
}