dns-server/main.go

// DNS Server API
// @title DNS Server API
// @version 1.0
// @description DNS服务器API文档
// @termsOfService http://swagger.io/terms/
// @contact.name API Support
// @contact.email support@example.com
// @license.name Apache 2.0
// @license.url http://www.apache.org/licenses/LICENSE-2.0.html
// @host localhost:8080
// @BasePath /api
package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"

	"dns-server/config"
	"dns-server/dns"
	"dns-server/http"
	"dns-server/logger"
	"dns-server/shield"
)

// createDefaultConfig 创建默认配置文件
func createDefaultConfig(configFile string) error {
	// 默认配置内容
	defaultConfig := `{
  "dns": {
    "port": 53,
    "upstreamDNS": [
      "223.5.5.5:53",
      "223.6.6.6:53"
    ],
    "timeout": 5000,
    "statsFile": "./data/stats.json",
    "saveInterval": 300
  },
  "http": {
    "port": 8080,
    "host": "0.0.0.0",
    "enableAPI": true,
    "username": "admin",
    "password": "admin"
  },
  "shield": {
    "localRulesFile": "data/rules.txt",
    "blacklists": [
      {
        "name": "AdGuard DNS filter",
        "url": "https://gitea.amazehome.xyz/AMAZEHOME/hosts-and-filters/raw/branch/main/filter.txt",
        "enabled": true
      },
      {
        "name": "Adaway Default Blocklist",
        "url": "https://gitea.amazehome.xyz/AMAZEHOME/hosts-and-Filters/raw/branch/main/hosts/adaway.txt",
        "enabled": true
      },
      {
        "name": "CHN-anti-AD",
        "url": "https://gitea.amazehome.xyz/AMAZEHOME/hosts-and-Filters/raw/branch/main/list/easylist.txt",
        "enabled": true
      },
      {
        "name": "My GitHub Rules",
        "url": "https://gitea.amazehome.xyz/AMAZEHOME/hosts-and-filters/raw/branch/main/rules/costomize.txt",
        "enabled": true
      }
    ],
    "updateInterval": 3600,
    "hostsFile": "data/hosts.txt",
    "blockMethod": "NXDOMAIN",
    "customBlockIP": "",
    "statsFile": "./data/shield_stats.json",
    "statsSaveInterval": 60,
    "remoteRulesCacheDir": "./data/remote_rules"
  },
  "log": {
    "file": "logs/dns-server.log",
    "level": "debug",
    "maxSize": 100,
    "maxBackups": 10,
    "maxAge": 30
  }
}`

	// 写入默认配置到文件
	return os.WriteFile(configFile, []byte(defaultConfig), 0644)
}

// createRequiredFiles 创建所需的文件和文件夹
func createRequiredFiles(cfg *config.Config) error {
	// 创建数据文件夹
	dataDir := "./data"
	if err := os.MkdirAll(dataDir, 0755); err != nil {
		return fmt.Errorf("创建数据文件夹失败: %w", err)
	}

	// 创建远程规则缓存文件夹
	if err := os.MkdirAll(cfg.Shield.RemoteRulesCacheDir, 0755); err != nil {
		return fmt.Errorf("创建远程规则缓存文件夹失败: %w", err)
	}

	// 创建日志文件夹
	logDir := filepath.Dir(cfg.Log.File)
	if logDir != "." {
		if err := os.MkdirAll(logDir, 0755); err != nil {
			return fmt.Errorf("创建日志文件夹失败: %w", err)
		}
	}

	// 创建本地规则文件
	if _, err := os.Stat(cfg.Shield.LocalRulesFile); os.IsNotExist(err) {
		if err := os.WriteFile(cfg.Shield.LocalRulesFile, []byte("# 本地规则文件\n# 格式：域名\n# 例如：example.com\n"), 0644); err != nil {
			return fmt.Errorf("创建本地规则文件失败: %w", err)
		}
	}

	// 创建Hosts文件
	if _, err := os.Stat(cfg.Shield.HostsFile); os.IsNotExist(err) {
		if err := os.WriteFile(cfg.Shield.HostsFile, []byte("# Hosts文件\n# 格式：IP 域名\n# 例如：127.0.0.1 localhost\n"), 0644); err != nil {
			return fmt.Errorf("创建Hosts文件失败: %w", err)
		}
	}

	// 创建统计数据文件
	if _, err := os.Stat(cfg.DNS.StatsFile); os.IsNotExist(err) {
		if err := os.WriteFile(cfg.DNS.StatsFile, []byte("{}"), 0644); err != nil {
			return fmt.Errorf("创建统计数据文件失败: %w", err)
		}
	}

	// 创建Shield统计数据文件
	if _, err := os.Stat(cfg.Shield.StatsFile); os.IsNotExist(err) {
		if err := os.WriteFile(cfg.Shield.StatsFile, []byte("{}"), 0644); err != nil {
			return fmt.Errorf("创建Shield统计数据文件失败: %w", err)
		}
	}

	return nil
}

