实现日志功能
This commit is contained in:
244
dns/server.go
244
dns/server.go
@@ -29,15 +29,24 @@ type BlockedDomain struct {
|
||||
LastSeen time.Time
|
||||
}
|
||||
|
||||
// ClientStats 客户端统计
|
||||
|
||||
type ClientStats struct {
|
||||
IP string
|
||||
Count int64
|
||||
LastSeen time.Time
|
||||
}
|
||||
|
||||
// StatsData 用于持久化的统计数据结构
|
||||
type StatsData struct {
|
||||
Stats *Stats `json:"stats"`
|
||||
BlockedDomains map[string]*BlockedDomain `json:"blockedDomains"`
|
||||
ResolvedDomains map[string]*BlockedDomain `json:"resolvedDomains"`
|
||||
HourlyStats map[string]int64 `json:"hourlyStats"`
|
||||
DailyStats map[string]int64 `json:"dailyStats"`
|
||||
MonthlyStats map[string]int64 `json:"monthlyStats"`
|
||||
LastSaved time.Time `json:"lastSaved"`
|
||||
Stats *Stats `json:"stats"`
|
||||
BlockedDomains map[string]*BlockedDomain `json:"blockedDomains"`
|
||||
ResolvedDomains map[string]*BlockedDomain `json:"resolvedDomains"`
|
||||
ClientStats map[string]*ClientStats `json:"clientStats"`
|
||||
HourlyStats map[string]int64 `json:"hourlyStats"`
|
||||
DailyStats map[string]int64 `json:"dailyStats"`
|
||||
MonthlyStats map[string]int64 `json:"monthlyStats"`
|
||||
LastSaved time.Time `json:"lastSaved"`
|
||||
}
|
||||
|
||||
// Server DNS服务器
|
||||
@@ -46,6 +55,7 @@ type Server struct {
|
||||
shieldConfig *config.ShieldConfig
|
||||
shieldManager *shield.ShieldManager
|
||||
server *dns.Server
|
||||
tcpServer *dns.Server
|
||||
resolver *dns.Client
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
@@ -55,6 +65,8 @@ type Server struct {
|
||||
blockedDomains map[string]*BlockedDomain
|
||||
resolvedDomainsMutex sync.RWMutex
|
||||
resolvedDomains map[string]*BlockedDomain // 用于记录解析的域名
|
||||
clientStatsMutex sync.RWMutex
|
||||
clientStats map[string]*ClientStats // 用于记录客户端统计
|
||||
hourlyStatsMutex sync.RWMutex
|
||||
hourlyStats map[string]int64 // 按小时统计屏蔽数量
|
||||
dailyStatsMutex sync.RWMutex
|
||||
@@ -64,20 +76,22 @@ type Server struct {
|
||||
saveTicker *time.Ticker // 用于定时保存数据
|
||||
startTime time.Time // 服务器启动时间
|
||||
saveDone chan struct{} // 用于通知保存协程停止
|
||||
stopped bool // 服务器是否已经停止
|
||||
stoppedMutex sync.Mutex // 保护stopped标志的互斥锁
|
||||
}
|
||||
|
||||
// Stats DNS服务器统计信息
|
||||
type Stats struct {
|
||||
Queries int64
|
||||
Blocked int64
|
||||
Allowed int64
|
||||
Errors int64
|
||||
LastQuery time.Time
|
||||
AvgResponseTime float64 // 平均响应时间(ms)
|
||||
TotalResponseTime int64 // 总响应时间
|
||||
QueryTypes map[string]int64 // 查询类型统计
|
||||
SourceIPs map[string]bool // 活跃来源IP
|
||||
CpuUsage float64 // CPU使用率(%)
|
||||
Queries int64
|
||||
Blocked int64
|
||||
Allowed int64
|
||||
Errors int64
|
||||
LastQuery time.Time
|
||||
AvgResponseTime float64 // 平均响应时间(ms)
|
||||
TotalResponseTime int64 // 总响应时间
|
||||
QueryTypes map[string]int64 // 查询类型统计
|
||||
SourceIPs map[string]bool // 活跃来源IP
|
||||
CpuUsage float64 // CPU使用率(%)
|
||||
}
|
||||
|
||||
// NewServer 创建DNS服务器实例
|
||||
@@ -95,41 +109,59 @@ func NewServer(config *config.DNSConfig, shieldConfig *config.ShieldConfig, shie
|
||||
cancel: cancel,
|
||||
startTime: time.Now(), // 记录服务器启动时间
|
||||
stats: &Stats{
|
||||
Queries: 0,
|
||||
Blocked: 0,
|
||||
Allowed: 0,
|
||||
Errors: 0,
|
||||
AvgResponseTime: 0,
|
||||
Queries: 0,
|
||||
Blocked: 0,
|
||||
Allowed: 0,
|
||||
Errors: 0,
|
||||
AvgResponseTime: 0,
|
||||
TotalResponseTime: 0,
|
||||
QueryTypes: make(map[string]int64),
|
||||
SourceIPs: make(map[string]bool),
|
||||
CpuUsage: 0,
|
||||
QueryTypes: make(map[string]int64),
|
||||
SourceIPs: make(map[string]bool),
|
||||
CpuUsage: 0,
|
||||
},
|
||||
blockedDomains: make(map[string]*BlockedDomain),
|
||||
resolvedDomains: make(map[string]*BlockedDomain),
|
||||
clientStats: make(map[string]*ClientStats),
|
||||
hourlyStats: make(map[string]int64),
|
||||
dailyStats: make(map[string]int64),
|
||||
monthlyStats: make(map[string]int64),
|
||||
saveDone: make(chan struct{}),
|
||||
stopped: false, // 初始化为未停止状态
|
||||
}
|
||||
|
||||
|
||||
// 加载已保存的统计数据
|
||||
server.loadStatsData()
|
||||
|
||||
|
||||
return server
|
||||
|
||||
|
||||
}
|
||||
|
||||
// Start 启动DNS服务器
|
||||
func (s *Server) Start() error {
|
||||
// 重新初始化上下文和取消函数
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
s.ctx = ctx
|
||||
s.cancel = cancel
|
||||
|
||||
// 重新初始化saveDone通道
|
||||
s.saveDone = make(chan struct{})
|
||||
|
||||
// 重置stopped标志
|
||||
s.stoppedMutex.Lock()
|
||||
s.stopped = false
|
||||
s.stoppedMutex.Unlock()
|
||||
|
||||
// 更新服务器启动时间
|
||||
s.startTime = time.Now()
|
||||
|
||||
s.server = &dns.Server{
|
||||
Addr: fmt.Sprintf(":%d", s.config.Port),
|
||||
Net: "udp",
|
||||
Handler: dns.HandlerFunc(s.handleDNSRequest),
|
||||
}
|
||||
|
||||
// 启动TCP服务器(用于大型响应)
|
||||
tcpServer := &dns.Server{
|
||||
// 保存TCP服务器实例,以便在Stop方法中关闭
|
||||
s.tcpServer = &dns.Server{
|
||||
Addr: fmt.Sprintf(":%d", s.config.Port),
|
||||
Net: "tcp",
|
||||
Handler: dns.HandlerFunc(s.handleDNSRequest),
|
||||
@@ -138,6 +170,9 @@ func (s *Server) Start() error {
|
||||
// 启动CPU使用率监控
|
||||
go s.startCpuUsageMonitor()
|
||||
|
||||
// 启动自动保存功能
|
||||
go s.startAutoSave()
|
||||
|
||||
// 启动UDP服务
|
||||
go func() {
|
||||
logger.Info(fmt.Sprintf("DNS UDP服务器启动,监听端口: %d", s.config.Port))
|
||||
@@ -150,7 +185,7 @@ func (s *Server) Start() error {
|
||||
// 启动TCP服务
|
||||
go func() {
|
||||
logger.Info(fmt.Sprintf("DNS TCP服务器启动,监听端口: %d", s.config.Port))
|
||||
if err := tcpServer.ListenAndServe(); err != nil {
|
||||
if err := s.tcpServer.ListenAndServe(); err != nil {
|
||||
logger.Error("DNS TCP服务器启动失败", "error", err)
|
||||
s.cancel()
|
||||
}
|
||||
@@ -163,31 +198,44 @@ func (s *Server) Start() error {
|
||||
|
||||
// Stop 停止DNS服务器
|
||||
func (s *Server) Stop() {
|
||||
// 检查服务器是否已经停止
|
||||
s.stoppedMutex.Lock()
|
||||
if s.stopped {
|
||||
s.stoppedMutex.Unlock()
|
||||
return // 服务器已经停止,直接返回
|
||||
}
|
||||
// 标记服务器为已停止状态
|
||||
s.stopped = true
|
||||
s.stoppedMutex.Unlock()
|
||||
|
||||
// 发送停止信号给保存协程
|
||||
close(s.saveDone)
|
||||
|
||||
|
||||
// 最后保存一次数据
|
||||
s.saveStatsData()
|
||||
|
||||
|
||||
// 停止服务器
|
||||
s.cancel()
|
||||
if s.server != nil {
|
||||
s.server.Shutdown()
|
||||
}
|
||||
if s.tcpServer != nil {
|
||||
s.tcpServer.Shutdown()
|
||||
}
|
||||
logger.Info("DNS服务器已停止")
|
||||
}
|
||||
|
||||
// handleDNSRequest 处理DNS请求
|
||||
func (s *Server) handleDNSRequest(w dns.ResponseWriter, r *dns.Msg) {
|
||||
startTime := time.Now()
|
||||
|
||||
|
||||
// 获取来源IP
|
||||
sourceIP := w.RemoteAddr().String()
|
||||
// 提取IP地址部分,去掉端口
|
||||
if idx := strings.LastIndex(sourceIP, ":"); idx >= 0 {
|
||||
sourceIP = sourceIP[:idx]
|
||||
}
|
||||
|
||||
|
||||
// 更新来源IP统计
|
||||
s.updateStats(func(stats *Stats) {
|
||||
stats.Queries++
|
||||
@@ -195,6 +243,9 @@ func (s *Server) handleDNSRequest(w dns.ResponseWriter, r *dns.Msg) {
|
||||
stats.SourceIPs[sourceIP] = true
|
||||
})
|
||||
|
||||
// 更新客户端统计
|
||||
s.updateClientStats(sourceIP)
|
||||
|
||||
// 只处理递归查询
|
||||
if r.RecursionDesired == false {
|
||||
response := new(dns.Msg)
|
||||
@@ -202,7 +253,7 @@ func (s *Server) handleDNSRequest(w dns.ResponseWriter, r *dns.Msg) {
|
||||
response.RecursionAvailable = true
|
||||
response.SetRcode(r, dns.RcodeRefused)
|
||||
w.WriteMsg(response)
|
||||
|
||||
|
||||
// 计算响应时间
|
||||
responseTime := time.Since(startTime).Milliseconds()
|
||||
s.updateStats(func(stats *Stats) {
|
||||
@@ -315,11 +366,6 @@ func (s *Server) handleBlockedResponse(w dns.ResponseWriter, r *dns.Msg, domain
|
||||
// 更新被屏蔽域名统计
|
||||
s.updateBlockedDomainStats(domain)
|
||||
|
||||
// 更新总体统计
|
||||
s.updateStats(func(stats *Stats) {
|
||||
stats.Blocked++
|
||||
})
|
||||
|
||||
response := new(dns.Msg)
|
||||
response.SetReply(r)
|
||||
response.RecursionAvailable = true
|
||||
@@ -432,19 +478,19 @@ func (s *Server) updateBlockedDomainStats(domain string) {
|
||||
|
||||
// 更新统计数据
|
||||
now := time.Now()
|
||||
|
||||
|
||||
// 更新小时统计
|
||||
hourKey := now.Format("2006-01-02-15")
|
||||
s.hourlyStatsMutex.Lock()
|
||||
s.hourlyStats[hourKey]++
|
||||
s.hourlyStatsMutex.Unlock()
|
||||
|
||||
|
||||
// 更新每日统计
|
||||
dayKey := now.Format("2006-01-02")
|
||||
s.dailyStatsMutex.Lock()
|
||||
s.dailyStats[dayKey]++
|
||||
s.dailyStatsMutex.Unlock()
|
||||
|
||||
|
||||
// 更新每月统计
|
||||
monthKey := now.Format("2006-01")
|
||||
s.monthlyStatsMutex.Lock()
|
||||
@@ -452,6 +498,23 @@ func (s *Server) updateBlockedDomainStats(domain string) {
|
||||
s.monthlyStatsMutex.Unlock()
|
||||
}
|
||||
|
||||
// updateClientStats 更新客户端统计
|
||||
func (s *Server) updateClientStats(ip string) {
|
||||
s.clientStatsMutex.Lock()
|
||||
defer s.clientStatsMutex.Unlock()
|
||||
|
||||
if entry, exists := s.clientStats[ip]; exists {
|
||||
entry.Count++
|
||||
entry.LastSeen = time.Now()
|
||||
} else {
|
||||
s.clientStats[ip] = &ClientStats{
|
||||
IP: ip,
|
||||
Count: 1,
|
||||
LastSeen: time.Now(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// updateResolvedDomainStats 更新解析域名统计
|
||||
func (s *Server) updateResolvedDomainStats(domain string) {
|
||||
s.resolvedDomainsMutex.Lock()
|
||||
@@ -500,16 +563,16 @@ func (s *Server) GetStats() *Stats {
|
||||
|
||||
// 返回统计信息的副本
|
||||
return &Stats{
|
||||
Queries: s.stats.Queries,
|
||||
Blocked: s.stats.Blocked,
|
||||
Allowed: s.stats.Allowed,
|
||||
Errors: s.stats.Errors,
|
||||
LastQuery: s.stats.LastQuery,
|
||||
AvgResponseTime: s.stats.AvgResponseTime,
|
||||
Queries: s.stats.Queries,
|
||||
Blocked: s.stats.Blocked,
|
||||
Allowed: s.stats.Allowed,
|
||||
Errors: s.stats.Errors,
|
||||
LastQuery: s.stats.LastQuery,
|
||||
AvgResponseTime: s.stats.AvgResponseTime,
|
||||
TotalResponseTime: s.stats.TotalResponseTime,
|
||||
QueryTypes: queryTypesCopy,
|
||||
SourceIPs: sourceIPsCopy,
|
||||
CpuUsage: s.stats.CpuUsage,
|
||||
QueryTypes: queryTypesCopy,
|
||||
SourceIPs: sourceIPsCopy,
|
||||
CpuUsage: s.stats.CpuUsage,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -582,6 +645,29 @@ func (s *Server) GetRecentBlockedDomains(limit int) []BlockedDomain {
|
||||
return domains
|
||||
}
|
||||
|
||||
// GetTopClients 获取TOP客户端列表
|
||||
func (s *Server) GetTopClients(limit int) []ClientStats {
|
||||
s.clientStatsMutex.RLock()
|
||||
defer s.clientStatsMutex.RUnlock()
|
||||
|
||||
// 转换为切片
|
||||
clients := make([]ClientStats, 0, len(s.clientStats))
|
||||
for _, entry := range s.clientStats {
|
||||
clients = append(clients, *entry)
|
||||
}
|
||||
|
||||
// 按请求次数排序
|
||||
sort.Slice(clients, func(i, j int) bool {
|
||||
return clients[i].Count > clients[j].Count
|
||||
})
|
||||
|
||||
// 返回限制数量
|
||||
if len(clients) > limit {
|
||||
return clients[:limit]
|
||||
}
|
||||
return clients
|
||||
}
|
||||
|
||||
// GetHourlyStats 获取每小时统计数据
|
||||
func (s *Server) GetHourlyStats() map[string]int64 {
|
||||
s.hourlyStatsMutex.RLock()
|
||||
@@ -667,19 +753,26 @@ func (s *Server) loadStatsData() {
|
||||
s.hourlyStats = statsData.HourlyStats
|
||||
}
|
||||
s.hourlyStatsMutex.Unlock()
|
||||
|
||||
|
||||
s.dailyStatsMutex.Lock()
|
||||
if statsData.DailyStats != nil {
|
||||
s.dailyStats = statsData.DailyStats
|
||||
}
|
||||
s.dailyStatsMutex.Unlock()
|
||||
|
||||
|
||||
s.monthlyStatsMutex.Lock()
|
||||
if statsData.MonthlyStats != nil {
|
||||
s.monthlyStats = statsData.MonthlyStats
|
||||
}
|
||||
s.monthlyStatsMutex.Unlock()
|
||||
|
||||
// 加载客户端统计数据
|
||||
s.clientStatsMutex.Lock()
|
||||
if statsData.ClientStats != nil {
|
||||
s.clientStats = statsData.ClientStats
|
||||
}
|
||||
s.clientStatsMutex.Unlock()
|
||||
|
||||
logger.Info("统计数据加载成功")
|
||||
}
|
||||
|
||||
@@ -689,18 +782,25 @@ func (s *Server) saveStatsData() {
|
||||
return
|
||||
}
|
||||
|
||||
// 创建数据目录
|
||||
statsDir := filepath.Dir(s.config.StatsFile)
|
||||
err := os.MkdirAll(statsDir, 0755)
|
||||
// 获取绝对路径以避免工作目录问题
|
||||
statsFilePath, err := filepath.Abs(s.config.StatsFile)
|
||||
if err != nil {
|
||||
logger.Error("创建统计数据目录失败", "error", err)
|
||||
logger.Error("获取统计文件绝对路径失败", "path", s.config.StatsFile, "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
// 创建数据目录
|
||||
statsDir := filepath.Dir(statsFilePath)
|
||||
err = os.MkdirAll(statsDir, 0755)
|
||||
if err != nil {
|
||||
logger.Error("创建统计数据目录失败", "dir", statsDir, "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
// 收集所有统计数据
|
||||
statsData := &StatsData{
|
||||
Stats: s.GetStats(),
|
||||
LastSaved: time.Now(),
|
||||
Stats: s.GetStats(),
|
||||
LastSaved: time.Now(),
|
||||
}
|
||||
|
||||
// 复制域名数据
|
||||
@@ -724,14 +824,14 @@ func (s *Server) saveStatsData() {
|
||||
statsData.HourlyStats[k] = v
|
||||
}
|
||||
s.hourlyStatsMutex.RUnlock()
|
||||
|
||||
|
||||
s.dailyStatsMutex.RLock()
|
||||
statsData.DailyStats = make(map[string]int64)
|
||||
for k, v := range s.dailyStats {
|
||||
statsData.DailyStats[k] = v
|
||||
}
|
||||
s.dailyStatsMutex.RUnlock()
|
||||
|
||||
|
||||
s.monthlyStatsMutex.RLock()
|
||||
statsData.MonthlyStats = make(map[string]int64)
|
||||
for k, v := range s.monthlyStats {
|
||||
@@ -739,6 +839,14 @@ func (s *Server) saveStatsData() {
|
||||
}
|
||||
s.monthlyStatsMutex.RUnlock()
|
||||
|
||||
// 复制客户端统计数据
|
||||
s.clientStatsMutex.RLock()
|
||||
statsData.ClientStats = make(map[string]*ClientStats)
|
||||
for k, v := range s.clientStats {
|
||||
statsData.ClientStats[k] = v
|
||||
}
|
||||
s.clientStatsMutex.RUnlock()
|
||||
|
||||
// 序列化数据
|
||||
jsonData, err := json.MarshalIndent(statsData, "", " ")
|
||||
if err != nil {
|
||||
@@ -747,13 +855,13 @@ func (s *Server) saveStatsData() {
|
||||
}
|
||||
|
||||
// 写入文件
|
||||
err = ioutil.WriteFile(s.config.StatsFile, jsonData, 0644)
|
||||
err = os.WriteFile(statsFilePath, jsonData, 0644)
|
||||
if err != nil {
|
||||
logger.Error("保存统计数据到文件失败", "error", err)
|
||||
logger.Error("保存统计数据到文件失败", "file", statsFilePath, "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
logger.Info("统计数据保存成功")
|
||||
logger.Info("统计数据保存成功", "file", statsFilePath)
|
||||
}
|
||||
|
||||
// startCpuUsageMonitor 启动CPU使用率监控
|
||||
@@ -778,7 +886,7 @@ func (s *Server) startCpuUsageMonitor() {
|
||||
cpuUsage = 0.0
|
||||
logger.Error("获取系统CPU使用率失败", "error", err)
|
||||
}
|
||||
|
||||
|
||||
s.updateStats(func(stats *Stats) {
|
||||
stats.CpuUsage = cpuUsage
|
||||
})
|
||||
@@ -798,7 +906,7 @@ func getSystemCpuUsage(prevIdle, prevTotal *uint64) (float64, error) {
|
||||
defer file.Close()
|
||||
|
||||
var cpuUser, cpuNice, cpuSystem, cpuIdle, cpuIowait, cpuIrq, cpuSoftirq, cpuSteal uint64
|
||||
_, err = fmt.Fscanf(file, "cpu %d %d %d %d %d %d %d %d",
|
||||
_, err = fmt.Fscanf(file, "cpu %d %d %d %d %d %d %d %d",
|
||||
&cpuUser, &cpuNice, &cpuSystem, &cpuIdle, &cpuIowait, &cpuIrq, &cpuSoftirq, &cpuSteal)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
|
||||
Reference in New Issue
Block a user