实现API节点下载更新IP库、基本的分析函数

internal/db/models/sys_setting_dao.go

@@ -1,7 +1,9 @@
 package models
 
 import (
+	"encoding/json"
 	"fmt"
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
 	_ "github.com/go-sql-driver/mysql"
 	"github.com/iwind/TeaGo/Tea"
 	"github.com/iwind/TeaGo/dbs"
@@ -109,3 +111,20 @@ func (this *SysSettingDAO) CompareInt64Setting(code string, anotherValue int64)
 	}
 	return 0, nil
 }
+
+// 读取全局配置
+func (this *SysSettingDAO) ReadGlobalConfig() (*serverconfigs.GlobalConfig, error) {
+	globalConfigData, err := this.ReadSetting(SettingCodeServerGlobalConfig)
+	if err != nil {
+		return nil, err
+	}
+	if len(globalConfigData) == 0 {
+		return &serverconfigs.GlobalConfig{}, nil
+	}
+	config := &serverconfigs.GlobalConfig{}
+	err = json.Unmarshal(globalConfigData, config)
+	if err != nil {
+		return nil, err
+	}
+	return config, nil
+}

internal/iplibrary/init.go

@@ -0,0 +1,5 @@
+package iplibrary
+
+func init() {
+
+}

internal/iplibrary/library_interface.go

@@ -0,0 +1,12 @@
+package iplibrary
+
+type LibraryInterface interface {
+	// 加载数据库文件
+	Load(dbPath string) error
+
+	// 查询IP
+	Lookup(ip string) (*Result, error)
+
+	// 关闭数据库文件
+	Close()
+}

internal/iplibrary/library_ip2region.go

@@ -0,0 +1,56 @@
+package iplibrary
+
+import (
+	"fmt"
+	"github.com/TeaOSLab/EdgeAPI/internal/errors"
+	"github.com/iwind/TeaGo/logs"
+	"github.com/lionsoul2014/ip2region/binding/golang/ip2region"
+)
+
+type IP2RegionLibrary struct {
+	db *ip2region.Ip2Region
+}
+
+func (this *IP2RegionLibrary) Load(dbPath string) error {
+	db, err := ip2region.New(dbPath)
+	if err != nil {
+		return err
+	}
+	this.db = db
+
+	return nil
+}
+
+func (this *IP2RegionLibrary) Lookup(ip string) (*Result, error) {
+	if this.db == nil {
+		return nil, errors.New("library has not been loaded")
+	}
+
+	defer func() {
+		// 防止panic发生
+		err := recover()
+		if err != nil {
+			logs.Println("[IP2RegionLibrary]panic: " + fmt.Sprintf("%#v", err))
+		}
+	}()
+
+	info, err := this.db.MemorySearch(ip)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Result{
+		CityId:   info.CityId,
+		Country:  info.Country,
+		Region:   info.Region,
+		Province: info.Province,
+		City:     info.City,
+		ISP:      info.ISP,
+	}, nil
+}
+
+func (this *IP2RegionLibrary) Close() {
+	if this.db != nil {
+		this.db.Close()
+	}
+}

internal/iplibrary/library_ip2region_test.go

@@ -0,0 +1,55 @@
+package iplibrary
+
+import (
+	"github.com/iwind/TeaGo/Tea"
+	_ "github.com/iwind/TeaGo/bootstrap"
+	"github.com/iwind/TeaGo/logs"
+	"github.com/iwind/TeaGo/rands"
+	"runtime"
+	"strconv"
+	"testing"
+	"time"
+)
+
+func TestIP2RegionLibrary_Lookup(t *testing.T) {
+	library := &IP2RegionLibrary{}
+	err := library.Load(Tea.Root + "/resources/ipdata/ip2region/ip2region.db")
+	if err != nil {
+		t.Fatal(err)
+	}
+	result, err := library.Lookup("114.240.223.47")
+	if err != nil {
+		t.Fatal(err)
+	}
+	logs.PrintAsJSON(result, t)
+}
+
+func TestIP2RegionLibrary_Memory(t *testing.T) {
+	library := &IP2RegionLibrary{}
+	err := library.Load(Tea.Root + "/resources/ipdata/ip2region/ip2region.db")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	before := time.Now()
+
+	for i := 0; i < 1_000_000; i++ {
+		_, _ = library.Lookup(strconv.Itoa(rands.Int(0, 254)) + "." + strconv.Itoa(rands.Int(0, 254)) + "." + strconv.Itoa(rands.Int(0, 254)) + "." + strconv.Itoa(rands.Int(0, 254)))
+	}
+
+	t.Log("cost:", time.Since(before).Seconds()*1000, "ms")
+}
+
+func BenchmarkIP2RegionLibrary_Lookup(b *testing.B) {
+	runtime.GOMAXPROCS(1)
+
+	library := &IP2RegionLibrary{}
+	err := library.Load(Tea.Root + "/resources/ipdata/ip2region/ip2region.db")
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	for i := 0; i < b.N; i++ {
+		_, _ = library.Lookup(strconv.Itoa(rands.Int(0, 254)) + "." + strconv.Itoa(rands.Int(0, 254)) + "." + strconv.Itoa(rands.Int(0, 254)) + "." + strconv.Itoa(rands.Int(0, 254)))
+	}
+}