func main() {
	// 命令行参数解析
	var configFile string
	var daemonMode bool
	flag.StringVar(&configFile, "config", "config.json", "配置文件路径")
	flag.BoolVar(&daemonMode, "daemon", false, "以守护进程模式运行")
	flag.Parse()

	// 如果是守护进程模式，创建守护进程
	if daemonMode {
		if err := daemonize(); err != nil {
			log.Fatalf("创建守护进程失败: %v", err)
		}
		// 父进程退出
		os.Exit(0)
	}

	// 检查配置文件是否存在，如果不存在则创建默认配置文件
	if _, err := os.Stat(configFile); os.IsNotExist(err) {
		log.Printf("配置文件 %s 不存在，正在创建默认配置文件...", configFile)
		if err := createDefaultConfig(configFile); err != nil {
			log.Fatalf("创建默认配置文件失败: %v", err)
		}
		log.Printf("默认配置文件 %s 创建成功", configFile)
	}

	// 初始化配置
	var cfg *config.Config
	var err error
	cfg, err = config.LoadConfig(configFile)
	if err != nil {
		log.Fatalf("加载配置失败: %v", err)
	}

	// 创建所需的文件和文件夹
	log.Println("正在创建所需的文件和文件夹...")
	if err := createRequiredFiles(cfg); err != nil {
		log.Fatalf("创建所需文件和文件夹失败: %v", err)
	}
	log.Println("所需文件和文件夹创建成功")

	// 初始化日志系统
	if err := logger.InitLogger(cfg.Log.File, cfg.Log.Level, 0, 0, 0); err != nil {
		log.Fatalf("初始化日志系统失败: %v", err)
	}
	defer logger.Close()

	// 初始化屏蔽管理系统
	shieldManager := shield.NewShieldManager(&cfg.Shield)
	if err := shieldManager.LoadRules(); err != nil {
		logger.Error("加载屏蔽规则失败", "error", err)
	}

	// 启动DNS服务器
	dnsServer := dns.NewServer(&cfg.DNS, &cfg.Shield, shieldManager)
	go func() {
		if err := dnsServer.Start(); err != nil {
			logger.Error("DNS服务器启动失败", "error", err)
			os.Exit(1)
		}
	}()

	// 启动HTTP控制台服务器
	httpServer := http.NewServer(cfg, dnsServer, shieldManager)
	go func() {
		if err := httpServer.Start(); err != nil {
			logger.Error("HTTP控制台服务器启动失败", "error", err)
		}
	}()

	// 启动定时更新任务
	go shieldManager.StartAutoUpdate()

	logger.Info(fmt.Sprintf("DNS服务器已启动，监听端口: %d", cfg.DNS.Port))
	logger.Info(fmt.Sprintf("HTTP控制台已启动，监听端口: %d", cfg.HTTP.Port))

	// 监听信号
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	<-sigCh

	// 清理资源
	log.Println("正在关闭服务...")
	dnsServer.Stop()
	httpServer.Stop()
	shieldManager.StopAutoUpdate()
	// 守护进程模式下不需要删除PID文件

	log.Println("服务已关闭")
}

// daemonize 创建守护进程
func daemonize() error {
	// 1. 第一次fork，创建子进程，父进程退出
	pid, err := syscall.Fork()
	if err != nil {
		return fmt.Errorf("第一次fork失败: %w", err)
	}
	if pid > 0 {
		// 父进程退出
		os.Exit(0)
	}

	// 2. 设置文件权限掩码
	syscall.Umask(0)

	// 3. 创建新会话
	_, err = syscall.Setsid()
	if err != nil {
		return fmt.Errorf("创建新会话失败: %w", err)
	}

	// 4. 第二次fork，确保进程不是会话组长，避免获取控制终端
	pid, err = syscall.Fork()
	if err != nil {
		return fmt.Errorf("第二次fork失败: %w", err)
	}
	if pid > 0 {
		// 父进程退出
		os.Exit(0)
	}

	// 5. 切换工作目录到根目录
	err = syscall.Chdir("/")
	if err != nil {
		return fmt.Errorf("切换工作目录失败: %w", err)
	}

	// 6. 关闭所有打开的文件描述符
	for fd := 0; fd < 1024; fd++ {
		syscall.Close(fd)
	}

	// 7. 重定向标准输入、输出、错误到/dev/null
	nullFile, err := os.OpenFile("/dev/null", os.O_RDWR, 0)
	if err != nil {
		return fmt.Errorf("打开/dev/null失败: %w", err)
	}
	defer nullFile.Close()

	// 重定向标准输入、输出、错误
	syscall.Dup2(int(nullFile.Fd()), 0) // stdin
	syscall.Dup2(int(nullFile.Fd()), 1) // stdout
	syscall.Dup2(int(nullFile.Fd()), 2) // stderr

	// 8. 处理SIGCHLD信号，避免产生僵尸进程
	syscall.Signal(syscall.SIGCHLD, syscall.SIG_IGN)

	// 9. 写入PID文件
	pid = syscall.Getpid()
	pidFile, err := os.OpenFile("/var/run/dns-server.pid", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return fmt.Errorf("打开PID文件失败: %w", err)
	}
	_, err = fmt.Fprintf(pidFile, "%d\n", pid)
	if err != nil {
		pidFile.Close()
		return fmt.Errorf("写入PID文件失败: %w", err)
	}
	pidFile.Close()

	return nil
}