package trafficlogger import ( "encoding/json" "net/http" "strconv" "sync" "github.com/apernet/hysteria/core/v2/server" ) const ( indexHTML = ` Hysteria Traffic Stats API Server

This is a Hysteria Traffic Stats API server.

Check the documentation for usage.

` ) // TrafficStatsServer implements both server.TrafficLogger and http.Handler // to provide a simple HTTP API to get the traffic stats per user. type TrafficStatsServer interface { server.TrafficLogger http.Handler } func NewTrafficStatsServer(secret string) TrafficStatsServer { return &trafficStatsServerImpl{ StatsMap: make(map[string]*trafficStatsEntry), KickMap: make(map[string]struct{}), OnlineMap: make(map[string]int), Secret: secret, } } type trafficStatsServerImpl struct { Mutex sync.RWMutex StatsMap map[string]*trafficStatsEntry OnlineMap map[string]int KickMap map[string]struct{} Secret string } type trafficStatsEntry struct { Tx uint64 `json:"tx"` Rx uint64 `json:"rx"` } func (s *trafficStatsServerImpl) LogTraffic(id string, tx, rx uint64) (ok bool) { s.Mutex.Lock() defer s.Mutex.Unlock() _, ok = s.KickMap[id] if ok { delete(s.KickMap, id) return false } entry, ok := s.StatsMap[id] if !ok { entry = &trafficStatsEntry{} s.StatsMap[id] = entry } entry.Tx += tx entry.Rx += rx return true } // LogOnlineStateChanged updates the online state to the online map. func (s *trafficStatsServerImpl) LogOnlineState(id string, online bool) { s.Mutex.Lock() defer s.Mutex.Unlock() if online { s.OnlineMap[id]++ } else { s.OnlineMap[id]-- if s.OnlineMap[id] <= 0 { delete(s.OnlineMap, id) } } } func (s *trafficStatsServerImpl) ServeHTTP(w http.ResponseWriter, r *http.Request) { if s.Secret != "" && r.Header.Get("Authorization") != s.Secret { http.Error(w, "unauthorized", http.StatusUnauthorized) return } if r.Method == http.MethodGet && r.URL.Path == "/" { _, _ = w.Write([]byte(indexHTML)) return } if r.Method == http.MethodGet && r.URL.Path == "/traffic" { s.getTraffic(w, r) return } if r.Method == http.MethodPost && r.URL.Path == "/kick" { s.kick(w, r) return } if r.Method == http.MethodGet && r.URL.Path == "/online" { s.getOnline(w, r) return } http.NotFound(w, r) } func (s *trafficStatsServerImpl) getTraffic(w http.ResponseWriter, r *http.Request) { bClear, _ := strconv.ParseBool(r.URL.Query().Get("clear")) var jb []byte var err error if bClear { s.Mutex.Lock() jb, err = json.Marshal(s.StatsMap) s.StatsMap = make(map[string]*trafficStatsEntry) s.Mutex.Unlock() } else { s.Mutex.RLock() jb, err = json.Marshal(s.StatsMap) s.Mutex.RUnlock() } if err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) return } w.Header().Set("Content-Type", "application/json; charset=utf-8") _, _ = w.Write(jb) } func (s *trafficStatsServerImpl) getOnline(w http.ResponseWriter, r *http.Request) { s.Mutex.RLock() defer s.Mutex.RUnlock() jb, err := json.Marshal(s.OnlineMap) if err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) return } w.Header().Set("Content-Type", "application/json; charset=utf-8") _, _ = w.Write(jb) } func (s *trafficStatsServerImpl) kick(w http.ResponseWriter, r *http.Request) { var ids []string err := json.NewDecoder(r.Body).Decode(&ids) if err != nil { http.Error(w, err.Error(), http.StatusBadRequest) return } s.Mutex.Lock() for _, id := range ids { s.KickMap[id] = struct{}{} } s.Mutex.Unlock() w.WriteHeader(http.StatusOK) }