waf-platform/EdgeCommon/pkg/iplibrary/reader_maxmind.go

// 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"
	"sync"
)

// MaxMindReader MaxMind GeoIP2 Reader
type MaxMindReader struct {
	db          *geoip2.Reader
	dbASN       *geoip2.Reader // ASN 数据库（可选）
	meta        *Meta
	initialized bool
	mutex       sync.RWMutex

	// 临时文件路径，Destroy 时自动清理
	tmpFiles []string
}

// 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")
	}

	// 直接从内存字节加载，避免在 /tmp 持续生成 mmdb 临时文件。
	db, err := geoip2.FromBytes(cityDBData)
	if err != nil {
		return nil, fmt.Errorf("open MaxMind city database failed: %w", err)
	}

	reader := &MaxMindReader{
		db: db,
	}

	// 加载 ASN 数据库（可选）
	if len(asnDBData) > 0 {
		dbASN, err := geoip2.FromBytes(asnDBData)
		if err == nil {
			reader.dbASN = dbASN
		}
	}

	// 初始化 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
	}
	for _, f := range this.tmpFiles {
		os.Remove(f)
	}
	this.tmpFiles = nil
	this.initialized = false
}

// maxMindItem MaxMind 查询结果项
type maxMindItem struct {
	Record geoip2.City
	ASN    uint
	ASNOrg string
}