github.com/metacubex/mihomo@v1.18.5/hub/route/connections.go (about)

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