refactor(online): track count instead of raddr

This commit is contained in:
Haruue 2024-05-11 11:05:32 +08:00
parent 88eef7617f
commit ba9b3cdebb
No known key found for this signature in database
GPG Key ID: F6083B28CBCBC148
3 changed files with 17 additions and 63 deletions

View File

@ -196,6 +196,5 @@ type EventLogger interface {
// The implementation of this interface must be thread-safe. // The implementation of this interface must be thread-safe.
type TrafficLogger interface { type TrafficLogger interface {
Log(id string, tx, rx uint64) (ok bool) Log(id string, tx, rx uint64) (ok bool)
LogOnline(id string, addr net.Addr) LogOnlineStateChanged(id string, online bool)
LogOffline(id string, addr net.Addr)
} }

View File

@ -83,7 +83,7 @@ func (s *serverImpl) handleClient(conn quic.Connection) {
err := h3s.ServeQUICConn(conn) err := h3s.ServeQUICConn(conn)
// If the client is authenticated, we need to log the disconnect event // If the client is authenticated, we need to log the disconnect event
if handler.authenticated && s.config.EventLogger != nil { if handler.authenticated && s.config.EventLogger != nil {
s.config.TrafficLogger.LogOffline(handler.authID, conn.RemoteAddr()) s.config.TrafficLogger.LogOnlineStateChanged(handler.authID, false)
s.config.EventLogger.Disconnect(conn.RemoteAddr(), handler.authID, err) s.config.EventLogger.Disconnect(conn.RemoteAddr(), handler.authID, err)
} }
_ = conn.CloseWithError(closeErrCodeOK, "") _ = conn.CloseWithError(closeErrCodeOK, "")
@ -155,7 +155,7 @@ func (h *h3sHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(protocol.StatusAuthOK) w.WriteHeader(protocol.StatusAuthOK)
// Call event logger // Call event logger
if h.config.EventLogger != nil { if h.config.EventLogger != nil {
h.config.TrafficLogger.LogOnline(id, h.conn.RemoteAddr()) h.config.TrafficLogger.LogOnlineStateChanged(id, true)
h.config.EventLogger.Connect(h.conn.RemoteAddr(), id, actualTx) h.config.EventLogger.Connect(h.conn.RemoteAddr(), id, actualTx)
} }
// Initialize UDP session manager (if UDP is enabled) // Initialize UDP session manager (if UDP is enabled)

View File

@ -2,11 +2,11 @@ package trafficlogger
import ( import (
"encoding/json" "encoding/json"
"github.com/apernet/hysteria/core/server"
"net"
"net/http" "net/http"
"strconv" "strconv"
"sync" "sync"
"github.com/apernet/hysteria/core/server"
) )
const ( const (
@ -24,7 +24,7 @@ func NewTrafficStatsServer(secret string) TrafficStatsServer {
return &trafficStatsServerImpl{ return &trafficStatsServerImpl{
StatsMap: make(map[string]*trafficStatsEntry), StatsMap: make(map[string]*trafficStatsEntry),
KickMap: make(map[string]struct{}), KickMap: make(map[string]struct{}),
OnlineMap: make(map[string]map[string]bool), OnlineMap: make(map[string]int),
Secret: secret, Secret: secret,
} }
} }
@ -32,7 +32,7 @@ func NewTrafficStatsServer(secret string) TrafficStatsServer {
type trafficStatsServerImpl struct { type trafficStatsServerImpl struct {
Mutex sync.RWMutex Mutex sync.RWMutex
StatsMap map[string]*trafficStatsEntry StatsMap map[string]*trafficStatsEntry
OnlineMap map[string]map[string]bool OnlineMap map[string]int
KickMap map[string]struct{} KickMap map[string]struct{}
Secret string Secret string
} }
@ -63,40 +63,19 @@ func (s *trafficStatsServerImpl) Log(id string, tx, rx uint64) (ok bool) {
return true return true
} }
// LogOnline adds the user to the online map. // LogOnlineStateChanged updates the online state to the online map.
func (s *trafficStatsServerImpl) LogOnline(id string, addr net.Addr) { func (s *trafficStatsServerImpl) LogOnlineStateChanged(id string, online bool) {
s.Mutex.Lock() s.Mutex.Lock()
defer s.Mutex.Unlock() defer s.Mutex.Unlock()
if _, ok := s.OnlineMap[id]; !ok {
s.OnlineMap[id] = make(map[string]bool)
}
userIp, _, err := net.SplitHostPort(addr.String())
if err != nil {
return
}
s.OnlineMap[id][userIp] = true
}
// LogOffline removes the user from the online map. if online {
func (s *trafficStatsServerImpl) LogOffline(id string, addr net.Addr) { s.OnlineMap[id]++
s.Mutex.Lock() } else {
defer s.Mutex.Unlock() s.OnlineMap[id]--
userIp, _, err := net.SplitHostPort(addr.String()) if s.OnlineMap[id] <= 0 {
if err != nil {
return
}
if onlineUsers, ok := s.OnlineMap[id]; ok {
if !onlineUsers[userIp] {
//if the user's ip is not in the online map, delete the whole entry
delete(s.OnlineMap, id)
return
}
delete(onlineUsers, userIp)
if len(onlineUsers) == 0 {
delete(s.OnlineMap, id) delete(s.OnlineMap, id)
} }
} }
} }
func (s *trafficStatsServerImpl) ServeHTTP(w http.ResponseWriter, r *http.Request) { func (s *trafficStatsServerImpl) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@ -146,34 +125,10 @@ func (s *trafficStatsServerImpl) getTraffic(w http.ResponseWriter, r *http.Reque
} }
func (s *trafficStatsServerImpl) getOnline(w http.ResponseWriter, r *http.Request) { func (s *trafficStatsServerImpl) getOnline(w http.ResponseWriter, r *http.Request) {
var jb []byte
var err error
bClear, _ := strconv.ParseBool(r.URL.Query().Get("clear"))
OnlineSet := make(map[string][]string)
if bClear {
s.Mutex.Lock()
for id, addrs := range s.OnlineMap {
for addr := range addrs {
OnlineSet[id] = append(OnlineSet[id], addr)
}
}
s.OnlineMap = make(map[string]map[string]bool)
s.Mutex.Unlock()
} else {
s.Mutex.RLock() s.Mutex.RLock()
for id, addrs := range s.OnlineMap { defer s.Mutex.RUnlock()
for addr := range addrs {
OnlineSet[id] = append(OnlineSet[id], addr)
}
}
s.Mutex.RUnlock()
}
jb, err = json.Marshal(OnlineSet)
jb, err := json.Marshal(s.OnlineMap)
if err != nil { if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError) http.Error(w, err.Error(), http.StatusInternalServerError)
return return