Initial commit (code only without large binaries)

EdgeCommon/pkg/iplibrary/reader_maxmind.go

@@ -0,0 +1,198 @@
+// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package iplibrary
+
+import (
+	"fmt"
+	"github.com/oschwald/geoip2-golang"
+	"net"
+	"os"
+	"path/filepath"
+	"sync"
+	"time"
+)
+
+// MaxMindReader MaxMind GeoIP2 Reader
+type MaxMindReader struct {
+	db          *geoip2.Reader
+	dbASN       *geoip2.Reader // ASN 数据库（可选）
+	meta        *Meta
+	initialized bool
+	mutex       sync.RWMutex
+}
+
+// NewMaxMindReader 创建 MaxMind Reader
+// 参数：
+//   - cityDBPath: City 数据库路径（包含 Country 和 City 信息）
+//   - asnDBPath: ASN 数据库路径（可选，用于获取 ISP 信息）
+func NewMaxMindReader(cityDBPath, asnDBPath string) (*MaxMindReader, error) {
+	if len(cityDBPath) == 0 {
+		return nil, fmt.Errorf("city database path is required")
+	}
+
+	// 打开 City 数据库（包含 Country 信息）
+	db, err := geoip2.Open(cityDBPath)
+	if err != nil {
+		return nil, fmt.Errorf("open MaxMind city database failed: %w", err)
+	}
+
+	reader := &MaxMindReader{
+		db: db,
+	}
+
+	// 打开 ASN 数据库（可选）
+	if len(asnDBPath) > 0 {
+		dbASN, err := geoip2.Open(asnDBPath)
+		if err == nil {
+			reader.dbASN = dbASN
+		}
+		// 如果 ASN 数据库打开失败，不影响主功能，只记录错误
+	}
+
+	// 初始化 Meta
+	reader.meta = reader.initMeta()
+
+	reader.initialized = true
+
+	return reader, nil
+}
+
+// NewMaxMindReaderFromBytes 从字节数据创建 MaxMind Reader
+// 参数：
+//   - cityDBData: City 数据库字节数据（包含 Country 和 City 信息）
+//   - asnDBData: ASN 数据库字节数据（可选，用于获取 ISP 信息）
+func NewMaxMindReaderFromBytes(cityDBData, asnDBData []byte) (*MaxMindReader, error) {
+	if len(cityDBData) == 0 {
+		return nil, fmt.Errorf("city database data is required")
+	}
+
+	// 创建临时文件，使用更唯一的文件名避免冲突
+	tmpDir := os.TempDir()
+	pid := os.Getpid()
+	// 使用时间戳增加唯一性，避免同一进程多次调用时的冲突
+	timestamp := time.Now().UnixNano()
+	cityTmpFile := filepath.Join(tmpDir, fmt.Sprintf("geolite2-city-%d-%d.mmdb", pid, timestamp))
+	asnTmpFile := filepath.Join(tmpDir, fmt.Sprintf("geolite2-asn-%d-%d.mmdb", pid, timestamp))
+
+	// 如果临时文件已存在，先删除（可能是之前崩溃留下的）
+	os.Remove(cityTmpFile)
+	os.Remove(asnTmpFile)
+
+	// 写入 City 数据库到临时文件
+	if err := os.WriteFile(cityTmpFile, cityDBData, 0644); err != nil {
+		return nil, fmt.Errorf("write city database to temp file failed: %w", err)
+	}
+
+	// 打开 City 数据库
+	db, err := geoip2.Open(cityTmpFile)
+	if err != nil {
+		os.Remove(cityTmpFile)
+		return nil, fmt.Errorf("open MaxMind city database failed: %w", err)
+	}
+
+	reader := &MaxMindReader{
+		db: db,
+	}
+
+	// 写入并打开 ASN 数据库（可选）
+	if len(asnDBData) > 0 {
+		if err := os.WriteFile(asnTmpFile, asnDBData, 0644); err == nil {
+			dbASN, err := geoip2.Open(asnTmpFile)
+			if err == nil {
+				reader.dbASN = dbASN
+			} else {
+				// ASN 数据库打开失败，清理临时文件但不影响主功能
+				os.Remove(asnTmpFile)
+			}
+		}
+	}
+
+	// 初始化 Meta
+	reader.meta = reader.initMeta()
+
+	reader.initialized = true
+
+	return reader, nil
+}
+
+// initMeta 初始化 Meta
+func (this *MaxMindReader) initMeta() *Meta {
+	meta := &Meta{
+		Version:   3, // MaxMind 版本
+		Code:      "maxmind",
+		Author:    "MaxMind",
+		Countries: []*Country{},
+		Provinces: []*Province{},
+		Cities:    []*City{},
+		Towns:     []*Town{},
+		Providers: []*Provider{},
+		CreatedAt: 0,
+	}
+	meta.Init()
+	return meta
+}
+
+// Lookup 查询 IP 信息
+func (this *MaxMindReader) Lookup(ip net.IP) *QueryResult {
+	if !this.initialized || this.db == nil {
+		return &QueryResult{}
+	}
+
+	this.mutex.RLock()
+	defer this.mutex.RUnlock()
+
+	// 查询 City 数据库
+	record, err := this.db.City(ip)
+	if err != nil {
+		return &QueryResult{}
+	}
+
+	// 构建查询结果
+	result := &QueryResult{
+		item: &maxMindItem{
+			Record: *record,
+		},
+		meta: this.meta,
+	}
+
+	// 如果有 ASN 数据库，查询 ASN 信息
+	if this.dbASN != nil {
+		asnRecord, err := this.dbASN.ASN(ip)
+		if err == nil {
+			if maxItem, ok := result.item.(*maxMindItem); ok {
+				maxItem.ASN = asnRecord.AutonomousSystemNumber
+				maxItem.ASNOrg = asnRecord.AutonomousSystemOrganization
+			}
+		}
+	}
+
+	return result
+}
+
+// Meta 获取 Meta 信息
+func (this *MaxMindReader) Meta() *Meta {
+	return this.meta
+}
+
+// Destroy 销毁 Reader
+func (this *MaxMindReader) Destroy() {
+	this.mutex.Lock()
+	defer this.mutex.Unlock()
+
+	if this.db != nil {
+		this.db.Close()
+		this.db = nil
+	}
+	if this.dbASN != nil {
+		this.dbASN.Close()
+		this.dbASN = nil
+	}
+	this.initialized = false
+}
+
+// maxMindItem MaxMind 查询结果项
+type maxMindItem struct {
+	Record geoip2.City
+	ASN    uint
+	ASNOrg string
+}