WAF规则匹配后的IP也会上报/实现IP全局名单/将名单存储到本地数据库，提升读写速度

internal/iplibrary/README.md

@@ -0,0 +1,5 @@
+# IPList
+List Check Order:
+~~~
+Global List --> Node List--> Server List --> WAF List --> Bind List
+~~~

internal/iplibrary/ip_list.go

@@ -7,6 +7,9 @@ import (
 	"sync"
 )
 
+var GlobalBlackIPList = NewIPList()
+var GlobalWhiteIPList = NewIPList()
+
 // IPList IP名单
 // TODO IP名单可以分片关闭，这样让每一片的数据量减少，查询更快
 type IPList struct {

internal/iplibrary/ip_list_db.go

@@ -0,0 +1,145 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package iplibrary
+
+import (
+	"database/sql"
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
+	"github.com/iwind/TeaGo/Tea"
+	_ "github.com/mattn/go-sqlite3"
+	"os"
+	"path/filepath"
+)
+
+type IPListDB struct {
+	db *sql.DB
+
+	itemTableName   string
+	deleteItemStmt  *sql.Stmt
+	insertItemStmt  *sql.Stmt
+	selectItemsStmt *sql.Stmt
+
+	dir string
+}
+
+func NewIPListDB() (*IPListDB, error) {
+	var db = &IPListDB{
+		itemTableName: "ipItems",
+		dir:           filepath.Clean(Tea.Root + "/data"),
+	}
+	err := db.init()
+	return db, err
+}
+
+func (this *IPListDB) init() error {
+	// 检查目录是否存在
+	_, err := os.Stat(this.dir)
+	if err != nil {
+		err = os.MkdirAll(this.dir, 0777)
+		if err != nil {
+			return err
+		}
+		remotelogs.Println("CACHE", "create cache dir '"+this.dir+"'")
+	}
+
+	db, err := sql.Open("sqlite3", "file:"+this.dir+"/ip_list.db?cache=shared&mode=rwc&_journal_mode=WAL")
+	if err != nil {
+		return err
+	}
+	db.SetMaxOpenConns(1)
+	this.db = db
+
+	// 初始化数据库
+	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS "` + this.itemTableName + `" (
+  "id" integer NOT NULL PRIMARY KEY AUTOINCREMENT,
+  "listId" integer DEFAULT 0,
+  "listType" varchar(32),
+  "isGlobal" integer(1) DEFAULT 0,
+  "type" varchar(16),
+  "itemId" integer DEFAULT 0,
+  "ipFrom" varchar(64) DEFAULT 0,
+  "ipTo" varchar(64) DEFAULT 0,
+  "expiredAt" integer DEFAULT 0,
+  "eventLevel" varchar(32),
+  "isDeleted" integer(1) DEFAULT 0,
+  "version" integer DEFAULT 0,
+  "nodeId" integer DEFAULT 0,
+  "serverId" integer DEFAULT 0
+);
+
+CREATE INDEX IF NOT EXISTS "ip_list_itemId"
+ON "` + this.itemTableName + `" (
+  "itemId" ASC
+);
+
+CREATE INDEX IF NOT EXISTS "ip_list_expiredAt"
+ON "` + this.itemTableName + `" (
+  "expiredAt" ASC
+);
+`)
+	if err != nil {
+		return err
+	}
+
+	// 初始化SQL语句
+	this.deleteItemStmt, err = this.db.Prepare(`DELETE FROM "` + this.itemTableName + `" WHERE "itemId"=?`)
+	if err != nil {
+		return err
+	}
+
+	this.insertItemStmt, err = this.db.Prepare(`INSERT INTO "` + this.itemTableName + `" ("listId", "listType", "isGlobal", "type", "itemId", "ipFrom", "ipTo", "expiredAt", "eventLevel", "isDeleted", "version", "nodeId", "serverId") VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`)
+	if err != nil {
+		return err
+	}
+
+	this.selectItemsStmt, err = this.db.Prepare(`SELECT "listId", "listType", "isGlobal", "type", "itemId", "ipFrom", "ipTo", "expiredAt", "eventLevel", "isDeleted", "version", "nodeId", "serverId" FROM "` + this.itemTableName + `" ORDER BY "version" ASC, "itemId" ASC LIMIT ?, ?`)
+	if err != nil {
+		return err
+	}
+
+	this.db = db
+
+	return nil
+}
+
+func (this *IPListDB) AddItem(item *pb.IPItem) error {
+	_, err := this.deleteItemStmt.Exec(item.Id)
+	if err != nil {
+		return err
+	}
+	_, err = this.insertItemStmt.Exec(item.ListId, item.ListType, item.IsGlobal, item.Type, item.Id, item.IpFrom, item.IpTo, item.ExpiredAt, item.EventLevel, item.IsDeleted, item.Version, item.NodeId, item.ServerId)
+	return err
+}
+
+func (this *IPListDB) ReadItems(offset int64, size int64) (items []*pb.IPItem, err error) {
+	rows, err := this.selectItemsStmt.Query(offset, size)
+	if err != nil {
+		return nil, err
+	}
+	defer func() {
+		_ = rows.Close()
+	}()
+
+	for rows.Next() {
+		//  "listId", "listType", "isGlobal", "type", "itemId", "ipFrom", "ipTo", "expiredAt", "eventLevel", "isDeleted", "version", "nodeId", "serverId"
+		var pbItem = &pb.IPItem{}
+		err = rows.Scan(&pbItem.ListId, &pbItem.ListType, &pbItem.IsGlobal, &pbItem.Type, &pbItem.Id, &pbItem.IpFrom, &pbItem.IpTo, &pbItem.ExpiredAt, &pbItem.EventLevel, &pbItem.IsDeleted, &pbItem.Version, &pbItem.NodeId, &pbItem.ServerId)
+		if err != nil {
+			return nil, err
+		}
+		items = append(items, pbItem)
+	}
+	return
+}
+
+func (this *IPListDB) Close() error {
+	if this.db != nil {
+		_ = this.deleteItemStmt.Close()
+		_ = this.insertItemStmt.Close()
+		_ = this.selectItemsStmt.Close()
+
+		return this.db.Close()
+	}
+	return nil
+}

internal/iplibrary/ip_list_db_test.go

@@ -0,0 +1,60 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package iplibrary
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	_ "github.com/iwind/TeaGo/bootstrap"
+	"github.com/iwind/TeaGo/logs"
+	"testing"
+	"time"
+)
+
+func TestIPListDB_AddItem(t *testing.T) {
+	db, err := NewIPListDB()
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = db.AddItem(&pb.IPItem{
+		Id:                            1,
+		IpFrom:                        "192.168.1.101",
+		IpTo:                          "",
+		Version:                       1024,
+		ExpiredAt:                     time.Now().Unix(),
+		Reason:                        "",
+		ListId:                        2,
+		IsDeleted:                     true,
+		Type:                          "ipv4",
+		EventLevel:                    "error",
+		ListType:                      "black",
+		IsGlobal:                      true,
+		CreatedAt:                     0,
+		NodeId:                        11,
+		ServerId:                      22,
+		SourceNodeId:                  0,
+		SourceServerId:                0,
+		SourceHTTPFirewallPolicyId:    0,
+		SourceHTTPFirewallRuleGroupId: 0,
+		SourceHTTPFirewallRuleSetId:   0,
+		SourceServer:                  nil,
+		SourceHTTPFirewallPolicy:      nil,
+		SourceHTTPFirewallRuleGroup:   nil,
+		SourceHTTPFirewallRuleSet:     nil,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Log("ok")
+}
+
+func TestIPListDB_ReadItems(t *testing.T) {
+	db, err := NewIPListDB()
+	if err != nil {
+		t.Fatal(err)
+	}
+	items, err := db.ReadItems(0, 2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	logs.PrintAsJSON(items, t)
+}

internal/iplibrary/list_utils.go

@@ -0,0 +1,55 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package iplibrary
+
+import (
+	"github.com/TeaOSLab/EdgeNode/internal/utils"
+)
+
+// AllowIP 检查IP是否被允许访问
+func AllowIP(ip string, serverId int64) bool {
+	var ipLong = utils.IP2Long(ip)
+	if ipLong == 0 {
+		return false
+	}
+
+	// check white lists
+	if GlobalWhiteIPList.Contains(ipLong) {
+		return true
+	}
+
+	if serverId > 0 {
+		var list = SharedServerListManager.FindWhiteList(serverId, false)
+		if list != nil && list.Contains(ipLong) {
+			return true
+		}
+	}
+
+	// check black lists
+	if GlobalBlackIPList.Contains(ipLong) {
+		return false
+	}
+
+	if serverId > 0 {
+		var list = SharedServerListManager.FindBlackList(serverId, false)
+		if list != nil && list.Contains(ipLong) {
+			return false
+		}
+	}
+
+	return true
+}
+
+// AllowIPStrings 检查一组IP是否被允许访问
+func AllowIPStrings(ipStrings []string, serverId int64) bool {
+	if len(ipStrings) == 0 {
+		return true
+	}
+	for _, ip := range ipStrings {
+		isAllowed := AllowIP(ip, serverId)
+		if !isAllowed {
+			return false
+		}
+	}
+	return true
+}

internal/iplibrary/list_utils_test.go

@@ -0,0 +1,20 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package iplibrary
+
+import (
+	"testing"
+	"time"
+)
+
+func TestIPIsAllowed(t *testing.T) {
+	manager := NewIPListManager()
+	manager.init()
+
+	var before = time.Now()
+	defer func() {
+		t.Log(time.Since(before).Seconds()*1000, "ms")
+	}()
+	t.Log(AllowIP("127.0.0.1", 0))
+	t.Log(AllowIP("127.0.0.1", 23))
+}

internal/iplibrary/manager_ip_list.go

@@ -8,9 +8,6 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/utils"
 	"github.com/TeaOSLab/EdgeNode/internal/waf"
 	"github.com/iwind/TeaGo/Tea"
-	"github.com/iwind/TeaGo/types"
-	"io/ioutil"
-	"os"
 	"sync"
 	"time"
 )
@@ -24,13 +21,9 @@ func init() {
 	})
 }
 
-var versionCacheFile = "ip_list_version.cache"
-
 // IPListManager IP名单管理
 type IPListManager struct {
-	// 缓存文件
-	// 每行一个数据：id|from|to|expiredAt
-	cacheFile string
+	db *IPListDB
 
 	version  int64
 	pageSize int64
@@ -41,17 +34,13 @@ type IPListManager struct {
 
 func NewIPListManager() *IPListManager {
 	return &IPListManager{
-		cacheFile: Tea.Root + "/configs/ip_list.cache",
-		pageSize:  500,
-		listMap:   map[int64]*IPList{},
+		pageSize: 500,
+		listMap:  map[int64]*IPList{},
 	}
 }
 
 func (this *IPListManager) Start() {
-	// TODO 从缓存当中读取数据
-
-	// 从缓存中读取位置
-	this.version = this.readLocalVersion()
+	this.init()
 
 	// 第一次读取
 	err := this.loop()
@@ -60,6 +49,9 @@ func (this *IPListManager) Start() {
 	}
 
 	ticker := time.NewTicker(60 * time.Second)
+	if Tea.IsTesting() {
+		ticker = time.NewTicker(10 * time.Second)
+	}
 	events.On(events.EventQuit, func() {
 		ticker.Stop()
 	})
@@ -88,6 +80,31 @@ func (this *IPListManager) Start() {
 	}
 }
 
+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()
@@ -119,11 +136,53 @@ func (this *IPListManager) fetch() (hasNext bool, err error) {
 	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 {
-		list, ok := this.listMap[item.ListId]
-		if !ok {
+		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
 		}
@@ -133,18 +192,13 @@ func (this *IPListManager) fetch() (hasNext bool, err error) {
 		if item.IsDeleted {
 			list.Delete(item.Id)
 
-			// 从临时名单中删除
-			if len(item.IpFrom) > 0 && len(item.IpTo) == 0 {
-				switch item.ListType {
-				case "black":
-					waf.SharedIPBlackList.RemoveIP(item.IpFrom)
-				case "white":
-					waf.SharedIPWhiteList.RemoveIP(item.IpFrom)
-				}
-			}
+			// 从WAF名单中删除
+			waf.SharedIPBlackList.RemoveIP(item.IpFrom, item.ServerId)
 
 			// 操作事件
-			SharedActionManager.DeleteItem(item.ListType, item)
+			if shouldExecute {
+				SharedActionManager.DeleteItem(item.ListType, item)
+			}
 
 			continue
 		}
@@ -159,8 +213,10 @@ func (this *IPListManager) fetch() (hasNext bool, err error) {
 		})
 
 		// 事件操作
-		SharedActionManager.DeleteItem(item.ListType, item)
-		SharedActionManager.AddItem(item.ListType, item)
+		if shouldExecute {
+			SharedActionManager.DeleteItem(item.ListType, item)
+			SharedActionManager.AddItem(item.ListType, item)
+		}
 	}
 
 	for changedList := range changedLists {
@@ -169,38 +225,4 @@ func (this *IPListManager) fetch() (hasNext bool, err error) {
 
 	this.locker.Unlock()
 	this.version = items[len(items)-1].Version
-
-	// 写入版本号到缓存当中
-	this.updateLocalVersion(this.version)
-
-	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) readLocalVersion() int64 {
-	data, err := ioutil.ReadFile(Tea.ConfigFile(versionCacheFile))
-	if err != nil || len(data) == 0 {
-		return 0
-	}
-	return types.Int64(string(data))
-}
-
-func (this *IPListManager) updateLocalVersion(version int64) {
-	fp, err := os.OpenFile(Tea.ConfigFile(versionCacheFile), os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0666)
-	if err != nil {
-		remotelogs.Warn("IP_LIST", "write local version cache failed: "+err.Error())
-		return
-	}
-	_, err = fp.WriteString(types.String(version))
-	if err != nil {
-		remotelogs.Warn("IP_LIST", "write local version cache failed: "+err.Error())
-		return
-	}
-	_ = fp.Close()
 }

internal/iplibrary/manager_ip_list_test.go

@@ -1,10 +1,36 @@
 package iplibrary
 
-import "testing"
+import (
+	"github.com/TeaOSLab/EdgeNode/internal/utils"
+	"github.com/iwind/TeaGo/logs"
+	"testing"
+	"time"
+)
+
+func TestIPListManager_init(t *testing.T) {
+	manager := NewIPListManager()
+	manager.init()
+	t.Log(manager.listMap)
+	t.Log(SharedServerListManager.blackMap)
+	logs.PrintAsJSON(GlobalBlackIPList.sortedItems, t)
+}
+
+func TestIPListManager_check(t *testing.T) {
+	manager := NewIPListManager()
+	manager.init()
+
+	var before = time.Now()
+	defer func() {
+		t.Log(time.Since(before).Seconds()*1000, "ms")
+	}()
+	t.Log(SharedServerListManager.FindBlackList(23, true).Contains(utils.IP2Long("127.0.0.2")))
+	t.Log(GlobalBlackIPList.Contains(utils.IP2Long("127.0.0.6")))
+}
 
 func TestIPListManager_loop(t *testing.T) {
 	manager := NewIPListManager()
-	manager.pageSize = 2
+	manager.Start()
+	manager.pageSize = 10
 	err := manager.loop()
 	if err != nil {
 		t.Fatal(err)

internal/iplibrary/server_list_manager.go

@@ -0,0 +1,61 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package iplibrary
+
+import "sync"
+
+var SharedServerListManager = NewServerListManager()
+
+// ServerListManager 服务相关名单
+type ServerListManager struct {
+	whiteMap map[int64]*IPList // serverId => *List
+	blackMap map[int64]*IPList // serverId => *List
+
+	locker sync.RWMutex
+}
+
+func NewServerListManager() *ServerListManager {
+	return &ServerListManager{
+		whiteMap: map[int64]*IPList{},
+		blackMap: map[int64]*IPList{},
+	}
+}
+
+func (this *ServerListManager) FindWhiteList(serverId int64, autoCreate bool) *IPList {
+	this.locker.RLock()
+	list, ok := this.whiteMap[serverId]
+	this.locker.RUnlock()
+	if ok {
+		return list
+	}
+
+	if autoCreate {
+		list = NewIPList()
+		this.locker.Lock()
+		this.whiteMap[serverId] = list
+		this.locker.Unlock()
+
+		return list
+	}
+	return nil
+}
+
+func (this *ServerListManager) FindBlackList(serverId int64, autoCreate bool) *IPList {
+	this.locker.RLock()
+	list, ok := this.blackMap[serverId]
+	this.locker.RUnlock()
+	if ok {
+		return list
+	}
+
+	if autoCreate {
+		list = NewIPList()
+		this.locker.Lock()
+		this.blackMap[serverId] = list
+		this.locker.Unlock()
+
+		return list
+	}
+
+	return nil
+}