sk-matrix-project/backend/golang/matrix-dnslog-service/cmd/service/main.go
2023-03-24 16:12:11 +08:00

63 lines
1.6 KiB
Go

package service
import (
"github.com/miekg/dns"
"github.com/spf13/viper"
"matrix-common/constants"
"matrix-common/pkg/logger"
"matrix-dnslog-service/pkg/config"
"matrix-dnslog-service/pkg/handler"
"strings"
)
func Main() {
viper.SetConfigName(constants.ConfigFileName)
viper.SetConfigType(constants.ConfigType)
for _, path := range constants.ConfigPaths {
viper.AddConfigPath(path)
}
if err := viper.ReadInConfig(); err != nil {
if _, ok := err.(viper.ConfigFileNotFoundError); ok {
_ = config.GenerateConfig()
logger.Log().Fatalf("未找到配置文件, 已生成示例配置文件于运行路径下")
} else {
logger.Log().Fatalf("配置解析失败 %s", err)
}
}
conf := &config.Conf{}
err := viper.Unmarshal(conf)
if err != nil {
logger.Log().Fatalf("配置文件解析失败: %s, 请检查配置是否有误", err)
}
serv := &dns.Server{Addr: conf.Addr, Net: "udp"}
if conf.Proxy {
logger.Log().Infof("[Proxy] dns 代理已启用 代理将转发至服务器: %s", conf.ProxyServer)
dns.HandleFunc(".", func(w dns.ResponseWriter, r *dns.Msg) {
qname := r.Question[0].Name
logger.Log().Infof("[Proxy] 查询域名 %s\n", qname)
resp, _ := dns.Exchange(r, conf.ProxyServer)
_ = w.WriteMsg(resp)
})
}
dns.HandleFunc(conf.DomainPattern, func(w dns.ResponseWriter, r *dns.Msg) {
qname := r.Question[0].Name
// 去除后缀 只保留前缀
name := strings.TrimSuffix(qname, conf.DomainPattern)
// 处理数据
handler.ProcessHandler(name)
m := new(dns.Msg)
m.SetReply(r)
_ = w.WriteMsg(m)
})
if err := serv.ListenAndServe(); err != nil {
logger.Log().Fatalf("服务启动失败: %v\n", err)
}
}