waf-platform/EdgeDNS/internal/nodes/manager_key.go

// 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
	readyCh  chan struct{} // 初始加载完成后关闭
}

// 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),
		readyCh:      make(chan struct{}),
	}
}

// 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())
		}
	}

	// 通知初始加载完成
	close(this.readyCh)

	// 更新
	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()
}