增加websocket,数据实时显示

go.mod

@@ -5,6 +5,7 @@ go 1.23.0
 toolchain go1.24.10
 
 require (
+	github.com/gorilla/websocket v1.5.1
 	github.com/miekg/dns v1.1.68
 	github.com/sirupsen/logrus v1.9.3
 )

go.sum

@@ -3,6 +3,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
+github.com/gorilla/websocket v1.5.1 h1:gmztn0JnHVt9JZquRuzLw3g4wouNVzKL15iLr/zn/QY=
+github.com/gorilla/websocket v1.5.1/go.mod h1:x3kM2JMyaluk02fnUJpQuwD2dCS5NDG2ZHL0uE0tcaY=
 github.com/miekg/dns v1.1.68 h1:jsSRkNozw7G/mnmXULynzMNIsgY2dHC8LO6U6Ij2JEA=
 github.com/miekg/dns v1.1.68/go.mod h1:fujopn7TB3Pu3JM69XaawiU0wqjpL9/8xGop5UrTPps=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=

http/server.go

@@ -7,8 +7,10 @@ import (
 	"net/http"
 	"sort"
 	"strings"
+	"sync"
 	"time"
 
+	"github.com/gorilla/websocket"
 	"dns-server/config"
 	"dns-server/dns"
 	"dns-server/logger"
@@ -22,16 +24,37 @@ type Server struct {
 	dnsServer     *dns.Server
 	shieldManager *shield.ShieldManager
 	server        *http.Server
+	
+	// WebSocket相关字段
+	upgrader      websocket.Upgrader
+	clients       map[*websocket.Conn]bool
+	clientsMutex  sync.Mutex
+	broadcastChan chan []byte
 }
 
 // NewServer 创建HTTP服务器实例
 func NewServer(globalConfig *config.Config, dnsServer *dns.Server, shieldManager *shield.ShieldManager) *Server {
-	return &Server{
+	server := &Server{
 		globalConfig:  globalConfig,
 		config:        &globalConfig.HTTP,
 		dnsServer:     dnsServer,
 		shieldManager: shieldManager,
+		upgrader: websocket.Upgrader{
+			ReadBufferSize:  1024,
+			WriteBufferSize: 1024,
+			// 允许所有CORS请求
+			CheckOrigin: func(r *http.Request) bool {
+				return true
+			},
+		},
+		clients:       make(map[*websocket.Conn]bool),
+		broadcastChan: make(chan []byte, 100),
 	}
+	
+	// 启动广播协程
+	go server.startBroadcastLoop()
+	
+	return server
 }
 
 // Start 启动HTTP服务器
@@ -55,6 +78,8 @@ func (s *Server) Start() error {
 		mux.HandleFunc("/api/daily-stats", s.handleDailyStats)
 		mux.HandleFunc("/api/monthly-stats", s.handleMonthlyStats)
 		mux.HandleFunc("/api/query/type", s.handleQueryTypeStats)
+		// WebSocket端点
+		mux.HandleFunc("/ws/stats", s.handleWebSocketStats)
 	}
 
 	// 静态文件服务（可后续添加前端界面）
@@ -133,6 +158,174 @@ func (s *Server) handleStats(w http.ResponseWriter, r *http.Request) {
 	json.NewEncoder(w).Encode(stats)
 }
 
+// WebSocket相关方法
+
+// handleWebSocketStats 处理WebSocket连接，用于实时推送统计数据
+func (s *Server) handleWebSocketStats(w http.ResponseWriter, r *http.Request) {
+	// 升级HTTP连接为WebSocket
+	conn, err := s.upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		logger.Error(fmt.Sprintf("WebSocket升级失败: %v", err))
+		return
+	}
+	defer conn.Close()
+
+	// 将新客户端添加到客户端列表
+	s.clientsMutex.Lock()
+	s.clients[conn] = true
+	clientCount := len(s.clients)
+	s.clientsMutex.Unlock()
+
+	logger.Info(fmt.Sprintf("新WebSocket客户端连接，当前连接数: %d", clientCount))
+
+	// 发送初始数据
+	if err := s.sendInitialStats(conn); err != nil {
+		logger.Error(fmt.Sprintf("发送初始数据失败: %v", err))
+		return
+	}
+
+	// 定期发送更新数据
+	ticker := time.NewTicker(500 * time.Millisecond) // 每500ms检查一次数据变化
+	defer ticker.Stop()
+
+	// 最后一次发送的数据快照，用于检测变化
+	var lastStats map[string]interface{}
+
+	// 保持连接并定期发送数据
+	for {
+		select {
+		case <-ticker.C:
+			// 获取最新统计数据
+			currentStats := s.buildStatsData()
+			
+			// 检查数据是否有变化
+			if !s.areStatsEqual(lastStats, currentStats) {
+				// 数据有变化，发送更新
+				data, err := json.Marshal(map[string]interface{}{
+					"type":  "stats_update",
+					"data":  currentStats,
+					"time":  time.Now(),
+				})
+				if err != nil {
+					logger.Error(fmt.Sprintf("序列化统计数据失败: %v", err))
+					continue
+				}
+				
+				if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
+					logger.Error(fmt.Sprintf("发送WebSocket消息失败: %v", err))
+					return
+				}
+				
+				// 更新最后发送的数据
+				lastStats = currentStats
+			}
+		case <-r.Context().Done():
+			// 客户端断开连接
+			s.clientsMutex.Lock()
+			delete(s.clients, conn)
+			clientCount := len(s.clients)
+			s.clientsMutex.Unlock()
+			logger.Info(fmt.Sprintf("WebSocket客户端断开连接，当前连接数: %d", clientCount))
+			return
+		}
+	}
+}
+
+// sendInitialStats 发送初始统计数据
+func (s *Server) sendInitialStats(conn *websocket.Conn) error {
+	stats := s.buildStatsData()
+	data, err := json.Marshal(map[string]interface{}{
+		"type":  "initial_data",
+		"data":  stats,
+		"time":  time.Now(),
+	})
+	if err != nil {
+		return err
+	}
+	return conn.WriteMessage(websocket.TextMessage, data)
+}
+
+// buildStatsData 构建统计数据
+func (s *Server) buildStatsData() map[string]interface{} {
+	dnsStats := s.dnsServer.GetStats()
+	shieldStats := s.shieldManager.GetStats()
+
+	// 获取最常用查询类型
+	topQueryType := "-"
+	maxCount := int64(0)
+	if len(dnsStats.QueryTypes) > 0 {
+		for queryType, count := range dnsStats.QueryTypes {
+			if count > maxCount {
+				maxCount = count
+				topQueryType = queryType
+			}
+		}
+	}
+
+	// 获取活跃来源IP数量
+	activeIPCount := len(dnsStats.SourceIPs)
+
+	// 格式化平均响应时间
+	formattedResponseTime := float64(int(dnsStats.AvgResponseTime*100)) / 100
+	
+	return map[string]interface{}{
+		"dns": map[string]interface{}{
+			"Queries":         dnsStats.Queries,
+			"Blocked":         dnsStats.Blocked,
+			"Allowed":         dnsStats.Allowed,
+			"Errors":          dnsStats.Errors,
+			"LastQuery":       dnsStats.LastQuery,
+			"AvgResponseTime": formattedResponseTime,
+			"TotalResponseTime": dnsStats.TotalResponseTime,
+			"QueryTypes":      dnsStats.QueryTypes,
+			"SourceIPs":       dnsStats.SourceIPs,
+			"CpuUsage":        dnsStats.CpuUsage,
+		},
+		"shield":        shieldStats,
+		"topQueryType":  topQueryType,
+		"activeIPs":     activeIPCount,
+		"avgResponseTime": formattedResponseTime,
+		"cpuUsage":      dnsStats.CpuUsage,
+	}
+}
+
+// areStatsEqual 检查两次统计数据是否相等（用于检测变化）
+func (s *Server) areStatsEqual(stats1, stats2 map[string]interface{}) bool {
+	if stats1 == nil || stats2 == nil {
+		return false
+	}
+	
+	// 只比较关键数值，避免频繁更新
+	if dns1, ok1 := stats1["dns"].(map[string]interface{}); ok1 {
+		if dns2, ok2 := stats2["dns"].(map[string]interface{}); ok2 {
+			// 检查主要计数器
+			if dns1["Queries"] != dns2["Queries"] ||
+			   dns1["Blocked"] != dns2["Blocked"] ||
+			   dns1["Allowed"] != dns2["Allowed"] ||
+			   dns1["Errors"] != dns2["Errors"] {
+				return false
+			}
+		}
+	}
+	
+	return true
+}
+
+// startBroadcastLoop 启动广播循环
+func (s *Server) startBroadcastLoop() {
+	for message := range s.broadcastChan {
+		s.clientsMutex.Lock()
+		for client := range s.clients {
+			if err := client.WriteMessage(websocket.TextMessage, message); err != nil {
+				logger.Error(fmt.Sprintf("广播消息失败: %v", err))
+				client.Close()
+				delete(s.clients, client)
+			}
+		}
+		s.clientsMutex.Unlock()
+	}
+}
+
 // handleTopBlockedDomains 处理TOP屏蔽域名请求
 func (s *Server) handleTopBlockedDomains(w http.ResponseWriter, r *http.Request) {
 	if r.Method != http.MethodGet {

static/js/dashboard.js

@@ -5,6 +5,8 @@ let ratioChart = null;
 let dnsRequestsChart = null;
 let queryTypeChart = null; // 解析类型统计饼图
 let intervalId = null;
+let wsConnection = null;
+let wsReconnectTimer = null;
 // 存储统计卡片图表实例
 let statCardCharts = {};
 // 存储统计卡片历史数据
@@ -30,14 +32,260 @@ async function initDashboard() {
         // 初始化时间范围切换
         initTimeRangeToggle();
         
-        // 设置定时更新
+        // 建立WebSocket连接
-        intervalId = setInterval(loadDashboardData, 5000); // 每5秒更新一次
+        connectWebSocket();
+        
+        // 在页面卸载时清理资源
+        window.addEventListener('beforeunload', cleanupResources);
     } catch (error) {
         console.error('初始化仪表盘失败:', error);
         showNotification('初始化失败: ' + error.message, 'error');
     }
 }
 
+// 建立WebSocket连接
+function connectWebSocket() {
+    try {
+        // 构建WebSocket URL
+        const wsProtocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
+        const wsUrl = `${wsProtocol}//${window.location.host}/ws/stats`;
+        
+        console.log('正在连接WebSocket:', wsUrl);
+        
+        // 创建WebSocket连接
+        wsConnection = new WebSocket(wsUrl);
+        
+        // 连接打开事件
+        wsConnection.onopen = function() {
+            console.log('WebSocket连接已建立');
+            showNotification('实时数据更新已连接', 'success');
+            
+            // 清除重连计时器
+            if (wsReconnectTimer) {
+                clearTimeout(wsReconnectTimer);
+                wsReconnectTimer = null;
+            }
+        };
+        
+        // 接收消息事件
+        wsConnection.onmessage = function(event) {
+            try {
+                const data = JSON.parse(event.data);
+                
+                if (data.type === 'initial_data' || data.type === 'stats_update') {
+                    console.log('收到实时数据更新');
+                    processRealTimeData(data.data);
+                }
+            } catch (error) {
+                console.error('处理WebSocket消息失败:', error);
+            }
+        };
+        
+        // 连接关闭事件
+        wsConnection.onclose = function(event) {
+            console.warn('WebSocket连接已关闭，代码:', event.code);
+            wsConnection = null;
+            
+            // 设置重连
+            setupReconnect();
+        };
+        
+        // 连接错误事件
+        wsConnection.onerror = function(error) {
+            console.error('WebSocket连接错误:', error);
+        };
+        
+    } catch (error) {
+        console.error('创建WebSocket连接失败:', error);
+        // 如果WebSocket连接失败，回退到定时刷新
+        fallbackToIntervalRefresh();
+    }
+}
+
+// 设置重连逻辑
+function setupReconnect() {
+    if (wsReconnectTimer) {
+        return; // 已经有重连计时器在运行
+    }
+    
+    const reconnectDelay = 5000; // 5秒后重连
+    console.log(`将在${reconnectDelay}ms后尝试重新连接WebSocket`);
+    
+    wsReconnectTimer = setTimeout(() => {
+        connectWebSocket();
+    }, reconnectDelay);
+}
+
+// 处理实时数据更新
+function processRealTimeData(stats) {
+    try {
+        // 更新统计卡片
+        updateStatsCards(stats);
+        
+        // 获取查询类型统计数据
+        let queryTypeStats = null;
+        if (stats.dns && stats.dns.QueryTypes) {
+            queryTypeStats = Object.entries(stats.dns.QueryTypes).map(([type, count]) => ({
+                type,
+                count
+            }));
+        }
+        
+        // 更新图表数据
+        updateCharts(stats, queryTypeStats);
+        
+        // 更新卡片图表
+        updateStatCardCharts(stats);
+        
+        // 尝试从stats中获取总查询数等信息
+        if (stats.dns) {
+            totalQueries = stats.dns.Allowed + stats.dns.Blocked + (stats.dns.Errors || 0);
+            blockedQueries = stats.dns.Blocked;
+            errorQueries = stats.dns.Errors || 0;
+            allowedQueries = stats.dns.Allowed;
+        } else {
+            totalQueries = stats.totalQueries || 0;
+            blockedQueries = stats.blockedQueries || 0;
+            errorQueries = stats.errorQueries || 0;
+            allowedQueries = stats.allowedQueries || 0;
+        }
+        
+        // 更新新卡片数据
+        if (document.getElementById('avg-response-time')) {
+            const responseTime = stats.avgResponseTime ? stats.avgResponseTime.toFixed(2) + 'ms' : '---';
+            
+            // 计算响应时间趋势
+            let responsePercent = '---';
+            let trendClass = 'text-gray-400';
+            let trendIcon = '---';
+            
+            if (stats.avgResponseTime !== undefined && stats.avgResponseTime !== null) {
+                // 存储当前值用于下次计算趋势
+                const prevResponseTime = window.dashboardHistoryData.prevResponseTime || stats.avgResponseTime;
+                window.dashboardHistoryData.prevResponseTime = stats.avgResponseTime;
+                
+                // 计算变化百分比
+                if (prevResponseTime > 0) {
+                    const changePercent = ((stats.avgResponseTime - prevResponseTime) / prevResponseTime) * 100;
+                    responsePercent = Math.abs(changePercent).toFixed(1) + '%';
+                    
+                    // 设置趋势图标和颜色
+                    if (changePercent > 0) {
+                        trendIcon = '↓';
+                        trendClass = 'text-danger';
+                    } else if (changePercent < 0) {
+                        trendIcon = '↑';
+                        trendClass = 'text-success';
+                    } else {
+                        trendIcon = '•';
+                        trendClass = 'text-gray-500';
+                    }
+                }
+            }
+            
+            document.getElementById('avg-response-time').textContent = responseTime;
+            const responseTimePercentElem = document.getElementById('response-time-percent');
+            if (responseTimePercentElem) {
+                responseTimePercentElem.textContent = trendIcon + ' ' + responsePercent;
+                responseTimePercentElem.className = `text-sm flex items-center ${trendClass}`;
+            }
+        }
+        
+        if (document.getElementById('top-query-type')) {
+            const queryType = stats.topQueryType || '---';
+            
+            const queryPercentElem = document.getElementById('query-type-percentage');
+            if (queryPercentElem) {
+                queryPercentElem.textContent = '• ---';
+                queryPercentElem.className = 'text-sm flex items-center text-gray-500';
+            }
+            
+            document.getElementById('top-query-type').textContent = queryType;
+        }
+        
+        if (document.getElementById('active-ips')) {
+            const activeIPs = stats.activeIPs !== undefined ? formatNumber(stats.activeIPs) : '---';
+            
+            // 计算活跃IP趋势
+            let ipsPercent = '---';
+            let trendClass = 'text-gray-400';
+            let trendIcon = '---';
+            
+            if (stats.activeIPs !== undefined) {
+                const prevActiveIPs = window.dashboardHistoryData.prevActiveIPs || stats.activeIPs;
+                window.dashboardHistoryData.prevActiveIPs = stats.activeIPs;
+                
+                if (prevActiveIPs > 0) {
+                    const changePercent = ((stats.activeIPs - prevActiveIPs) / prevActiveIPs) * 100;
+                    ipsPercent = Math.abs(changePercent).toFixed(1) + '%';
+                    
+                    if (changePercent > 0) {
+                        trendIcon = '↑';
+                        trendClass = 'text-primary';
+                    } else if (changePercent < 0) {
+                        trendIcon = '↓';
+                        trendClass = 'text-secondary';
+                    } else {
+                        trendIcon = '•';
+                        trendClass = 'text-gray-500';
+                    }
+                }
+            }
+            
+            document.getElementById('active-ips').textContent = activeIPs;
+            const activeIpsPercentElem = document.getElementById('active-ips-percentage');
+            if (activeIpsPercentElem) {
+                activeIpsPercentElem.textContent = trendIcon + ' ' + ipsPercent;
+                activeIpsPercentElem.className = `text-sm flex items-center ${trendClass}`;
+            }
+        }
+        
+    } catch (error) {
+        console.error('处理实时数据失败:', error);
+    }
+}
+
+// 回退到定时刷新
+function fallbackToIntervalRefresh() {
+    console.warn('回退到定时刷新模式');
+    showNotification('实时更新连接失败，已切换到定时刷新模式', 'warning');
+    
+    // 如果已经有定时器，先清除
+    if (intervalId) {
+        clearInterval(intervalId);
+    }
+    
+    // 设置新的定时器
+    intervalId = setInterval(async () => {
+        try {
+            await loadDashboardData();
+        } catch (error) {
+            console.error('定时刷新失败:', error);
+        }
+    }, 5000); // 每5秒更新一次
+}
+
+// 清理资源
+function cleanupResources() {
+    // 清除WebSocket连接
+    if (wsConnection) {
+        wsConnection.close();
+        wsConnection = null;
+    }
+    
+    // 清除重连计时器
+    if (wsReconnectTimer) {
+        clearTimeout(wsReconnectTimer);
+        wsReconnectTimer = null;
+    }
+    
+    // 清除定时刷新
+    if (intervalId) {
+        clearInterval(intervalId);
+        intervalId = null;
+    }
+}
+
 // 加载仪表盘数据
 async function loadDashboardData() {
     console.log('开始加载仪表盘数据');