1.4.5.2

EdgeDNS/internal/nodes/manager_key.go

@@ -0,0 +1,284 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package nodes
+
+import (
+	"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"
+	"sync"
+	"time"
+)
+
+// KeyManager 密钥管理器
+type KeyManager struct {
+	domainKeyMap map[int64]*models.NSKeys // domainId => *NSKeys
+	zoneKeyMap   map[int64]*models.NSKeys // zoneId => *NSKeys
+
+	db      *dbs.DB
+	locker  sync.RWMutex
+	version int64
+
+	notifier chan bool
+}
+
+// NewKeyManager 获取密钥管理器
+func NewKeyManager(db *dbs.DB) *KeyManager {
+	return &KeyManager{
+		domainKeyMap: map[int64]*models.NSKeys{},
+		zoneKeyMap:   map[int64]*models.NSKeys{},
+		db:           db,
+		notifier:     make(chan bool, 8),
+	}
+}
+
+// Start 启动自动任务
+func (this *KeyManager) Start() {
+	remotelogs.Println("KEY_MANAGER", "starting ...")
+
+	// 从本地数据库中加载数据
+	err := this.Load()
+	if err != nil {
+		if rpc.IsConnError(err) {
+			remotelogs.Debug("KEY_MANAGER", "load failed: "+err.Error())
+		} else {
+			remotelogs.Error("KEY_MANAGER", "load failed: "+err.Error())
+		}
+	}
+
+	// 初始化运行
+	err = this.LoopAll()
+	if err != nil {
+		if rpc.IsConnError(err) {
+			remotelogs.Debug("KEY_MANAGER", "loop failed: "+err.Error())
+		} else {
+			remotelogs.Error("KEY_MANAGER", "loop failed: "+err.Error())
+		}
+	}
+
+	// 更新
+	var ticker = time.NewTicker(1 * time.Minute)
+	for {
+		select {
+		case <-ticker.C:
+		case <-this.notifier:
+		}
+
+		err := this.LoopAll()
+		if err != nil {
+			if rpc.IsConnError(err) {
+				remotelogs.Debug("KEY_MANAGER", "loop failed: "+err.Error())
+			} else {
+				remotelogs.Error("KEY_MANAGER", "loop failed: "+err.Error())
+			}
+		}
+	}
+}
+
+// Load 从数据库中加载数据
+func (this *KeyManager) Load() error {
+	var offset = 0
+	var size = 10000
+	for {
+		keys, err := this.db.ListKeys(offset, size)
+		if err != nil {
+			return err
+		}
+		if len(keys) == 0 {
+			break
+		}
+
+		this.locker.Lock()
+		for _, key := range keys {
+			if key.ZoneId > 0 {
+				keyList, ok := this.zoneKeyMap[key.ZoneId]
+				if ok {
+					keyList.Add(key)
+				} else {
+					keyList = models.NewNSKeys()
+					keyList.Add(key)
+					this.zoneKeyMap[key.ZoneId] = keyList
+				}
+			} else if key.DomainId > 0 {
+				keyList, ok := this.domainKeyMap[key.DomainId]
+				if ok {
+					keyList.Add(key)
+				} else {
+					keyList = models.NewNSKeys()
+					keyList.Add(key)
+					this.domainKeyMap[key.DomainId] = keyList
+				}
+			}
+
+			if key.Version > this.version {
+				this.version = key.Version
+			}
+		}
+		this.locker.Unlock()
+
+		offset += size
+	}
+
+	if this.version > 0 {
+		this.version++
+	}
+
+	return nil
+}
+
+func (this *KeyManager) LoopAll() error {
+	for {
+		hasNext, err := this.Loop()
+		if err != nil {
+			return err
+		}
+		if !hasNext {
+			break
+		}
+	}
+	return nil
+}
+
+// Loop 单次循环任务
+func (this *KeyManager) Loop() (hasNext bool, err error) {
+	client, err := rpc.SharedRPC()
+	if err != nil {
+		return false, err
+	}
+	resp, err := client.NSKeyRPC.ListNSKeysAfterVersion(client.Context(), &pb.ListNSKeysAfterVersionRequest{
+		Version: this.version,
+		Size:    20000,
+	})
+	if err != nil {
+		return false, err
+	}
+	var keys = resp.NsKeys
+	if len(keys) == 0 {
+		return false, nil
+	}
+	for _, key := range keys {
+		this.processKey(key)
+
+		if key.Version > this.version {
+			this.version = key.Version
+		}
+	}
+	this.version++
+	return true, nil
+}
+
+func (this *KeyManager) FindKeysWithDomain(domainId int64) []*models.NSKey {
+	this.locker.RLock()
+	defer this.locker.RUnlock()
+
+	keys, ok := this.domainKeyMap[domainId]
+	if ok {
+		return keys.All()
+	}
+	return nil
+}
+
+// NotifyUpdate 通知更新
+func (this *KeyManager) NotifyUpdate() {
+	select {
+	case this.notifier <- true:
+	default:
+	}
+}
+
+// 处理Key
+func (this *KeyManager) processKey(key *pb.NSKey) {
+	if key.NsDomain == nil && key.NsZone == nil {
+		return
+	}
+	if !key.IsOn || key.IsDeleted {
+		this.locker.Lock()
+		if key.NsDomain != nil {
+			list, ok := this.domainKeyMap[key.NsDomain.Id]
+			if ok {
+				list.Remove(key.Id)
+			}
+		}
+		if key.NsZone != nil {
+			list, ok := this.zoneKeyMap[key.NsZone.Id]
+			if ok {
+				list.Remove(key.Id)
+			}
+		}
+		this.locker.Unlock()
+
+		// 从数据库中删除
+		if this.db != nil {
+			err := this.db.DeleteKey(key.Id)
+			if err != nil {
+				remotelogs.Error("KEY_MANAGER", "delete key from db failed: "+err.Error())
+			}
+		}
+
+		return
+	}
+
+	var domainId int64
+	var zoneId int64
+	if key.NsDomain != nil {
+		domainId = key.NsDomain.Id
+	}
+	if key.NsZone != nil {
+		zoneId = key.NsZone.Id
+	}
+
+	// 存入数据库
+	if this.db != nil {
+		exists, err := this.db.ExistsKey(key.Id)
+		if err != nil {
+			remotelogs.Error("KEY_MANAGER", "query failed: "+err.Error())
+		} else {
+			if exists {
+				err = this.db.UpdateKey(key.Id, domainId, zoneId, key.Algo, key.Secret, key.SecretType, key.Version)
+				if err != nil {
+					remotelogs.Error("KEY_MANAGER", "update failed: "+err.Error())
+				}
+			} else {
+				err = this.db.InsertKey(key.Id, domainId, zoneId, key.Algo, key.Secret, key.SecretType, key.Version)
+				if err != nil {
+					remotelogs.Error("KEY_MANAGER", "insert failed: "+err.Error())
+				}
+			}
+		}
+	}
+
+	// 加入缓存Map
+	this.locker.Lock()
+	var nsKey = &models.NSKey{
+		Id:         key.Id,
+		DomainId:   domainId,
+		ZoneId:     zoneId,
+		Algo:       key.Algo,
+		Secret:     key.Secret,
+		SecretType: key.SecretType,
+		Version:    key.Version,
+	}
+
+	if zoneId > 0 {
+		keyList, ok := this.zoneKeyMap[zoneId]
+		if ok {
+			keyList.Add(nsKey)
+		} else {
+			keyList = models.NewNSKeys()
+			keyList.Add(nsKey)
+			this.zoneKeyMap[zoneId] = keyList
+		}
+	} else if domainId > 0 {
+		keyList, ok := this.domainKeyMap[domainId]
+		if ok {
+			keyList.Add(nsKey)
+		} else {
+			keyList = models.NewNSKeys()
+			keyList.Add(nsKey)
+			this.domainKeyMap[domainId] = keyList
+		}
+	}
+	this.locker.Unlock()
+}