internal/iplibrary/manager.go

@@ -0,0 +1,90 @@
+package iplibrary
+
+import (
+	"github.com/TeaOSLab/EdgeAPI/internal/db/models"
+	"github.com/TeaOSLab/EdgeAPI/internal/errors"
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
+	"github.com/iwind/TeaGo/Tea"
+	"github.com/iwind/TeaGo/dbs"
+	"github.com/iwind/TeaGo/files"
+	"github.com/iwind/TeaGo/logs"
+	"github.com/iwind/TeaGo/types"
+	"regexp"
+	"strings"
+)
+
+var SharedManager = NewManager()
+var SharedLibrary LibraryInterface
+
+func init() {
+	dbs.OnReady(func() {
+		// 初始化
+		library, err := SharedManager.Load()
+		if err != nil {
+			logs.Println("[IP_LIBRARY]" + err.Error())
+			return
+		}
+		SharedLibrary = library
+	})
+}
+
+type Manager struct {
+	code string
+}
+
+func NewManager() *Manager {
+	return &Manager{}
+}
+
+func (this *Manager) Load() (LibraryInterface, error) {
+	// 当前正在使用的IP库代号
+	config, err := models.SharedSysSettingDAO.ReadGlobalConfig()
+	if err != nil {
+		return nil, err
+	}
+	code := config.IPLibrary.Code
+	if len(code) == 0 {
+		code = serverconfigs.DefaultIPLibraryType
+	}
+
+	dir := Tea.Root + "/resources/ipdata/" + code
+	var lastVersion int64 = -1
+	lastFilename := ""
+	for _, file := range files.NewFile(dir).List() {
+		filename := file.Name()
+
+		reg := regexp.MustCompile(`^` + regexp.QuoteMeta(code) + `.(\d+)\.`)
+		if reg.MatchString(filename) { // 先查找有版本号的
+			result := reg.FindStringSubmatch(filename)
+			version := types.Int64(result[1])
+			if version > lastVersion {
+				lastVersion = version
+				lastFilename = filename
+			}
+		} else if strings.HasPrefix(filename, code+".") { // 后查找默认的
+			if lastVersion == -1 {
+				lastFilename = filename
+				lastVersion = 0
+			}
+		}
+	}
+
+	if len(lastFilename) == 0 {
+		return nil, errors.New("ip library file not found")
+	}
+
+	var libraryPtr LibraryInterface
+	switch code {
+	case serverconfigs.IPLibraryTypeIP2Region:
+		libraryPtr = &IP2RegionLibrary{}
+	default:
+		return nil, errors.New("invalid ip library code '" + code + "'")
+	}
+
+	err = libraryPtr.Load(dir + "/" + lastFilename)
+	if err != nil {
+		return nil, err
+	}
+
+	return libraryPtr, nil
+}

internal/iplibrary/manager_test.go

@@ -0,0 +1,26 @@
+package iplibrary
+
+import (
+	_ "github.com/iwind/TeaGo/bootstrap"
+	"github.com/iwind/TeaGo/dbs"
+	"testing"
+)
+
+func TestManager_Load(t *testing.T) {
+	dbs.NotifyReady()
+
+	manager := NewManager()
+	lib, err := manager.Load()
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Log(lib.Lookup("1.2.3.4"))
+	t.Log(lib.Lookup("2.3.4.5"))
+	t.Log(lib.Lookup("200.200.200.200"))
+	t.Log(lib.Lookup("202.106.0.20"))
+}
+
+func TestNewManager(t *testing.T) {
+	dbs.NotifyReady()
+	t.Log(SharedLibrary)
+}

