connections.go 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889
  1. package route
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "net/http"
  6. "strconv"
  7. "time"
  8. "github.com/metacubex/mihomo/tunnel/statistic"
  9. "github.com/go-chi/chi/v5"
  10. "github.com/go-chi/render"
  11. "github.com/gobwas/ws"
  12. "github.com/gobwas/ws/wsutil"
  13. )
  14. func connectionRouter() http.Handler {
  15. r := chi.NewRouter()
  16. r.Get("/", getConnections)
  17. r.Delete("/", closeAllConnections)
  18. r.Delete("/{id}", closeConnection)
  19. return r
  20. }
  21. func getConnections(w http.ResponseWriter, r *http.Request) {
  22. if !(r.Header.Get("Upgrade") == "websocket") {
  23. snapshot := statistic.DefaultManager.Snapshot()
  24. render.JSON(w, r, snapshot)
  25. return
  26. }
  27. conn, _, _, err := ws.UpgradeHTTP(r, w)
  28. if err != nil {
  29. return
  30. }
  31. intervalStr := r.URL.Query().Get("interval")
  32. interval := 1000
  33. if intervalStr != "" {
  34. t, err := strconv.Atoi(intervalStr)
  35. if err != nil {
  36. render.Status(r, http.StatusBadRequest)
  37. render.JSON(w, r, ErrBadRequest)
  38. return
  39. }
  40. interval = t
  41. }
  42. buf := &bytes.Buffer{}
  43. sendSnapshot := func() error {
  44. buf.Reset()
  45. snapshot := statistic.DefaultManager.Snapshot()
  46. if err := json.NewEncoder(buf).Encode(snapshot); err != nil {
  47. return err
  48. }
  49. return wsutil.WriteMessage(conn, ws.StateServerSide, ws.OpText, buf.Bytes())
  50. }
  51. if err := sendSnapshot(); err != nil {
  52. return
  53. }
  54. tick := time.NewTicker(time.Millisecond * time.Duration(interval))
  55. defer tick.Stop()
  56. for range tick.C {
  57. if err := sendSnapshot(); err != nil {
  58. break
  59. }
  60. }
  61. }
  62. func closeConnection(w http.ResponseWriter, r *http.Request) {
  63. id := chi.URLParam(r, "id")
  64. if c := statistic.DefaultManager.Get(id); c != nil {
  65. _ = c.Close()
  66. }
  67. render.NoContent(w, r)
  68. }
  69. func closeAllConnections(w http.ResponseWriter, r *http.Request) {
  70. statistic.DefaultManager.Range(func(c statistic.Tracker) bool {
  71. _ = c.Close()
  72. return true
  73. })
  74. render.NoContent(w, r)
  75. }