实现IP库制品管理API、自动更新程序

pkg/iplibrary/default_ip_library.go

@@ -12,18 +12,23 @@ import (
 //go:embed internal-ip-library.db
 var ipLibraryData []byte
 
-var library = NewIPLibrary()
+var defaultLibrary = NewIPLibrary()
 
-func Init() error {
-	return library.Init()
+func DefaultIPLibraryData() []byte {
+	return ipLibraryData
+}
+
+func InitDefault() error {
+	defaultLibrary.reader = nil
+	return defaultLibrary.InitFromData(ipLibraryData)
 }
 
 func Lookup(ip net.IP) *QueryResult {
-	return library.Lookup(ip)
+	return defaultLibrary.Lookup(ip)
 }
 
 func LookupIP(ip string) *QueryResult {
-	return library.LookupIP(ip)
+	return defaultLibrary.LookupIP(ip)
 }
 
 type IPLibrary struct {
@@ -34,11 +39,15 @@ func NewIPLibrary() *IPLibrary {
 	return &IPLibrary{}
 }
 
-func (this *IPLibrary) Init() error {
-	if len(ipLibraryData) == 0 || this.reader != nil {
+func NewIPLibraryWithReader(reader *Reader) *IPLibrary {
+	return &IPLibrary{reader: reader}
+}
+
+func (this *IPLibrary) InitFromData(data []byte) error {
+	if len(data) == 0 || this.reader != nil {
 		return nil
 	}
-	var reader = bytes.NewReader(ipLibraryData)
+	var reader = bytes.NewReader(data)
 	gzipReader, err := gzip.NewReader(reader)
 	if err != nil {
 		return err

pkg/iplibrary/default_ip_library_test.go

@@ -15,7 +15,7 @@ import (
 func TestIPLibrary_Init(t *testing.T) {
 	var lib = iplibrary.NewIPLibrary()
 
-	err := lib.Init()
+	err := lib.InitFromData(iplibrary.DefaultIPLibraryData())
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -29,7 +29,7 @@ func TestIPLibrary_Lookup(t *testing.T) {
 
 	var before = time.Now()
 
-	err := lib.Init()
+	err := lib.InitFromData(iplibrary.DefaultIPLibraryData())
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -59,7 +59,7 @@ func TestIPLibrary_Lookup(t *testing.T) {
 
 func TestIPLibrary_LookupIP(t *testing.T) {
 	var lib = iplibrary.NewIPLibrary()
-	err := lib.Init()
+	err := lib.InitFromData(iplibrary.DefaultIPLibraryData())
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -78,7 +78,7 @@ func TestIPLibrary_LookupIP(t *testing.T) {
 
 func BenchmarkIPLibrary_Lookup(b *testing.B) {
 	var lib = iplibrary.NewIPLibrary()
-	err := lib.Init()
+	err := lib.InitFromData(iplibrary.DefaultIPLibraryData())
 	if err != nil {
 		b.Fatal(err)
 	}

pkg/iplibrary/meta.go

@@ -34,6 +34,7 @@ type Provider struct {
 
 type Meta struct {
 	Version   int         `json:"version"` // IP库版本
+	Code      string      `json:"code"`    // 代号，用来区分不同的IP库
 	Author    string      `json:"author"`
 	Countries []*Country  `json:"countries"`
 	Provinces []*Province `json:"provinces"`

pkg/iplibrary/reader_file.go

@@ -4,6 +4,8 @@ package iplibrary
 
 import (
 	"compress/gzip"
+	"errors"
+	"io"
 	"net"
 	"os"
 )
@@ -21,9 +23,13 @@ func NewFileReader(path string) (*FileReader, error) {
 		_ = fp.Close()
 	}()
 
-	gzReader, err := gzip.NewReader(fp)
+	return NewFileDataReader(fp)
+}
+
+func NewFileDataReader(dataReader io.Reader) (*FileReader, error) {
+	gzReader, err := gzip.NewReader(dataReader)
 	if err != nil {
-		return nil, err
+		return nil, errors.New("create gzip reader failed: " + err.Error())
 	}
 
 	reader, err := NewReader(gzReader)
@@ -43,3 +49,7 @@ func (this *FileReader) Meta() *Meta {
 func (this *FileReader) Lookup(ip net.IP) *QueryResult {
 	return this.rawReader.Lookup(ip)
 }
+
+func (this *FileReader) RawReader() *Reader {
+	return this.rawReader
+}

pkg/iplibrary/updater.go

@@ -0,0 +1,268 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package iplibrary
+
+import (
+	"errors"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"io"
+	"os"
+	"strings"
+	"time"
+)
+
+type UpdaterSource interface {
+	// DataDir 文件目录
+	DataDir() string
+
+	// FindLatestFile 检查最新的IP库文件
+	FindLatestFile() (code string, fileId int64, err error)
+
+	// DownloadFile 下载文件
+	DownloadFile(fileId int64, writer io.Writer) error
+
+	// LogInfo 普通日志
+	LogInfo(message string)
+
+	// LogError 错误日志
+	LogError(err error)
+}
+
+type Updater struct {
+	source UpdaterSource
+
+	currentCode string
+	ticker      *time.Ticker
+
+	isUpdating bool
+}
+
+func NewUpdater(source UpdaterSource, interval time.Duration) *Updater {
+	return &Updater{
+		source: source,
+		ticker: time.NewTicker(interval),
+	}
+}
+
+func (this *Updater) Start() {
+	// 初始化
+	err := this.Init()
+	if err != nil {
+		this.source.LogError(err)
+	}
+
+	// 先运行一次
+	err = this.Loop()
+	if err != nil {
+		this.source.LogError(err)
+	}
+
+	// 开始定时运行
+	for range this.ticker.C {
+		err = this.Loop()
+		if err != nil {
+			this.source.LogError(err)
+		}
+	}
+}
+
+func (this *Updater) Init() error {
+	// 检查当前正在使用的IP库
+	var path = this.source.DataDir() + "/ip-library.db"
+	fp, err := os.Open(path)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return nil
+		}
+
+		return errors.New("read ip library file failed '" + err.Error() + "'")
+	}
+	defer func() {
+		_ = fp.Close()
+	}()
+
+	return this.loadFile(fp)
+}
+
+func (this *Updater) Loop() error {
+	if this.isUpdating {
+		return nil
+	}
+
+	this.isUpdating = true
+
+	defer func() {
+		this.isUpdating = false
+	}()
+
+	code, fileId, err := this.source.FindLatestFile()
+	if err != nil {
+		// 不提示连接错误
+		if this.isConnError(err) {
+			return nil
+		}
+		return err
+	}
+	if len(code) == 0 || fileId <= 0 {
+		// 还原到内置IP库
+		if len(this.currentCode) > 0 {
+			this.currentCode = ""
+			this.source.LogInfo("resetting to default ip library ...")
+
+			var defaultPath = this.source.DataDir() + "/ip-library.db"
+			_, err = os.Stat(defaultPath)
+			if err == nil {
+				err = os.Remove(defaultPath)
+				if err != nil {
+					this.source.LogError(errors.New("can not remove default 'ip-library.db'"))
+				}
+			}
+
+			err = InitDefault()
+			if err != nil {
+				this.source.LogError(errors.New("initialize default ip library failed: " + err.Error()))
+			}
+		}
+
+		return nil
+	}
+
+	// 下载
+	if this.currentCode == code {
+		// 不再重复下载
+		return nil
+	}
+
+	// 检查是否存在
+	var dir = this.source.DataDir()
+	var path = dir + "/ip-" + code + ".db"
+	stat, err := os.Stat(path)
+	if err == nil && !stat.IsDir() && stat.Size() > 0 {
+		fp, err := os.Open(path)
+		if err != nil {
+			return err
+		}
+
+		defer func() {
+			_ = fp.Close()
+		}()
+
+		err = this.loadFile(fp)
+		if err != nil {
+			// 尝试删除
+			_ = os.Remove(path)
+		} else {
+			this.currentCode = code
+
+			// 拷贝到 ip-library.db
+			err = this.createDefaultFile(path, dir)
+			if err != nil {
+				this.source.LogError(err)
+			}
+		}
+		return err
+	}
+
+	// write to file
+	fp, err := os.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0666)
+	if err != nil {
+		return errors.New("create ip library file failed: " + err.Error())
+	}
+
+	var isOk = false
+	defer func() {
+		if !isOk {
+			_ = os.Remove(path)
+		}
+	}()
+
+	err = this.source.DownloadFile(fileId, fp)
+	if err != nil {
+		_ = fp.Close()
+		return err
+	}
+	err = fp.Close()
+	if err != nil {
+		return nil
+	}
+
+	// load library from file
+	fp, err = os.Open(path)
+	if err != nil {
+		return nil
+	}
+	err = this.loadFile(fp)
+	_ = fp.Close()
+	if err != nil {
+		return errors.New("load file failed: " + err.Error())
+	}
+
+	isOk = true
+	this.currentCode = code
+
+	// 拷贝到 ip-library.db
+	err = this.createDefaultFile(path, dir)
+	if err != nil {
+		this.source.LogError(err)
+	}
+
+	return nil
+}
+
+func (this *Updater) loadFile(fp *os.File) error {
+	this.source.LogInfo("load ip library from '" + fp.Name() + "' ...")
+
+	fileReader, err := NewFileDataReader(fp)
+	if err != nil {
+		return errors.New("load ip library from reader failed: " + err.Error())
+	}
+
+	var reader = fileReader.RawReader()
+	defaultLibrary = NewIPLibraryWithReader(reader)
+	this.currentCode = reader.Meta().Code
+	return nil
+}
+
+func (this *Updater) createDefaultFile(sourcePath string, dir string) error {
+	sourceFp, err := os.Open(sourcePath)
+	if err != nil {
+		return errors.New("prepare to copy file to 'ip-library.db' failed: " + err.Error())
+	}
+	defer func() {
+		_ = sourceFp.Close()
+	}()
+
+	dstFp, err := os.Create(dir + "/ip-library.db")
+	if err != nil {
+		return errors.New("prepare to copy file to 'ip-library.db' failed: " + err.Error())
+	}
+	defer func() {
+		_ = dstFp.Close()
+	}()
+	_, err = io.Copy(dstFp, sourceFp)
+	if err != nil {
+		return errors.New("copy file to 'ip-library.db' failed: " + err.Error())
+	}
+	return nil
+}
+
+// isConnError 是否为连接错误
+func (this *Updater) isConnError(err error) bool {
+	if err == nil {
+		return false
+	}
+
+	// 检查是否为连接错误
+	statusErr, ok := status.FromError(err)
+	if ok {
+		var errorCode = statusErr.Code()
+		return errorCode == codes.Unavailable || errorCode == codes.Canceled
+	}
+
+	if strings.Contains(err.Error(), "code = Canceled") {
+		return true
+	}
+
+	return false
+}

pkg/iplibrary/updater_test.go

@@ -0,0 +1,53 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package iplibrary_test
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/iplibrary"
+	"github.com/iwind/TeaGo/Tea"
+	_ "github.com/iwind/TeaGo/bootstrap"
+	"io"
+	"testing"
+	"time"
+)
+
+type updaterSource struct {
+	t *testing.T
+}
+
+func (this *updaterSource) DataDir() string {
+	return Tea.Root + "/data"
+}
+
+func (this *updaterSource) FindLatestFile() (code string, fileId int64, err error) {
+	return "CODE", 1, nil
+}
+
+func (this *updaterSource) DownloadFile(fileId int64, writer io.Writer) error {
+	this.t.Log("downloading file:", fileId, "writer:", writer)
+	_, err := writer.Write(iplibrary.DefaultIPLibraryData())
+	return err
+}
+
+func (this *updaterSource) LogInfo(message string) {
+	this.t.Log(message)
+}
+
+func (this *updaterSource) LogError(err error) {
+	this.t.Fatal(err)
+}
+
+func TestNewUpdater(t *testing.T) {
+	var updater = iplibrary.NewUpdater(&updaterSource{
+		t: t,
+	}, 1*time.Minute)
+	err := updater.Init()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = updater.Loop()
+	if err != nil {
+		t.Fatal(err)
+	}
+}