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