EdgeNode/internal/iplibrary/manager_ip_list.go

package iplibrary

import (
	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
	"github.com/TeaOSLab/EdgeNode/internal/events"
	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
	"github.com/TeaOSLab/EdgeNode/internal/rpc"
	"github.com/TeaOSLab/EdgeNode/internal/utils"
	"github.com/TeaOSLab/EdgeNode/internal/waf"
	"github.com/iwind/TeaGo/Tea"
	"sync"
	"time"
)

var SharedIPListManager = NewIPListManager()
var IPListUpdateNotify = make(chan bool, 1)

func init() {
	events.On(events.EventLoaded, func() {
		go SharedIPListManager.Start()
	})
}

// IPListManager IP名单管理
type IPListManager struct {
	db *IPListDB

	version  int64
	pageSize int64

	listMap map[int64]*IPList
	locker  sync.Mutex
}

func NewIPListManager() *IPListManager {
	return &IPListManager{
		pageSize: 500,
		listMap:  map[int64]*IPList{},
	}
}

func (this *IPListManager) Start() {
	this.init()

	// 第一次读取
	err := this.loop()
	if err != nil {
		remotelogs.ErrorObject("IP_LIST_MANAGER", err)
	}

	ticker := time.NewTicker(60 * time.Second)
	if Tea.IsTesting() {
		ticker = time.NewTicker(10 * time.Second)
	}
	events.On(events.EventQuit, func() {
		ticker.Stop()
	})
	countErrors := 0
	for {
		select {
		case <-ticker.C:
		case <-IPListUpdateNotify:
		}
		err := this.loop()
		if err != nil {
			countErrors++

			remotelogs.ErrorObject("IP_LIST_MANAGER", err)

			// 连续错误小于3次的我们立即重试
			if countErrors <= 3 {
				select {
				case IPListUpdateNotify <- true:
				default:
				}
			}
		} else {
			countErrors = 0
		}
	}
}

func (this *IPListManager) init() {
	// 从数据库中当中读取数据
	db, err := NewIPListDB()
	if err != nil {
		remotelogs.Error("IP_LIST_MANAGER", "create ip list local database failed: "+err.Error())
	} else {
		this.db = db

		var offset int64 = 0
		var size int64 = 1000
		for {
			items, err := db.ReadItems(offset, size)
			if err != nil {
				remotelogs.Error("IP_LIST_MANAGER", "read ip list from local database failed: "+err.Error())
			} else {
				if len(items) == 0 {
					break
				}
				this.processItems(items, false)
			}
			offset += int64(len(items))
		}
	}
}

func (this *IPListManager) loop() error {
	for {
		hasNext, err := this.fetch()
		if err != nil {
			return err
		}
		if !hasNext {
			break
		}
		time.Sleep(1 * time.Second)
	}

	return nil
}

func (this *IPListManager) fetch() (hasNext bool, err error) {
	rpcClient, err := rpc.SharedRPC()
	if err != nil {
		return false, err
	}
	itemsResp, err := rpcClient.IPItemRPC().ListIPItemsAfterVersion(rpcClient.Context(), &pb.ListIPItemsAfterVersionRequest{
		Version: this.version,
		Size:    this.pageSize,
	})
	if err != nil {
		return false, err
	}
	items := itemsResp.IpItems
	if len(items) == 0 {
		return false, nil
	}

	// 保存到本地数据库
	if this.db != nil {
		for _, item := range items {
			err = this.db.AddItem(item)
			if err != nil {
				remotelogs.Error("IP_LIST_MANAGER", "insert item to local database failed: "+err.Error())
			}
		}
	}

	this.processItems(items, true)

	return true, nil
}

func (this *IPListManager) FindList(listId int64) *IPList {
	this.locker.Lock()
	list, _ := this.listMap[listId]
	this.locker.Unlock()
	return list
}

func (this *IPListManager) processItems(items []*pb.IPItem, shouldExecute bool) {
	this.locker.Lock()
	var changedLists = map[*IPList]bool{}
	for _, item := range items {
		var list *IPList
		// TODO 实现节点专有List
		if item.ServerId > 0 { // 服务专有List
			switch item.ListType {
			case "black":
				list = SharedServerListManager.FindBlackList(item.ServerId, true)
			case "white":
				list = SharedServerListManager.FindWhiteList(item.ServerId, true)
			}
		} else if item.IsGlobal { // 全局List
			switch item.ListType {
			case "black":
				list = GlobalBlackIPList
			case "white":
				list = GlobalWhiteIPList
			}
		} else { // 其他List
			list = this.listMap[item.ListId]
		}
		if list == nil {
			list = NewIPList()
			this.listMap[item.ListId] = list
		}

		changedLists[list] = true

		if item.IsDeleted {
			list.Delete(item.Id)

			// 从WAF名单中删除
			waf.SharedIPBlackList.RemoveIP(item.IpFrom, item.ServerId)

			// 操作事件
			if shouldExecute {
				SharedActionManager.DeleteItem(item.ListType, item)
			}

			continue
		}

		list.AddDelay(&IPItem{
			Id:         item.Id,
			Type:       item.Type,
			IPFrom:     utils.IP2Long(item.IpFrom),
			IPTo:       utils.IP2Long(item.IpTo),
			ExpiredAt:  item.ExpiredAt,
			EventLevel: item.EventLevel,
		})

		// 事件操作
		if shouldExecute {
			SharedActionManager.DeleteItem(item.ListType, item)
			SharedActionManager.AddItem(item.ListType, item)
		}
	}

	for changedList := range changedLists {
		changedList.Sort()
	}

	this.locker.Unlock()
	this.version = items[len(items)-1].Version
}