Initial commit (code only without large binaries)

EdgeCommon/pkg/iplibrary/default_ip_library.go

@@ -0,0 +1,178 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package iplibrary
+
+import (
+	_ "embed"
+	"fmt"
+	"github.com/iwind/TeaGo/Tea"
+	"github.com/iwind/TeaGo/logs"
+	"net"
+	"os"
+	"sync"
+)
+
+//go:embed GeoLite2-City.mmdb
+var maxMindCityDBData []byte
+
+//go:embed GeoLite2-ASN.mmdb
+var maxMindASNDBData []byte
+
+var defaultLibrary *IPLibrary
+var commonLibrary *IPLibrary
+
+var libraryLocker = &sync.Mutex{} // 为了保持加载顺序性
+
+// InitDefault 加载默认的IP库
+// 优先使用后台上传的 MaxMind 文件，如果没有上传则使用代码中嵌入的 MaxMind 库
+func InitDefault() error {
+	libraryLocker.Lock()
+	defer libraryLocker.Unlock()
+
+	// 1. 优先检查后台上传的 MaxMind 文件（data/iplibrary/ 目录）
+	iplibDir := Tea.Root + "/data/iplibrary"
+	cityDBPath := iplibDir + "/maxmind-city.mmdb"
+	asnDBPath := iplibDir + "/maxmind-asn.mmdb"
+
+	// 检查上传的文件是否存在
+	uploadedFileExists := false
+	if _, err := os.Stat(cityDBPath); err == nil {
+		uploadedFileExists = true
+	}
+
+	// 如果 commonLibrary 已存在，需要检查是否应该重新加载
+	if commonLibrary != nil {
+		if commonLibrary.reader != nil {
+			// 检查是否是 MaxMindReader
+			_, isMaxMind := commonLibrary.reader.(*MaxMindReader)
+			if isMaxMind {
+				// 如果之前使用的是上传的文件，但现在文件不存在了，需要重新加载
+				// 如果之前使用的是嵌入的库，且现在也没有上传的文件，可以继续使用
+				// 这里通过检查文件是否存在来判断：如果文件存在，且库已加载，直接使用
+				if uploadedFileExists {
+					defaultLibrary = commonLibrary
+					return nil
+				}
+				// 文件不存在了，需要重新加载（使用嵌入的库）
+				commonLibrary.Destroy()
+				commonLibrary = nil
+				defaultLibrary = nil
+			} else {
+				// 不是 MaxMind 库，销毁并重新初始化
+				commonLibrary.Destroy()
+				commonLibrary = nil
+				defaultLibrary = nil
+			}
+		} else {
+			commonLibrary = nil
+			defaultLibrary = nil
+		}
+	}
+
+	if uploadedFileExists {
+		// 检查 ASN 文件是否存在
+		asnPath := ""
+		if _, err := os.Stat(asnDBPath); err == nil {
+			asnPath = asnDBPath
+		}
+		reader, err := NewMaxMindReader(cityDBPath, asnPath)
+		if err == nil {
+			defaultLibrary = NewIPLibraryWithReader(reader)
+			commonLibrary = defaultLibrary
+			logs.Println("[IP_LIBRARY]uploaded MaxMind database loaded successfully from: " + cityDBPath)
+			return nil
+		}
+		// 上传的文件加载失败，继续使用嵌入的库
+		logs.Println("[IP_LIBRARY]failed to load uploaded MaxMind database: " + err.Error() + ", will use embedded database")
+	}
+
+	// 2. 使用嵌入的 MaxMind 数据库作为默认
+	if len(maxMindCityDBData) == 0 {
+		return fmt.Errorf("embedded MaxMind database is empty (this should not happen if build is correct), please upload MaxMind database file or rebuild the application")
+	}
+
+	reader, err := NewMaxMindReaderFromBytes(maxMindCityDBData, maxMindASNDBData)
+	if err != nil {
+		return fmt.Errorf("failed to load embedded MaxMind database: %w", err)
+	}
+
+	defaultLibrary = NewIPLibraryWithReader(reader)
+	commonLibrary = defaultLibrary
+	logs.Println("[IP_LIBRARY]embedded MaxMind database loaded successfully (size: " + fmt.Sprintf("%d", len(maxMindCityDBData)) + " bytes)")
+	return nil
+}
+
+// Lookup 查询IP信息
+func Lookup(ip net.IP) *QueryResult {
+	if defaultLibrary == nil {
+		// 如果 defaultLibrary 未初始化，尝试初始化
+		_ = InitDefault()
+		if defaultLibrary == nil {
+			return &QueryResult{}
+		}
+	}
+	return defaultLibrary.Lookup(ip)
+}
+
+// LookupIP 查询IP信息
+func LookupIP(ip string) *QueryResult {
+	if defaultLibrary == nil {
+		// 如果 defaultLibrary 未初始化，尝试初始化
+		_ = InitDefault()
+		if defaultLibrary == nil {
+			return &QueryResult{}
+		}
+	}
+	return defaultLibrary.LookupIP(ip)
+}
+
+// LookupIPSummaries 查询一组IP对应的区域描述
+func LookupIPSummaries(ipList []string) map[string]string /** ip => summary **/ {
+	var result = map[string]string{}
+	for _, ip := range ipList {
+		var region = LookupIP(ip)
+		if region != nil && region.IsOk() {
+			result[ip] = region.Summary()
+		}
+	}
+	return result
+}
+
+type IPLibrary struct {
+	reader ReaderInterface
+}
+
+func NewIPLibrary() *IPLibrary {
+	return &IPLibrary{}
+}
+
+func NewIPLibraryWithReader(reader ReaderInterface) *IPLibrary {
+	return &IPLibrary{reader: reader}
+}
+
+func (this *IPLibrary) Lookup(ip net.IP) *QueryResult {
+	if this.reader == nil {
+		return &QueryResult{}
+	}
+
+	var result = this.reader.Lookup(ip)
+	if result == nil {
+		result = &QueryResult{}
+	}
+
+	return result
+}
+
+func (this *IPLibrary) LookupIP(ip string) *QueryResult {
+	if this.reader == nil {
+		return &QueryResult{}
+	}
+	return this.Lookup(net.ParseIP(ip))
+}
+
+func (this *IPLibrary) Destroy() {
+	if this.reader != nil {
+		this.reader.Destroy()
+		this.reader = nil
+	}
+}