From 88eef7617fd7f9dc701e57c544005e2fc9ca8997 Mon Sep 17 00:00:00 2001
From: HynoR <20227709+HynoR@users.noreply.github.com>
Date: Thu, 9 May 2024 10:56:19 +0800
Subject: [PATCH] refactor getOnline feature

---
 core/server/config.go        |  2 +
 core/server/server.go        |  2 +
 extras/trafficlogger/http.go | 74 +++++++++++++++++++++++++++++-------
 3 files changed, 64 insertions(+), 14 deletions(-)

diff --git a/core/server/config.go b/core/server/config.go
index f647f0d..fe64201 100644
--- a/core/server/config.go
+++ b/core/server/config.go
@@ -196,4 +196,6 @@ type EventLogger interface {
 // The implementation of this interface must be thread-safe.
 type TrafficLogger interface {
 	Log(id string, tx, rx uint64) (ok bool)
+	LogOnline(id string, addr net.Addr)
+	LogOffline(id string, addr net.Addr)
 }
diff --git a/core/server/server.go b/core/server/server.go
index 78ab531..99351b0 100644
--- a/core/server/server.go
+++ b/core/server/server.go
@@ -83,6 +83,7 @@ func (s *serverImpl) handleClient(conn quic.Connection) {
 	err := h3s.ServeQUICConn(conn)
 	// If the client is authenticated, we need to log the disconnect event
 	if handler.authenticated && s.config.EventLogger != nil {
+		s.config.TrafficLogger.LogOffline(handler.authID, conn.RemoteAddr())
 		s.config.EventLogger.Disconnect(conn.RemoteAddr(), handler.authID, err)
 	}
 	_ = conn.CloseWithError(closeErrCodeOK, "")
@@ -154,6 +155,7 @@ func (h *h3sHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			w.WriteHeader(protocol.StatusAuthOK)
 			// Call event logger
 			if h.config.EventLogger != nil {
+				h.config.TrafficLogger.LogOnline(id, h.conn.RemoteAddr())
 				h.config.EventLogger.Connect(h.conn.RemoteAddr(), id, actualTx)
 			}
 			// Initialize UDP session manager (if UDP is enabled)
diff --git a/extras/trafficlogger/http.go b/extras/trafficlogger/http.go
index a5a8fc6..5f46e0e 100644
--- a/extras/trafficlogger/http.go
+++ b/extras/trafficlogger/http.go
@@ -2,12 +2,11 @@ package trafficlogger
 
 import (
 	"encoding/json"
+	"github.com/apernet/hysteria/core/server"
+	"net"
 	"net/http"
 	"strconv"
 	"sync"
-	"time"
-
-	"github.com/apernet/hysteria/core/server"
 )
 
 const (
@@ -25,7 +24,7 @@ func NewTrafficStatsServer(secret string) TrafficStatsServer {
 	return &trafficStatsServerImpl{
 		StatsMap:  make(map[string]*trafficStatsEntry),
 		KickMap:   make(map[string]struct{}),
-		OnlineMap: make(map[string]int64),
+		OnlineMap: make(map[string]map[string]bool),
 		Secret:    secret,
 	}
 }
@@ -33,7 +32,7 @@ func NewTrafficStatsServer(secret string) TrafficStatsServer {
 type trafficStatsServerImpl struct {
 	Mutex     sync.RWMutex
 	StatsMap  map[string]*trafficStatsEntry
-	OnlineMap map[string]int64
+	OnlineMap map[string]map[string]bool
 	KickMap   map[string]struct{}
 	Secret    string
 }
@@ -61,11 +60,45 @@ func (s *trafficStatsServerImpl) Log(id string, tx, rx uint64) (ok bool) {
 	entry.Tx += tx
 	entry.Rx += rx
 
-	s.OnlineMap[id] = time.Now().Unix()
-
 	return true
 }
 
+// LogOnline adds the user to the online map.
+func (s *trafficStatsServerImpl) LogOnline(id string, addr net.Addr) {
+	s.Mutex.Lock()
+	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.
+func (s *trafficStatsServerImpl) LogOffline(id string, addr net.Addr) {
+	s.Mutex.Lock()
+	defer s.Mutex.Unlock()
+	userIp, _, err := net.SplitHostPort(addr.String())
+	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)
+		}
+	}
+
+}
+
 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)
@@ -116,17 +149,30 @@ func (s *trafficStatsServerImpl) getOnline(w http.ResponseWriter, r *http.Reques
 	var jb []byte
 	var err error
 
-	timeNow := time.Now().Unix()
+	bClear, _ := strconv.ParseBool(r.URL.Query().Get("clear"))
+	OnlineSet := make(map[string][]string)
 
-	for id, lastSeen := range s.OnlineMap {
-		if timeNow-lastSeen > 180 {
-			delete(s.OnlineMap, id)
+	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()
+		for id, addrs := range s.OnlineMap {
+			for addr := range addrs {
+				OnlineSet[id] = append(OnlineSet[id], addr)
+			}
+		}
+		s.Mutex.RUnlock()
 	}
 
-	s.Mutex.RLock()
-	jb, err = json.Marshal(s.OnlineMap)
-	s.Mutex.RUnlock()
+	jb, err = json.Marshal(OnlineSet)
 
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)