godns/main.go

package main

import (
	"errors"
	"log"
	"math/rand"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/blang/semver"
	"github.com/miekg/dns"
	"github.com/rhysd/go-github-selfupdate/selfupdate"
	"github.com/yl2chen/cidranger"

	"godns/internal/handler"
	"godns/internal/model"
	"godns/internal/stats"
	"godns/internal/web"
	"godns/pkg/doh"
	"godns/pkg/logger"
)

var (
	version string

	config   *model.Config
	dataPath string
)

func main() {
	dataPath = detectDataPath()

	ipRanger := loadIPRanger(dataPath + "china_ip_list.txt")

	// 先创建一个临时 logger 用于读取配置
	tempLogger := logger.New(false)

	config = &model.Config{}
	if err := config.ReadInConfig(dataPath+"/config.json", ipRanger, tempLogger); err != nil {
		panic(err)
	}

	// 设置默认 Web 监听地址
	if config.WebAddr == "" {
		config.WebAddr = "0.0.0.0:8854"
	}

	// 根据配置创建正式的 logger 和 stats 实例
	debugLogger := logger.New(config.Debug)
	statsRecorder := stats.NewStats()

	// 加载持久化的统计数据
	if err := statsRecorder.Load(dataPath); err != nil {
		log.Printf("Failed to load stats from disk: %v", err)
	} else {
		log.Printf("Stats loaded successfully from disk")
	}

	// 更新 upstreams 的 logger 为正式的 logger
	for i := 0; i < len(config.Bootstrap); i++ {
		config.Bootstrap[i].SetLogger(debugLogger)
	}
	for i := 0; i < len(config.Upstreams); i++ {
		config.Upstreams[i].SetLogger(debugLogger)
	}

	// Bootstrap handler 不需要缓存，只是用于初始化连接
	bootstrapHandler := handler.NewHandler(model.StrategyAnyResult, false, config.Bootstrap, dataPath, debugLogger, nil)

	for i := 0; i < len(config.Upstreams); i++ {
		config.Upstreams[i].InitConnectionPool(bootstrapHandler.LookupIP)
	}

	server := &dns.Server{Addr: config.ServeAddr, Net: "udp"}
	serverTCP := &dns.Server{Addr: config.ServeAddr, Net: "tcp"}

	// 只有 upstream handler 需要缓存
	upstreamHandler := handler.NewHandler(config.Strategy, config.BuiltInCache, config.Upstreams, dataPath, debugLogger, statsRecorder)
	dns.HandleFunc(".", upstreamHandler.HandleRequest)

	// Setup graceful shutdown
	defer func() {
		// 保存统计数据
		log.Printf("Saving stats before shutdown...")
		if err := statsRecorder.Save(dataPath); err != nil {
			log.Printf("Error saving stats: %v", err)
		} else {
			log.Printf("Stats saved successfully")
		}

		// 关闭缓存
		if err := upstreamHandler.Close(); err != nil {
			log.Printf("Error closing cache: %v", err)
		}
	}()

	log.Println("==== DNS Server ====")
	log.Println("端口:", config.ServeAddr)
	log.Println("模式:", config.StrategyName())
	log.Println("数据:", dataPath)
	if config.BuiltInCache {
		log.Println("启用 BadgerDB 缓存: 最大 40MB")
	} else {
		log.Println("禁用缓存")
	}

	log.Println("版本:", version)

	// 创建更新检查通道
	checkUpdateCh := make(chan struct{}, 1)

	// 启动 Web 服务（监控面板 + DoH + pprof）
	webServerHandler := http.NewServeMux()

	// 注册监控面板路由
	var webUsername, webPassword string
	if config.WebAuth != nil {
		webUsername = config.WebAuth.Username
		webPassword = config.WebAuth.Password
	}
	webHandler := web.NewHandler(statsRecorder, version, checkUpdateCh, debugLogger, webUsername, webPassword)
	webHandler.RegisterRoutes(webServerHandler)

	// 如果启用 DoH，注册 DoH 路由
	if config.DohServer != nil {
		dohServer := doh.NewServer(config.DohServer.Username, config.DohServer.Password, upstreamHandler.HandleDnsMsg, statsRecorder)
		dohServer.RegisterRoutes(webServerHandler)
		log.Printf("DoH 服务: http://%s/dns-query", config.WebAddr)
	}

	// 如果启用 profiling，注册 pprof 路由
	if config.Profiling {
		webServerHandler.HandleFunc("/debug/", http.DefaultServeMux.ServeHTTP)
		log.Printf("性能分析: http://%s/debug/pprof/", config.WebAddr)
	}

	go http.ListenAndServe(config.WebAddr, webServerHandler)
	log.Printf("监控面板: http://%s/", config.WebAddr)

	// 定时保存统计数据（使用配置的间隔）
	statsSaveTicker := time.NewTicker(time.Duration(config.StatsSaveInterval) * time.Minute)
	defer statsSaveTicker.Stop()

	go func() {
		for range statsSaveTicker.C {
			if err := statsRecorder.Save(dataPath); err != nil {
				debugLogger.Printf("Failed to save stats to disk: %v", err)
			} else {
				debugLogger.Printf("Stats saved successfully to disk")
			}
		}
	}()

	stopCh := make(chan error)

	// 启动后台更新检查
	go checkUpdate(checkUpdateCh, stopCh, debugLogger)

	// 定时触发更新检查（生产者1：定时器）
	if version != "" {
		go func() {
			// 启动时立即检查一次
			select {
			case checkUpdateCh <- struct{}{}:
			default:
			}

			// 定时检查
			ticker := time.NewTicker(time.Duration(40+rand.Intn(20)) * time.Minute)
			defer ticker.Stop()
			for range ticker.C {
				select {
				case checkUpdateCh <- struct{}{}:
				default:
					// 如果通道已满，跳过本次
				}
			}
		}()
	}

	go func() {
		stopCh <- server.ListenAndServe()
	}()
	go func() {
		stopCh <- serverTCP.ListenAndServe()
	}()

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-sigCh
		log.Println("Shutting down...")
		stopCh <- errors.New("shutdown signal received")
	}()

	log.Printf("server stopped: %+v", <-stopCh)
}

