实现IP黑白名单、国家｜地区封禁、省份封禁

internal/iplibrary/manager_ip_list.go

@@ -0,0 +1,124 @@
+package iplibrary
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	"github.com/TeaOSLab/EdgeNode/internal/events"
+	"github.com/TeaOSLab/EdgeNode/internal/logs"
+	"github.com/TeaOSLab/EdgeNode/internal/rpc"
+	"github.com/iwind/TeaGo/Tea"
+	"sync"
+	"time"
+)
+
+var SharedIPListManager = NewIPListManager()
+
+func init() {
+	events.On(events.EventStart, func() {
+		go SharedIPListManager.Start()
+	})
+}
+
+// IP名单管理
+type IPListManager struct {
+	// 缓存文件
+	// 每行一个数据：id|from|to|expiredAt
+	cacheFile string
+
+	version  int64
+	pageSize int64
+
+	listMap map[int64]*IPList
+	locker  sync.Mutex
+}
+
+func NewIPListManager() *IPListManager {
+	return &IPListManager{
+		cacheFile: Tea.Root + "/configs/ip_list.cache",
+		pageSize:  1000,
+		listMap:   map[int64]*IPList{},
+	}
+}
+
+func (this *IPListManager) Start() {
+	// TODO 从缓存当中读取数据
+
+	// 第一次读取
+	err := this.loop()
+	if err != nil {
+		logs.Println("IP_LIST_MANAGER", err.Error())
+	}
+
+	ticker := time.NewTicker(60 * time.Second) // TODO 未来改成可以手动触发IP变更事件
+	events.On(events.EventQuit, func() {
+		ticker.Stop()
+	})
+	for range ticker.C {
+		err := this.loop()
+		if err != nil {
+			logs.Println("IP_LIST_MANAGER", err.Error())
+		}
+	}
+}
+
+func (this *IPListManager) loop() error {
+	for {
+		hasNext, err := this.fetch()
+		if err != nil {
+			return err
+		}
+		if !hasNext {
+			break
+		}
+	}
+
+	// TODO 写入到缓存当中
+
+	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
+	}
+	this.locker.Lock()
+	for _, item := range items {
+		list, ok := this.listMap[item.ListId]
+		if !ok {
+			list = NewIPList()
+			this.listMap[item.ListId] = list
+		}
+		if item.IsDeleted {
+			list.Delete(item.Id)
+			continue
+		}
+		list.Add(&IPItem{
+			Id:        item.Id,
+			IPFrom:    IP2Long(item.IpFrom),
+			IPTo:      IP2Long(item.IpTo),
+			ExpiredAt: item.ExpiredAt,
+		})
+	}
+	this.locker.Unlock()
+	this.version = items[len(items)-1].Version
+
+	return true, nil
+}
+
+func (this *IPListManager) FindList(listId int64) *IPList {
+	this.locker.Lock()
+	list, _ := this.listMap[listId]
+	this.locker.Unlock()
+	return list
+}