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

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