sub-store/internal/geoip/geoip.go

package geoip

import (
	"fmt"
	"net"
	"os"
	"path/filepath"
	"sync"

	"github.com/oschwald/maxminddb-golang"
)

// GeoResult holds the geographic lookup result for an IP.
type GeoResult struct {
	CountryCode  string // ISO 3166-1 alpha-2, e.g. "JP"
	CountryName  string // English country name, e.g. "Japan"
	CountryNameCN string // Chinese country name, e.g. "日本"
	CityName     string // English city name (may be empty)
	CityNameCN   string // Chinese city name (may be empty)
}

var (
	mu       sync.Mutex
	reader   *maxminddb.Reader
	loaded   bool
	mmdbPath string
)

// countryNameCN maps ISO country codes to Chinese names.
// Covers all 249 ISO 3166-1 countries.
var countryNameCN = map[string]string{
	"AD": "安道尔", "AE": "阿联酋", "AF": "阿富汗", "AG": "安提瓜", "AI": "安圭拉",
	"AL": "阿尔巴尼亚", "AM": "亚美尼亚", "AO": "安哥拉", "AQ": "南极洲", "AR": "阿根廷",
	"AS": "美属萨摩亚", "AT": "奥地利", "AU": "澳大利亚", "AW": "阿鲁巴", "AX": "奥兰群岛",
	"AZ": "阿塞拜疆", "BA": "波黑", "BB": "巴巴多斯", "BD": "孟加拉国", "BE": "比利时",
	"BF": "布基纳法索", "BG": "保加利亚", "BH": "巴林", "BI": "布隆迪", "BJ": "贝宁",
	"BL": "圣巴泰勒米", "BM": "百慕大", "BN": "文莱", "BO": "玻利维亚", "BQ": "荷兰加勒比",
	"BR": "巴西", "BS": "巴哈马", "BT": "不丹", "BV": "布韦岛", "BW": "博茨瓦纳",
	"BY": "白俄罗斯", "BZ": "伯利兹", "CA": "加拿大", "CC": "科科斯群岛", "CD": "刚果(金)",
	"CF": "中非", "CG": "刚果(布)", "CH": "瑞士", "CI": "科特迪瓦", "CK": "库克群岛",
	"CL": "智利", "CM": "喀麦隆", "CN": "中国", "CO": "哥伦比亚", "CR": "哥斯达黎加",
	"CU": "古巴", "CV": "佛得角", "CW": "库拉索", "CX": "圣诞岛", "CY": "塞浦路斯",
	"CZ": "捷克", "DE": "德国", "DJ": "吉布提", "DK": "丹麦", "DM": "多米尼克",
	"DO": "多米尼加", "DZ": "阿尔及利亚", "EC": "厄瓜多尔", "EE": "爱沙尼亚", "EG": "埃及",
	"EH": "西撒哈拉", "ER": "厄立特里亚", "ES": "西班牙", "ET": "埃塞俄比亚", "FI": "芬兰",
	"FJ": "斐济", "FK": "福克兰群岛", "FM": "密克罗尼西亚", "FO": "法罗群岛", "FR": "法国",
	"GA": "加蓬", "GB": "英国", "GD": "格林纳达", "GE": "格鲁吉亚", "GF": "法属圭亚那",
	"GG": "根西岛", "GH": "加纳", "GI": "直布罗陀", "GL": "格陵兰", "GM": "冈比亚",
	"GN": "几内亚", "GP": "瓜德罗普", "GQ": "赤道几内亚", "GR": "希腊", "GS": "南乔治亚",
	"GT": "危地马拉", "GU": "关岛", "GW": "几内亚比绍", "GY": "圭亚那", "HK": "香港",
	"HM": "赫德岛", "HN": "洪都拉斯", "HR": "克罗地亚", "HT": "海地", "HU": "匈牙利",
	"ID": "印度尼西亚", "IE": "爱尔兰", "IL": "以色列", "IM": "马恩岛", "IN": "印度",
	"IO": "英属印度洋领地", "IQ": "伊拉克", "IR": "伊朗", "IS": "冰岛", "IT": "意大利",
	"JE": "泽西岛", "JM": "牙买加", "JO": "约旦", "JP": "日本", "KE": "肯尼亚",
	"KG": "吉尔吉斯斯坦", "KH": "柬埔寨", "KI": "基里巴斯", "KM": "科摩罗", "KN": "圣基茨",
	"KP": "朝鲜", "KR": "韩国", "KW": "科威特", "KY": "开曼群岛", "KZ": "哈萨克斯坦",
	"LA": "老挝", "LB": "黎巴嫩", "LC": "圣卢西亚", "LI": "列支敦士登", "LK": "斯里兰卡",
	"LR": "利比里亚", "LS": "莱索托", "LT": "立陶宛", "LU": "卢森堡", "LV": "拉脱维亚",
	"LY": "利比亚", "MA": "摩洛哥", "MC": "摩纳哥", "MD": "摩尔多瓦", "ME": "黑山",
	"MF": "法属圣马丁", "MG": "马达加斯加", "MH": "马绍尔群岛", "MK": "北马其顿", "ML": "马里",
	"MM": "缅甸", "MN": "蒙古", "MO": "澳门", "MP": "北马里亚纳", "MQ": "马提尼克",
	"MR": "毛里塔尼亚", "MS": "蒙特塞拉特", "MT": "马耳他", "MU": "毛里求斯", "MV": "马尔代夫",
	"MW": "马拉维", "MX": "墨西哥", "MY": "马来西亚", "MZ": "莫桑比克", "NA": "纳米比亚",
	"NC": "新喀里多尼亚", "NE": "尼日尔", "NF": "诺福克岛", "NG": "尼日利亚", "NI": "尼加拉瓜",
	"NL": "荷兰", "NO": "挪威", "NP": "尼泊尔", "NR": "瑙鲁", "NU": "纽埃",
	"NZ": "新西兰", "OM": "阿曼", "PA": "巴拿马", "PE": "秘鲁", "PF": "法属波利尼西亚",
	"PG": "巴布亚新几内亚", "PH": "菲律宾", "PK": "巴基斯坦", "PL": "波兰", "PM": "圣皮埃尔",
	"PN": "皮特凯恩", "PR": "波多黎各", "PS": "巴勒斯坦", "PT": "葡萄牙", "PW": "帕劳",
	"PY": "巴拉圭", "QA": "卡塔尔", "RE": "留尼汪", "RO": "罗马尼亚", "RS": "塞尔维亚",
	"RU": "俄罗斯", "RW": "卢旺达", "SA": "沙特", "SB": "所罗门群岛", "SC": "塞舌尔",
	"SD": "苏丹", "SE": "瑞典", "SG": "新加坡", "SH": "圣赫勒拿", "SI": "斯洛文尼亚",
	"SJ": "斯瓦尔巴", "SK": "斯洛伐克", "SL": "塞拉利昂", "SM": "圣马力诺", "SN": "塞内加尔",
	"SO": "索马里", "SR": "苏里南", "SS": "南苏丹", "ST": "圣多美", "SV": "萨尔瓦多",
	"SX": "圣马丁", "SY": "叙利亚", "SZ": "斯威士兰", "TC": "特克斯群岛", "TD": "乍得",
	"TF": "法属南方领地", "TG": "多哥", "TH": "泰国", "TJ": "塔吉克斯坦", "TK": "托克劳",
	"TL": "东帝汶", "TM": "土库曼斯坦", "TN": "突尼斯", "TO": "汤加", "TR": "土耳其",
	"TT": "特立尼达", "TV": "图瓦卢", "TW": "台湾", "TZ": "坦桑尼亚", "UA": "乌克兰",
	"UG": "乌干达", "UM": "美属小群岛", "US": "美国", "UY": "乌拉圭", "UZ": "乌兹别克斯坦",
	"VA": "梵蒂冈", "VC": "圣文森特", "VE": "委内瑞拉", "VG": "英属维尔京", "VI": "美属维尔京",
	"VN": "越南", "VU": "瓦努阿图", "WF": "瓦利斯", "WS": "萨摩亚", "YE": "也门",
	"YT": "马约特", "ZA": "南非", "ZM": "赞比亚", "ZW": "津巴布韦",
}

