331 lines
7.9 KiB
Go
331 lines
7.9 KiB
Go
// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
|
||
|
||
package nodes
|
||
|
||
import (
|
||
"encoding/json"
|
||
"github.com/TeaOSLab/EdgeCommon/pkg/dnsconfigs"
|
||
"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
|
||
"github.com/TeaOSLab/EdgeDNS/internal/dbs"
|
||
"github.com/TeaOSLab/EdgeDNS/internal/models"
|
||
"github.com/TeaOSLab/EdgeDNS/internal/remotelogs"
|
||
"github.com/TeaOSLab/EdgeDNS/internal/rpc"
|
||
"github.com/iwind/TeaGo/types"
|
||
"strings"
|
||
"sync"
|
||
"sync/atomic"
|
||
"time"
|
||
)
|
||
|
||
// DomainManager 域名管理器
|
||
type DomainManager struct {
|
||
domainMap map[int64]*models.NSDomain // domainId => domain
|
||
namesMap map[string]map[int64]*models.NSDomain // domain name => { domainId => domain }
|
||
clusterId int64
|
||
|
||
db *dbs.DB
|
||
version int64
|
||
locker *sync.RWMutex
|
||
dbWriteFailures atomic.Int64 // DB 写入累计失败次数
|
||
|
||
notifier chan bool
|
||
readyCh chan struct{} // 初始加载完成后关闭
|
||
}
|
||
|
||
// NewDomainManager 获取域名管理器对象
|
||
func NewDomainManager(db *dbs.DB, clusterId int64) *DomainManager {
|
||
return &DomainManager{
|
||
db: db,
|
||
domainMap: map[int64]*models.NSDomain{},
|
||
namesMap: map[string]map[int64]*models.NSDomain{},
|
||
clusterId: clusterId,
|
||
notifier: make(chan bool, 8),
|
||
locker: &sync.RWMutex{},
|
||
readyCh: make(chan struct{}),
|
||
}
|
||
}
|
||
|
||
// Start 启动自动任务
|
||
func (this *DomainManager) Start() {
|
||
remotelogs.Println("DOMAIN_MANAGER", "starting ...")
|
||
|
||
// 从本地数据库中加载数据
|
||
err := this.Load()
|
||
if err != nil {
|
||
if rpc.IsConnError(err) {
|
||
remotelogs.Debug("DOMAIN_MANAGER", "load failed: "+err.Error())
|
||
} else {
|
||
remotelogs.Error("DOMAIN_MANAGER", "load failed: "+err.Error())
|
||
}
|
||
}
|
||
|
||
// 初始化运行
|
||
err = this.LoopAll()
|
||
if err != nil {
|
||
if rpc.IsConnError(err) {
|
||
remotelogs.Debug("DOMAIN_MANAGER", "loop failed: "+err.Error())
|
||
} else {
|
||
remotelogs.Error("DOMAIN_MANAGER", "loop failed: "+err.Error())
|
||
}
|
||
}
|
||
|
||
// 通知初始加载完成
|
||
close(this.readyCh)
|
||
|
||
// 更新
|
||
var ticker = time.NewTicker(20 * time.Second)
|
||
for {
|
||
select {
|
||
case <-ticker.C:
|
||
case <-this.notifier:
|
||
}
|
||
|
||
err = this.LoopAll()
|
||
if err != nil {
|
||
if rpc.IsConnError(err) {
|
||
remotelogs.Debug("DOMAIN_MANAGER", "loop failed: "+err.Error())
|
||
} else {
|
||
remotelogs.Error("DOMAIN_MANAGER", "loop failed: "+err.Error())
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
func (this *DomainManager) LoopAll() error {
|
||
for {
|
||
hasNext, err := this.Loop()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if !hasNext {
|
||
break
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// Load 从数据库中加载数据
|
||
func (this *DomainManager) Load() error {
|
||
var offset = 0
|
||
var size = 10000
|
||
for {
|
||
domains, err := this.db.ListDomains(this.clusterId, offset, size)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if len(domains) == 0 {
|
||
break
|
||
}
|
||
|
||
this.locker.Lock()
|
||
for _, domain := range domains {
|
||
this.domainMap[domain.Id] = domain
|
||
|
||
nameMap, ok := this.namesMap[domain.Name]
|
||
if ok {
|
||
nameMap[domain.Id] = domain
|
||
} else {
|
||
this.namesMap[domain.Name] = map[int64]*models.NSDomain{
|
||
domain.Id: domain,
|
||
}
|
||
}
|
||
|
||
if domain.Version > this.version {
|
||
this.version = domain.Version
|
||
}
|
||
}
|
||
this.locker.Unlock()
|
||
|
||
offset += size
|
||
}
|
||
|
||
if this.version > 0 {
|
||
this.version++
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// Loop 单次循环任务
|
||
func (this *DomainManager) Loop() (hasNext bool, err error) {
|
||
client, err := rpc.SharedRPC()
|
||
if err != nil {
|
||
return false, err
|
||
}
|
||
resp, err := client.NSDomainRPC.ListNSDomainsAfterVersion(client.Context(), &pb.ListNSDomainsAfterVersionRequest{
|
||
Version: this.version,
|
||
Size: 20000,
|
||
})
|
||
if err != nil {
|
||
return false, err
|
||
}
|
||
var domains = resp.NsDomains
|
||
if len(domains) == 0 {
|
||
return false, nil
|
||
}
|
||
for _, domain := range domains {
|
||
this.processDomain(domain)
|
||
if domain.Version > this.version {
|
||
this.version = domain.Version
|
||
}
|
||
}
|
||
this.version++
|
||
return true, nil
|
||
}
|
||
|
||
// FindDomain 根据名称查找域名
|
||
func (this *DomainManager) FindDomain(name string) (domain *models.NSDomain, ok bool) {
|
||
this.locker.RLock()
|
||
defer this.locker.RUnlock()
|
||
|
||
nameMap, ok := this.namesMap[name]
|
||
if !ok {
|
||
return nil, false
|
||
}
|
||
|
||
for _, domain2 := range nameMap {
|
||
return domain2, true
|
||
}
|
||
return
|
||
}
|
||
|
||
// FindDomainWithId 根据域名ID查询域名
|
||
func (this *DomainManager) FindDomainWithId(domainId int64) (domain *models.NSDomain) {
|
||
this.locker.RLock()
|
||
defer this.locker.RUnlock()
|
||
|
||
return this.domainMap[domainId]
|
||
}
|
||
|
||
// NotifyUpdate 通知更新
|
||
func (this *DomainManager) NotifyUpdate() {
|
||
select {
|
||
case this.notifier <- true:
|
||
default:
|
||
}
|
||
}
|
||
|
||
// SplitDomain 分解域名
|
||
func (this *DomainManager) SplitDomain(fullDomainName string) (rootDomain *models.NSDomain, recordName string) {
|
||
if len(fullDomainName) == 0 {
|
||
return
|
||
}
|
||
fullDomainName = strings.TrimSuffix(fullDomainName, ".") // 去除尾部的点（.）
|
||
fullDomainName = strings.ToLower(fullDomainName) // 转换为小写
|
||
var domainName = fullDomainName
|
||
var domain, ok = this.FindDomain(domainName)
|
||
if !ok {
|
||
for {
|
||
var index = strings.Index(domainName, ".")
|
||
if index < 0 {
|
||
break
|
||
}
|
||
domainName = domainName[index+1:]
|
||
domain, ok = this.FindDomain(domainName)
|
||
if ok {
|
||
recordName = fullDomainName[:len(fullDomainName)-len(domainName)-1]
|
||
break
|
||
}
|
||
}
|
||
}
|
||
return domain, recordName
|
||
}
|
||
|
||
// 处理域名
|
||
func (this *DomainManager) processDomain(domain *pb.NSDomain) {
|
||
if !domain.IsOn || domain.IsDeleted || domain.Status != dnsconfigs.NSDomainStatusVerified {
|
||
this.locker.Lock()
|
||
delete(this.domainMap, domain.Id)
|
||
|
||
nameMap, ok := this.namesMap[domain.Name]
|
||
if ok {
|
||
delete(nameMap, domain.Id)
|
||
if len(nameMap) == 0 {
|
||
delete(this.namesMap, domain.Name)
|
||
}
|
||
}
|
||
|
||
this.locker.Unlock()
|
||
|
||
// 从数据库中删除
|
||
if this.db != nil {
|
||
err := this.db.DeleteDomain(domain.Id)
|
||
if err != nil {
|
||
count := this.dbWriteFailures.Add(1)
|
||
remotelogs.Error("DOMAIN_MANAGER", "delete domain from db failed (total failures: "+types.String(count)+"): "+err.Error())
|
||
}
|
||
}
|
||
|
||
return
|
||
}
|
||
|
||
// 存入数据库
|
||
if this.db != nil {
|
||
exists, err := this.db.ExistsDomain(domain.Id)
|
||
if err != nil {
|
||
count := this.dbWriteFailures.Add(1)
|
||
remotelogs.Error("DOMAIN_MANAGER", "query failed (total failures: "+types.String(count)+"): "+err.Error())
|
||
} else {
|
||
if exists {
|
||
err = this.db.UpdateDomain(domain.Id, domain.NsCluster.Id, domain.UserId, domain.Name, domain.TsigJSON, domain.Version)
|
||
if err != nil {
|
||
count := this.dbWriteFailures.Add(1)
|
||
remotelogs.Error("DOMAIN_MANAGER", "update failed (total failures: "+types.String(count)+"): "+err.Error())
|
||
}
|
||
} else {
|
||
err = this.db.InsertDomain(domain.Id, domain.NsCluster.Id, domain.UserId, domain.Name, domain.TsigJSON, domain.Version)
|
||
if err != nil {
|
||
count := this.dbWriteFailures.Add(1)
|
||
remotelogs.Error("DOMAIN_MANAGER", "insert failed (total failures: "+types.String(count)+"): "+err.Error())
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// 同集群的才需要加载
|
||
if this.clusterId == domain.NsCluster.Id {
|
||
this.locker.Lock()
|
||
var tsigConfig = &dnsconfigs.NSTSIGConfig{}
|
||
if len(domain.TsigJSON) > 0 {
|
||
err := json.Unmarshal(domain.TsigJSON, tsigConfig)
|
||
if err != nil {
|
||
remotelogs.Error("DOMAIN_MANAGER", "decode TSIG json failed: "+err.Error()+", domain: "+domain.Name+", domainId: "+types.String(domain.Id)+", JSON: "+string(domain.TsigJSON))
|
||
}
|
||
}
|
||
|
||
var nsDomain = &models.NSDomain{
|
||
Id: domain.Id,
|
||
ClusterId: domain.NsCluster.Id,
|
||
UserId: domain.UserId,
|
||
Name: domain.Name,
|
||
TSIG: tsigConfig,
|
||
Version: domain.Version,
|
||
}
|
||
this.domainMap[domain.Id] = nsDomain
|
||
|
||
nameMap, ok := this.namesMap[domain.Name]
|
||
if ok {
|
||
nameMap[nsDomain.Id] = nsDomain
|
||
} else {
|
||
this.namesMap[domain.Name] = map[int64]*models.NSDomain{
|
||
nsDomain.Id: nsDomain,
|
||
}
|
||
}
|
||
this.locker.Unlock()
|
||
} else {
|
||
// 不同集群的删除域名
|
||
this.locker.Lock()
|
||
delete(this.domainMap, domain.Id)
|
||
|
||
nameMap, ok := this.namesMap[domain.Name]
|
||
if ok {
|
||
delete(nameMap, domain.Id)
|
||
if len(nameMap) == 0 {
|
||
delete(this.namesMap, domain.Name)
|
||
}
|
||
}
|
||
|
||
this.locker.Unlock()
|
||
}
|
||
}
|