feat(server): add traffic storage and statistics tracking for improved traffic management

cmd/server/main.go

@@ -84,6 +84,7 @@ func main() {
 	registry := plugin.NewRegistry()
 	server.SetPluginRegistry(registry)
 	server.SetJSPluginStore(clientStore) // 设置 JS 插件存储，用于客户端重连时恢复插件
+	server.SetTrafficStore(clientStore)  // 设置流量存储，用于记录流量统计
 
 	// 启动 Web 控制台
 	if cfg.Server.Web.Enabled {

internal/server/db/sqlite.go

@@ -397,14 +397,19 @@ func (s *SQLiteStore) Get24HourTraffic() (inbound, outbound int64, err error) {
 	return
 }
 
-// GetHourlyTraffic 获取每小时流量记录
+// GetHourlyTraffic 获取每小时流量记录（始终返回完整的 hours 小时数据）
 func (s *SQLiteStore) GetHourlyTraffic(hours int) ([]TrafficRecord, error) {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
 
-	cutoff := time.Now().Add(-time.Duration(hours) * time.Hour).Unix()
+	// 计算当前小时的起始时间戳
+	now := time.Now()
+	currentHour := time.Date(now.Year(), now.Month(), now.Day(), now.Hour(), 0, 0, 0, now.Location())
+
+	// 查询数据库中的记录
+	cutoff := currentHour.Add(-time.Duration(hours-1) * time.Hour).Unix()
 	rows, err := s.db.Query(`
-		SELECT hour_ts, inbound, outbound FROM traffic_stats 
+		SELECT hour_ts, inbound, outbound FROM traffic_stats
 		WHERE hour_ts >= ? ORDER BY hour_ts ASC
 	`, cutoff)
 	if err != nil {
@@ -412,13 +417,26 @@ func (s *SQLiteStore) GetHourlyTraffic(hours int) ([]TrafficRecord, error) {
 	}
 	defer rows.Close()
 
-	var records []TrafficRecord
+	// 将数据库记录放入 map 以便快速查找
+	dbRecords := make(map[int64]TrafficRecord)
 	for rows.Next() {
 		var r TrafficRecord
 		if err := rows.Scan(&r.Timestamp, &r.Inbound, &r.Outbound); err != nil {
 			return nil, err
 		}
-		records = append(records, r)
+		dbRecords[r.Timestamp] = r
 	}
+
+	// 生成完整的 hours 小时数据
+	records := make([]TrafficRecord, hours)
+	for i := 0; i < hours; i++ {
+		ts := currentHour.Add(-time.Duration(hours-1-i) * time.Hour).Unix()
+		if r, ok := dbRecords[ts]; ok {
+			records[i] = r
+		} else {
+			records[i] = TrafficRecord{Timestamp: ts, Inbound: 0, Outbound: 0}
+		}
+	}
+
 	return records, nil
 }

internal/server/tunnel/server.go

@@ -50,7 +50,8 @@ func generateClientID() string {
 // Server 隧道服务端
 type Server struct {
 	clientStore    db.ClientStore
-	jsPluginStore  db.JSPluginStore // JS 插件存储
+	jsPluginStore  db.JSPluginStore  // JS 插件存储
+	trafficStore   db.TrafficStore   // 流量存储
 	bindAddr       string
 	bindPort       int
 	token          string
@@ -161,6 +162,11 @@ func (s *Server) SetJSPluginStore(store db.JSPluginStore) {
 	s.jsPluginStore = store
 }
 
+// SetTrafficStore 设置流量存储
+func (s *Server) SetTrafficStore(store db.TrafficStore) {
+	s.trafficStore = store
+}
+
 // LoadJSPlugins 加载 JS 插件配置
 func (s *Server) LoadJSPlugins(plugins []JSPluginEntry) {
 	s.jsPlugins = plugins
@@ -536,7 +542,7 @@ func (s *Server) handleProxyConn(cs *ClientSession, conn net.Conn, rule protocol
 		return
 	}
 
-	relay.Relay(conn, stream)
+	relay.RelayWithStats(conn, stream, s.recordTraffic)
 }
 
 // heartbeatLoop 心跳检测循环
@@ -1224,7 +1230,7 @@ func (s *Server) handleClientPluginConn(cs *ClientSession, conn net.Conn, rule p
 		}
 	}
 
-	relay.Relay(conn, stream)
+	relay.RelayWithStats(conn, stream, s.recordTraffic)
 }
 
 // checkHTTPBasicAuth 检查 HTTP Basic Auth
@@ -1907,6 +1913,16 @@ func (s *Server) StopClientLogStream(sessionID string) {
 	s.logSessions.RemoveSession(sessionID)
 }
 
+// recordTraffic 记录流量统计
+func (s *Server) recordTraffic(inbound, outbound int64) {
+	if s.trafficStore == nil {
+		return
+	}
+	if err := s.trafficStore.AddTraffic(inbound, outbound); err != nil {
+		log.Printf("[Server] Record traffic error: %v", err)
+	}
+}
+
 // boolPtr 返回 bool 值的指针
 func boolPtr(b bool) *bool {
 	return &b

internal/server/tunnel/websocket.go

@@ -142,5 +142,5 @@ func (s *Server) handleWebsocketProxyConn(cs *ClientSession, conn net.Conn, rule
 		return
 	}
 
-	relay.Relay(conn, stream)
+	relay.RelayWithStats(conn, stream, s.recordTraffic)
 }