// countryCodeToEmoji converts an ISO 3166-1 alpha-2 code to a flag emoji.
func countryCodeToEmoji(code string) string {
	if len(code) != 2 {
		return "🏳️"
	}
	c1 := code[0]
	c2 := code[1]
	if c1 < 'A' || c1 > 'Z' || c2 < 'A' || c2 > 'Z' {
		return "🏳️"
	}
	// Regional indicator symbols: A=0x1F1E6, so subtract 'A' and add 0x1F1E6
	r1 := rune(0x1F1E6) + rune(c1-'A')
	r2 := rune(0x1F1E6) + rune(c2-'A')
	return string(r1) + string(r2)
}

// CountryNameCN returns the Chinese name for an ISO country code.
func CountryNameCN(code string) string {
	if name, ok := countryNameCN[code]; ok {
		return name
	}
	return ""
}

// CountryEmoji returns the flag emoji for an ISO country code.
func CountryEmoji(code string) string {
	return countryCodeToEmoji(code)
}

// SetMmdbPath sets the path to the GeoLite2-City mmdb file.
// Must be called before Lookup.
func SetMmdbPath(path string) {
	mu.Lock()
	defer mu.Unlock()
	if mmdbPath != path {
		mmdbPath = path
		loaded = false
	}
}

// ensureLoaded opens the mmdb reader on first use.
func ensureLoaded() error {
	mu.Lock()
	defer mu.Unlock()
	if loaded {
		return nil
	}
	if mmdbPath == "" {
		mmdbPath = filepath.Join("data", "GeoLite2-City.mmdb")
	}
	if _, err := os.Stat(mmdbPath); err != nil {
		return fmt.Errorf("GeoIP database not found at %s: %w", mmdbPath, err)
	}
	r, err := maxminddb.Open(mmdbPath)
	if err != nil {
		return fmt.Errorf("failed to open GeoIP database: %w", err)
	}
	reader = r
	loaded = true
	return nil
}

