diff --git a/dns_server b/dns_server new file mode 100755 index 0000000..a287171 Binary files /dev/null and b/dns_server differ diff --git a/go.mod b/go.mod index c5d4ddf..433ba33 100644 --- a/go.mod +++ b/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 ) diff --git a/go.sum b/go.sum index a445c83..068bbb2 100644 --- a/go.sum +++ b/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= diff --git a/http/server.go b/http/server.go index 3eaad95..bcd83e2 100644 --- a/http/server.go +++ b/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 { diff --git a/static/js/dashboard.js b/static/js/dashboard.js index c3d1f5d..3e119d6 100644 --- a/static/js/dashboard.js +++ b/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(); - // 设置定时更新 - intervalId = setInterval(loadDashboardData, 5000); // 每5秒更新一次 + // 建立WebSocket连接 + 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('开始加载仪表盘数据');