// checkUpdate 监听 channel 触发更新检查
func checkUpdate(checkCh <-chan struct{}, stopCh chan<- error, debugLogger logger.Logger) {
	for range checkCh {
		// 如果 version 为空，使用默认值
		ver := version
		if ver == "" {
			ver = "0.0.0"
		}
		v := semver.MustParse(ver)
		latest, err := selfupdate.UpdateSelf(v, "xofine/godns")
		if err != nil {
			debugLogger.Printf("Error checking for updates: %v", err)
			continue
		}
		if latest.Version.Equals(v) {
			debugLogger.Printf("No update available, current version: %s", v)
		} else {
			log.Printf("Updated to version: %s", latest.Version)
			stopCh <- errors.New("Server upgraded to " + latest.Version.String())
			return
		}
	}
}

func loadIPRanger(path string) cidranger.Ranger {
	ipRanger := cidranger.NewPCTrieRanger()

	content, err := os.ReadFile(path)
	if err != nil {
		panic(err)
	}
	lines := strings.Split(string(content), "\n")

	for i := 0; i < len(lines); i++ {
		if strings.TrimSpace(lines[i]) == "" {
			continue
		}
		_, network, err := net.ParseCIDR(lines[i])
		if err != nil {
			panic(err)
		}
		if err := ipRanger.Insert(cidranger.NewBasicRangerEntry(*network)); err != nil {
			panic(err)
		}
	}

	return ipRanger
}

func detectDataPath() string {
	ex, err := os.Executable()
	if err != nil {
		panic(err)
	}
	pwd, err := os.Getwd()
	if err != nil {
		panic(err)
	}
	pathList := []string{filepath.Dir(ex), pwd}

	for _, path := range pathList {
		if f, err := os.Stat(path + "/data/china_ip_list.txt"); err == nil {
			if f.Size() == 1024*200 {
				panic("离线IP库 china_ip_list.txt 文件损坏，请重新下载")
			}
			return path + "/data/"
		}
	}

	panic("没有检测到IP数据 data/china_ip_list.txt")
}