// Lookup returns geographic info for an IP address.
// Returns nil if the database is not available or the IP is not found.
func Lookup(ipStr string) *GeoResult {
	if ipStr == "" {
		return nil
	}
	ip := net.ParseIP(ipStr)
	if ip == nil {
		return nil
	}
	if err := ensureLoaded(); err != nil {
		return nil
	}

	var record struct {
		Country struct {
			IsoCode string `maxminddb:"iso_code"`
			Names   map[string]string `maxminddb:"names"`
		} `maxminddb:"country"`
		City struct {
			Names map[string]string `maxminddb:"names"`
		} `maxminddb:"city"`
	}

	if err := reader.Lookup(ip, &record); err != nil {
		return nil
	}

	code := record.Country.IsoCode
	if code == "" {
		return nil
	}

	result := &GeoResult{
		CountryCode:   code,
		CountryName:   record.Country.Names["en"],
		CountryNameCN: record.Country.Names["zh-CN"],
		CityName:      record.City.Names["en"],
		CityNameCN:    record.City.Names["zh-CN"],
	}

	// Fallback: if zh-CN name is missing from mmdb, use our built-in map
	if result.CountryNameCN == "" {
		result.CountryNameCN = CountryNameCN(code)
	}

	return result
}

// LookupHost resolves a hostname to its first IP and then looks up geo info.
func LookupHost(host string) *GeoResult {
	// If it's already an IP, lookup directly
	if ip := net.ParseIP(host); ip != nil {
		return Lookup(ip.String())
	}
	// Resolve domain
	ips, err := net.LookupIP(host)
	if err != nil || len(ips) == 0 {
		return nil
	}
	return Lookup(ips[0].String())
}