internal/iplibrary/result.go

@@ -0,0 +1,10 @@
+package iplibrary
+
+type Result struct {
+	CityId   int64
+	Country  string
+	Region   string
+	Province string
+	City     string
+	ISP      string
+}

internal/iplibrary/updater.go

@@ -0,0 +1,124 @@
+package iplibrary
+
+import (
+	"fmt"
+	"github.com/TeaOSLab/EdgeAPI/internal/db/models"
+	"github.com/TeaOSLab/EdgeAPI/internal/errors"
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
+	"github.com/iwind/TeaGo/Tea"
+	"github.com/iwind/TeaGo/dbs"
+	"github.com/iwind/TeaGo/logs"
+	"os"
+	"time"
+)
+
+func init() {
+	dbs.OnReady(func() {
+		updater := NewUpdater()
+		updater.Start()
+	})
+}
+
+// IP库更新程序
+type Updater struct {
+}
+
+// 获取新对象
+func NewUpdater() *Updater {
+	return &Updater{}
+}
+
+// 开始更新
+func (this *Updater) Start() {
+	// 这里不需要太频繁检查更新，因为通常不需要更新IP库
+	ticker := time.NewTicker(1 * time.Hour)
+	go func() {
+		for range ticker.C {
+			err := this.loop()
+			if err != nil {
+				logs.Println("[IP_LIBRARY]" + err.Error())
+			}
+		}
+	}()
+}
+
+// 单次任务
+func (this *Updater) loop() error {
+	config, err := models.SharedSysSettingDAO.ReadGlobalConfig()
+	if err != nil {
+		return err
+	}
+	code := config.IPLibrary.Code
+	if len(code) == 0 {
+		code = serverconfigs.DefaultIPLibraryType
+	}
+	lib, err := models.SharedIPLibraryDAO.FindLatestIPLibraryWithType(code)
+	if err != nil {
+		return err
+	}
+	if lib == nil {
+		return nil
+	}
+
+	typeInfo := serverconfigs.FindIPLibraryWithType(code)
+	if typeInfo == nil {
+		return errors.New("invalid ip library code '" + code + "'")
+	}
+
+	path := Tea.Root + "/resources/ipdata/" + code + "/" + code + "." + fmt.Sprintf("%d", lib.CreatedAt) + typeInfo.GetString("ext")
+
+	// 是否已经存在
+	_, err = os.Stat(path)
+	if err == nil {
+		return nil
+	}
+
+	// 开始下载
+	chunkIds, err := models.SharedFileChunkDAO.FindAllFileChunkIds(int64(lib.FileId))
+	if err != nil {
+		return err
+	}
+	isOk := false
+
+	fp, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0666)
+	if err != nil {
+		return err
+	}
+
+	defer func() {
+		// 如果保存不成功就直接删除
+		if !isOk {
+			_ = fp.Close()
+			_ = os.Remove(path)
+		}
+	}()
+	for _, chunkId := range chunkIds {
+		chunk, err := models.SharedFileChunkDAO.FindFileChunk(chunkId)
+		if err != nil {
+			return err
+		}
+		if chunk == nil {
+			continue
+		}
+		_, err = fp.Write([]byte(chunk.Data))
+		if err != nil {
+			return err
+		}
+	}
+
+	err = fp.Close()
+	if err != nil {
+		return err
+	}
+
+	// 重新加载
+	library, err := SharedManager.Load()
+	if err != nil {
+		return err
+	}
+	SharedLibrary = library
+
+	isOk = true
+
+	return nil
+}

internal/iplibrary/updater_test.go

@@ -0,0 +1,18 @@
+package iplibrary
+
+import (
+	_ "github.com/iwind/TeaGo/bootstrap"
+	"github.com/iwind/TeaGo/dbs"
+	"testing"
+)
+
+func TestUpdater_loop(t *testing.T) {
+	dbs.NotifyReady()
+
+	updater := NewUpdater()
+	err := updater.loop()
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Log("ok")
+}