137 lines
2.7 KiB
Go
137 lines
2.7 KiB
Go
package geoip
|
|
|
|
import (
|
|
"fmt"
|
|
"net"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/oschwald/geoip2-golang"
|
|
)
|
|
|
|
const cacheTTL = 30 * time.Minute
|
|
|
|
type Location struct {
|
|
CountryCode string
|
|
Country string
|
|
Province string
|
|
City string
|
|
Region string
|
|
}
|
|
|
|
type cacheEntry struct {
|
|
location Location
|
|
found bool
|
|
expireAt time.Time
|
|
}
|
|
|
|
var (
|
|
reader *geoip2.Reader
|
|
cacheMu sync.Mutex
|
|
cache = map[string]cacheEntry{}
|
|
cacheCap = 10_000
|
|
)
|
|
|
|
// Init opens an optional MaxMind GeoLite2/GeoIP2 City database.
|
|
func Init(path string) error {
|
|
Close()
|
|
path = strings.TrimSpace(path)
|
|
if path == "" {
|
|
return nil
|
|
}
|
|
opened, err := geoip2.Open(path)
|
|
if err != nil {
|
|
return fmt.Errorf("打开 GeoIP 数据库失败: %w", err)
|
|
}
|
|
reader = opened
|
|
return nil
|
|
}
|
|
|
|
func Close() {
|
|
if reader != nil {
|
|
_ = reader.Close()
|
|
reader = nil
|
|
}
|
|
cacheMu.Lock()
|
|
cache = map[string]cacheEntry{}
|
|
cacheMu.Unlock()
|
|
}
|
|
|
|
func Enabled() bool {
|
|
return reader != nil
|
|
}
|
|
|
|
func Lookup(ip string) (Location, bool) {
|
|
parsed := net.ParseIP(strings.TrimSpace(ip))
|
|
if parsed == nil || reader == nil {
|
|
return Location{}, false
|
|
}
|
|
|
|
now := time.Now()
|
|
cacheMu.Lock()
|
|
if cached, ok := cache[ip]; ok && now.Before(cached.expireAt) {
|
|
cacheMu.Unlock()
|
|
return cached.location, cached.found
|
|
}
|
|
cacheMu.Unlock()
|
|
|
|
city, err := reader.City(parsed)
|
|
if err != nil {
|
|
return Location{}, false
|
|
}
|
|
location := Location{
|
|
CountryCode: strings.ToUpper(city.Country.IsoCode),
|
|
Country: localizedName(city.Country.Names),
|
|
City: localizedName(city.City.Names),
|
|
}
|
|
if len(city.Subdivisions) > 0 {
|
|
location.Province = localizedName(city.Subdivisions[len(city.Subdivisions)-1].Names)
|
|
}
|
|
location.Region = buildRegion(location)
|
|
found := location.Country != "" || location.Province != "" || location.City != ""
|
|
|
|
cacheMu.Lock()
|
|
if len(cache) >= cacheCap {
|
|
for key, item := range cache {
|
|
if now.After(item.expireAt) {
|
|
delete(cache, key)
|
|
}
|
|
}
|
|
}
|
|
if len(cache) < cacheCap {
|
|
cache[ip] = cacheEntry{location: location, found: found, expireAt: now.Add(cacheTTL)}
|
|
}
|
|
cacheMu.Unlock()
|
|
return location, found
|
|
}
|
|
|
|
func localizedName(names map[string]string) string {
|
|
for _, key := range []string{"zh-CN", "zh", "en"} {
|
|
if value := strings.TrimSpace(names[key]); value != "" {
|
|
return value
|
|
}
|
|
}
|
|
for _, value := range names {
|
|
if value = strings.TrimSpace(value); value != "" {
|
|
return value
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func buildRegion(location Location) string {
|
|
parts := make([]string, 0, 3)
|
|
for _, value := range []string{location.Country, location.Province, location.City} {
|
|
value = strings.TrimSpace(value)
|
|
if value == "" {
|
|
continue
|
|
}
|
|
if len(parts) > 0 && parts[len(parts)-1] == value {
|
|
continue
|
|
}
|
|
parts = append(parts, value)
|
|
}
|
|
return strings.Join(parts, " ")
|
|
}
|