diff --git a/atsf_server/common/constants.go b/atsf_server/common/constants.go index 3f7d7d65..f78912a4 100644 --- a/atsf_server/common/constants.go +++ b/atsf_server/common/constants.go @@ -53,6 +53,7 @@ var NodeOfflineThreshold = 2 * time.Minute // V3 operational settings (hot-reloadable via Option table) var AgentHeartbeatInterval = 10000 // milliseconds var AgentUpdateRepo = "Rain-kl/ATSFlare" +var GeoIPProvider = "disabled" // V5 OpenResty performance settings (hot-reloadable via Option table) var OpenRestyWorkerProcesses = "auto" diff --git a/atsf_server/controller/option.go b/atsf_server/controller/option.go index d223f5b5..c1af66f5 100644 --- a/atsf_server/controller/option.go +++ b/atsf_server/controller/option.go @@ -5,6 +5,7 @@ import ( "atsflare/model" "atsflare/service" "atsflare/utils" + "atsflare/utils/geoip" "encoding/json" "fmt" "github.com/gin-gonic/gin" @@ -62,6 +63,16 @@ func validateBooleanOption(key string, value string) error { } } +func validateGeoIPOption(key string, value string) error { + if key != "GeoIPProvider" { + return nil + } + if !geoip.IsValidProvider(value) { + return fmt.Errorf("%s 仅支持 disabled、mmdb、ip-api、geojs、ipinfo", key) + } + return nil +} + func validateOpenRestyOption(key string, value string) error { trimmed := strings.TrimSpace(value) @@ -250,6 +261,13 @@ func UpdateOption(c *gin.Context) { }) return } + if err = validateGeoIPOption(option.Key, option.Value); err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } err = model.UpdateOption(option.Key, option.Value) if err != nil { c.JSON(http.StatusOK, gin.H{ diff --git a/atsf_server/go.mod b/atsf_server/go.mod index 61744ab8..890822d4 100644 --- a/atsf_server/go.mod +++ b/atsf_server/go.mod @@ -55,6 +55,7 @@ require ( github.com/mattn/go-isatty v0.0.19 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/oschwald/maxminddb-golang v1.13.1 // indirect github.com/pelletier/go-toml/v2 v2.0.8 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect diff --git a/atsf_server/go.sum b/atsf_server/go.sum index ad295b64..f16a7378 100644 --- a/atsf_server/go.sum +++ b/atsf_server/go.sum @@ -141,6 +141,8 @@ github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE= github.com/onsi/ginkgo v1.16.5/go.mod h1:+E8gABHa3K6zRBolWtd+ROzc/U5bkGt0FwiG042wbpU= github.com/onsi/gomega v1.18.1 h1:M1GfJqGRrBrrGGsbxzV5dqM2U2ApXefZCQpkukxYRLE= github.com/onsi/gomega v1.18.1/go.mod h1:0q+aL8jAiMXy9hbwj2mr5GziHiwhAIQpFmmtT5hitRs= +github.com/oschwald/maxminddb-golang v1.13.1 h1:G3wwjdN9JmIK2o/ermkHM+98oX5fS+k5MbwsmL4MRQE= +github.com/oschwald/maxminddb-golang v1.13.1/go.mod h1:K4pgV9N/GcK694KSTmVSDTODk4IsCNThNdTmnaBZ/F8= github.com/pelletier/go-toml/v2 v2.0.1/go.mod h1:r9LEWfGN8R5k0VXJ+0BkIe7MYkRdwZOjgMj2KwnJFUo= github.com/pelletier/go-toml/v2 v2.0.8 h1:0ctb6s9mE31h0/lhu+J6OPmVeDxJn+kYnJc2jZR9tGQ= github.com/pelletier/go-toml/v2 v2.0.8/go.mod h1:vuYfssBdrU2XDZ9bYydBu6t+6a6PYNcZljzZR9VXg+4= @@ -166,6 +168,7 @@ github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/stretchr/testify v1.8.3 h1:RP3t2pwF7cMEbC1dqtB6poj3niw/9gnV4Cjg5oW5gtY= github.com/stretchr/testify v1.8.3/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= +github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= github.com/swaggo/files v1.0.1 h1:J1bVJ4XHZNq0I46UU90611i9/YzdrF7x92oX1ig5IdE= github.com/swaggo/files v1.0.1/go.mod h1:0qXmMNH6sXNf+73t65aKeB+ApmgxdnkQzVTAj2uaMUg= github.com/swaggo/gin-swagger v1.6.1 h1:Ri06G4gc9N4t4k8hekMigJ9zKTFSlqj/9paAQCQs7cY= diff --git a/atsf_server/main.go b/atsf_server/main.go index 6a865447..968c9c37 100644 --- a/atsf_server/main.go +++ b/atsf_server/main.go @@ -6,6 +6,7 @@ import ( "atsflare/middleware" "atsflare/model" "atsflare/router" + "atsflare/utils/geoip" "embed" "fmt" "github.com/gin-contrib/sessions" @@ -65,6 +66,7 @@ func main() { // Initialize options model.InitOptionMap() + geoip.InitGeoIP() // Initialize HTTP server server := gin.Default() diff --git a/atsf_server/model/option.go b/atsf_server/model/option.go index 5be3dbd6..4c753a02 100644 --- a/atsf_server/model/option.go +++ b/atsf_server/model/option.go @@ -2,6 +2,7 @@ package model import ( "atsflare/common" + "atsflare/utils/geoip" "strconv" "strings" "time" @@ -54,6 +55,7 @@ func InitOptionMap() { common.OptionMap["AgentHeartbeatInterval"] = strconv.Itoa(common.AgentHeartbeatInterval) common.OptionMap["NodeOfflineThreshold"] = strconv.Itoa(int(common.NodeOfflineThreshold.Milliseconds())) common.OptionMap["AgentUpdateRepo"] = common.AgentUpdateRepo + common.OptionMap["GeoIPProvider"] = common.GeoIPProvider common.OptionMap["OpenRestyWorkerProcesses"] = common.OpenRestyWorkerProcesses common.OptionMap["OpenRestyWorkerConnections"] = strconv.Itoa(common.OpenRestyWorkerConnections) common.OptionMap["OpenRestyWorkerRlimitNofile"] = strconv.Itoa(common.OpenRestyWorkerRlimitNofile) @@ -122,8 +124,8 @@ func UpdateOption(key string, value string) error { } func updateOptionMap(key string, value string) { + shouldRefreshGeoIP := false common.OptionMapRWMutex.Lock() - defer common.OptionMapRWMutex.Unlock() if common.OptionMap == nil { common.OptionMap = make(map[string]string) } @@ -206,6 +208,11 @@ func updateOptionMap(key string, value string) { if value != "" { common.AgentUpdateRepo = value } + case "GeoIPProvider": + if geoip.IsValidProvider(value) { + common.GeoIPProvider = value + shouldRefreshGeoIP = true + } case "OpenRestyWorkerProcesses": if strings.TrimSpace(value) != "" { common.OpenRestyWorkerProcesses = value @@ -363,4 +370,8 @@ func updateOptionMap(key string, value string) { common.CriticalRateLimitDuration = v } } + common.OptionMapRWMutex.Unlock() + if shouldRefreshGeoIP { + geoip.InitGeoIP() + } } diff --git a/atsf_server/utils/geoip/emptyProvider.go b/atsf_server/utils/geoip/emptyProvider.go new file mode 100644 index 00000000..c0cfb7cb --- /dev/null +++ b/atsf_server/utils/geoip/emptyProvider.go @@ -0,0 +1,25 @@ +package geoip + +import ( + "fmt" + "net" +) + +type EmptyProvider struct{} + +func (e *EmptyProvider) Name() string { + return "EmptyProvider" +} + +func (e *EmptyProvider) Initialize() error { + return nil +} +func (e *EmptyProvider) GetGeoInfo(ip net.IP) (*GeoInfo, error) { + return nil, fmt.Errorf("you are using an empty GeoIP provider, please set a valid provider") +} +func (e *EmptyProvider) UpdateDatabase() error { + return fmt.Errorf("you are using an empty GeoIP provider, please set a valid provider") +} +func (e *EmptyProvider) Close() error { + return nil +} diff --git a/atsf_server/utils/geoip/geoip.go b/atsf_server/utils/geoip/geoip.go new file mode 100644 index 00000000..65d9b39d --- /dev/null +++ b/atsf_server/utils/geoip/geoip.go @@ -0,0 +1,204 @@ +package geoip + +import ( + "atsflare/common" + "fmt" + "log/slog" + "net" + "strings" + "sync" + "time" + "unicode" +) + +var CurrentProvider GeoIPService +var geoCache *providerCache +var providerMutex sync.RWMutex + +const ( + ProviderDisabled = "disabled" + ProviderMaxMind = "mmdb" + ProviderIPAPI = "ip-api" + ProviderGeoJS = "geojs" + ProviderIPInfo = "ipinfo" +) + +type GeoInfo struct { + ISOCode string + Name string +} + +func init() { + CurrentProvider = &EmptyProvider{} + geoCache = newProviderCache(48 * time.Hour) +} + +// GeoIPService 接口定义了获取地理位置信息的核心方法。 +type GeoIPService interface { + Name() string + GetGeoInfo(ip net.IP) (*GeoInfo, error) + UpdateDatabase() error + Close() error +} + +type cachedGeoInfo struct { + info *GeoInfo + expiresAt time.Time +} + +type providerCache struct { + mu sync.RWMutex + items map[string]cachedGeoInfo + duration time.Duration +} + +func newProviderCache(duration time.Duration) *providerCache { + return &providerCache{ + items: make(map[string]cachedGeoInfo), + duration: duration, + } +} + +func (c *providerCache) Get(key string) (*GeoInfo, bool) { + c.mu.RLock() + entry, ok := c.items[key] + c.mu.RUnlock() + if !ok { + return nil, false + } + if time.Now().After(entry.expiresAt) { + c.mu.Lock() + delete(c.items, key) + c.mu.Unlock() + return nil, false + } + return entry.info, true +} + +func (c *providerCache) Set(key string, info *GeoInfo) { + c.mu.Lock() + c.items[key] = cachedGeoInfo{ + info: info, + expiresAt: time.Now().Add(c.duration), + } + c.mu.Unlock() +} + +func (c *providerCache) Flush() { + c.mu.Lock() + c.items = make(map[string]cachedGeoInfo) + c.mu.Unlock() +} + +func GetRegionUnicodeEmoji(isoCode string) string { + if len(isoCode) != 2 { + return "" + } + isoCode = strings.ToUpper(isoCode) + + if !unicode.IsLetter(rune(isoCode[0])) || !unicode.IsLetter(rune(isoCode[1])) { + return "" + } + + rune1 := rune(0x1F1E6 + (rune(isoCode[0]) - 'A')) + rune2 := rune(0x1F1E6 + (rune(isoCode[1]) - 'A')) + return string(rune1) + string(rune2) +} + +func InitGeoIP() { + providerName := normalizeProvider(common.GeoIPProvider) + nextProvider, err := newProvider(providerName) + if err != nil { + slog.Error("initialize GeoIP provider failed", "provider", providerName, "error", err) + nextProvider = &EmptyProvider{} + } + setProvider(nextProvider) + if providerName == ProviderDisabled { + slog.Info("GeoIP provider disabled") + return + } + slog.Info("GeoIP provider configured", "provider", CurrentProvider.Name()) +} + +func GetGeoInfo(ip net.IP) (*GeoInfo, error) { + if ip == nil { + return nil, fmt.Errorf("IP address cannot be nil") + } + provider := getProvider() + cacheKey := provider.Name() + ":" + ip.String() + + if cachedInfo, found := geoCache.Get(cacheKey); found { + return cachedInfo, nil + } + + info, err := provider.GetGeoInfo(ip) + if err == nil && info != nil { + geoCache.Set(cacheKey, info) + } + return info, err +} + +func UpdateDatabase() error { + err := getProvider().UpdateDatabase() + if err == nil { + geoCache.Flush() + slog.Info("GeoIP cache cleared due to database update.") + } + return err +} + +func IsValidProvider(provider string) bool { + switch normalizeProvider(provider) { + case ProviderDisabled, ProviderMaxMind, ProviderIPAPI, ProviderGeoJS, ProviderIPInfo: + return true + default: + return false + } +} + +func normalizeProvider(provider string) string { + normalized := strings.TrimSpace(strings.ToLower(provider)) + if normalized == "" { + return ProviderDisabled + } + return normalized +} + +func newProvider(provider string) (GeoIPService, error) { + switch provider { + case ProviderDisabled: + return &EmptyProvider{}, nil + case ProviderMaxMind: + return NewMaxMindGeoIPService() + case ProviderIPAPI: + return NewIPAPIService() + case ProviderGeoJS: + return NewGeoJSService() + case ProviderIPInfo: + return NewIPInfoService() + default: + return nil, fmt.Errorf("unsupported GeoIP provider %q", provider) + } +} + +func setProvider(provider GeoIPService) { + providerMutex.Lock() + previous := CurrentProvider + CurrentProvider = provider + providerMutex.Unlock() + geoCache.Flush() + if previous != nil && previous != provider { + if err := previous.Close(); err != nil { + slog.Warn("close previous GeoIP provider failed", "error", err) + } + } +} + +func getProvider() GeoIPService { + providerMutex.RLock() + defer providerMutex.RUnlock() + if CurrentProvider == nil { + return &EmptyProvider{} + } + return CurrentProvider +} diff --git a/atsf_server/utils/geoip/geoip_test.go b/atsf_server/utils/geoip/geoip_test.go new file mode 100644 index 00000000..0e64d0af --- /dev/null +++ b/atsf_server/utils/geoip/geoip_test.go @@ -0,0 +1,81 @@ +package geoip + +import ( + "net" + "testing" +) + +type fakeProvider struct { + calls int +} + +func (f *fakeProvider) Name() string { + return "fake" +} + +func (f *fakeProvider) GetGeoInfo(ip net.IP) (*GeoInfo, error) { + f.calls++ + return &GeoInfo{ + ISOCode: "CN", + Name: "China", + }, nil +} + +func (f *fakeProvider) UpdateDatabase() error { + return nil +} + +func (f *fakeProvider) Close() error { + return nil +} + +func TestGetGeoInfoCachesByProviderAndIP(t *testing.T) { + originalProvider := CurrentProvider + geoCache.Flush() + fake := &fakeProvider{} + CurrentProvider = fake + defer func() { + CurrentProvider = originalProvider + }() + + ip := net.ParseIP("8.8.8.8") + record, err := GetGeoInfo(ip) + if err != nil { + t.Fatalf("expected nil error, got %v", err) + } + if record == nil || record.ISOCode != "CN" { + t.Fatalf("expected cached record, got %#v", record) + } + + _, err = GetGeoInfo(ip) + if err != nil { + t.Fatalf("expected nil error on second call, got %v", err) + } + if fake.calls != 1 { + t.Fatalf("expected provider to be called once, got %d", fake.calls) + } +} + +func TestUnicodeEmoji(t *testing.T) { + emoji := GetRegionUnicodeEmoji("CN") + if emoji != "🇨🇳" { + t.Errorf("expected emoji for CN, got %s", emoji) + } +} + +func TestIsValidProvider(t *testing.T) { + cases := map[string]bool{ + "disabled": true, + "mmdb": true, + "ip-api": true, + "geojs": true, + "ipinfo": true, + "unknown": false, + } + + for provider, want := range cases { + if got := IsValidProvider(provider); got != want { + t.Fatalf("provider %s validity mismatch: want %v, got %v", provider, want, got) + } + } +} diff --git a/atsf_server/utils/geoip/geojs.go b/atsf_server/utils/geoip/geojs.go new file mode 100644 index 00000000..769182c9 --- /dev/null +++ b/atsf_server/utils/geoip/geojs.go @@ -0,0 +1,80 @@ +package geoip + +import ( + "encoding/json" + "fmt" + "net" + "net/http" + "time" +) + +// GeoJSService 使用 geojs.io 服务实现 GeoIPService 接口。 +type GeoJSService struct { + Client *http.Client +} + +// geoJSResponse 定义了 geojs.io 服务返回的 JSON 响应的结构。 +// 我们只定义我们需要的字段。 +type geoJSResponse struct { + Country string `json:"country"` + CountryCode string `json:"country_code"` + // 可以根据需要添加其他字段，例如: + // City string `json:"city"` + // Region string `json:"region"` +} + +// NewGeoJSService 创建并返回一个 GeoJSService 的新实例。 +func NewGeoJSService() (*GeoJSService, error) { + return &GeoJSService{ + Client: &http.Client{ + Timeout: 5 * time.Second, // 设置一个合理的超时时间 + }, + }, nil +} + +// Name 返回服务的名称。 +func (s *GeoJSService) Name() string { + return "geojs.io" +} + +// GetGeoInfo 使用 geojs.io 服务检索给定 IP 地址的地理位置信息。 +func (s *GeoJSService) GetGeoInfo(ip net.IP) (*GeoInfo, error) { + // GeoJS 的 API 端点 + apiURL := fmt.Sprintf("https://get.geojs.io/v1/ip/geo/%s.json", ip.String()) + + resp, err := s.Client.Get(apiURL) + if err != nil { + return nil, fmt.Errorf("failed to get geo info from geojs.io: %w", err) + } + defer resp.Body.Close() + + // 检查响应状态码 + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("geojs.io returned non-200 status code: %d", resp.StatusCode) + } + + var apiResp geoJSResponse + if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil { + return nil, fmt.Errorf("failed to decode geojs.io response: %w", err) + } + + // 检查国家代码是否为空，因为 geojs 对无效/私有IP可能返回200 OK但内容为空 + if apiResp.CountryCode == "" { + return nil, fmt.Errorf("geojs.io returned empty geo info for ip: %s", ip.String()) + } + + return &GeoInfo{ + ISOCode: apiResp.CountryCode, + Name: apiResp.Country, + }, nil +} + +// UpdateDatabase 对于 geojs.io 是一个空操作，因为它是一个 Web 服务。 +func (s *GeoJSService) UpdateDatabase() error { + return nil +} + +// Close 对于 geojs.io 是一个空操作。 +func (s *GeoJSService) Close() error { + return nil +} diff --git a/atsf_server/utils/geoip/ipapi.go b/atsf_server/utils/geoip/ipapi.go new file mode 100644 index 00000000..08fb139d --- /dev/null +++ b/atsf_server/utils/geoip/ipapi.go @@ -0,0 +1,84 @@ +package geoip + +import ( + "encoding/json" + "fmt" + "net" + "net/http" + "time" +) + +// IPAPIService 使用 ip-api.com 服务实现 GeoIPService 接口。 +type IPAPIService struct { + Client *http.Client +} + +// ipAPIResponse 定义了 ip-api.com 服务返回的 JSON 响应的结构。 +type ipAPIResponse struct { + Status string `json:"status"` + Message string `json:"message"` // 当 status 为 fail 时出现 + Country string `json:"country"` + CountryCode string `json:"countryCode"` + Region string `json:"region"` + RegionName string `json:"regionName"` + City string `json:"city"` + Zip string `json:"zip"` + Lat float64 `json:"lat"` + Lon float64 `json:"lon"` + Timezone string `json:"timezone"` + ISP string `json:"isp"` + Org string `json:"org"` + As string `json:"as"` + Query string `json:"query"` +} + +func (s *IPAPIService) Name() string { + return "ip-api.com" +} + +// NewIPAPIService 创建并返回一个 IPAPIService 的新实例。 +func NewIPAPIService() (*IPAPIService, error) { + return &IPAPIService{ + Client: &http.Client{ + Timeout: 5 * time.Second, // 设置请求超时 + }, + }, nil +} + +// GetGeoInfo 使用 ip-api.com 服务检索给定 IP 地址的地理位置信息。 +func (s *IPAPIService) GetGeoInfo(ip net.IP) (*GeoInfo, error) { + // API URL, 使用 fields 参数来仅请求需要的字段 + apiURL := fmt.Sprintf("http://ip-api.com/json/%s?fields=status,message,country,countryCode", ip.String()) + + resp, err := s.Client.Get(apiURL) + if err != nil { + return nil, fmt.Errorf("failed to get geo info from ip-api.com: %w", err) + } + defer resp.Body.Close() + + var apiResp ipAPIResponse + if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil { + return nil, fmt.Errorf("failed to decode ip-api.com response: %w", err) + } + + if apiResp.Status != "success" { + return nil, fmt.Errorf("ip-api.com returned an error: %s", apiResp.Message) + } + + return &GeoInfo{ + ISOCode: apiResp.CountryCode, + Name: apiResp.Country, + }, nil +} + +// UpdateDatabase 对于 ip-api.com 是一个空操作，因为它是一个 Web 服务。 +func (s *IPAPIService) UpdateDatabase() error { + // 无需执行任何操作，因为数据由外部服务提供 + return nil +} + +// Close 对于 ip-api.com 是一个空操作，因为没有需要关闭的持久连接。 +func (s *IPAPIService) Close() error { + // 无需执行任何操作 + return nil +} diff --git a/atsf_server/utils/geoip/ipinfo.go b/atsf_server/utils/geoip/ipinfo.go new file mode 100644 index 00000000..0f1d32ed --- /dev/null +++ b/atsf_server/utils/geoip/ipinfo.go @@ -0,0 +1,91 @@ +package geoip + +import ( + "encoding/json" + "fmt" + "net" + "net/http" + "time" +) + +// IPInfoService 使用 ipinfo.io 服务实现 GeoIPService 接口。 +type IPInfoService struct { + Client *http.Client + // 每天 1000 次请求，限制由 IP 地址的所有人共享。 + // APIToken string +} + +// ipInfoResponse 定义了 ipinfo.io 服务返回的 JSON 响应的结构，只包含免费额度可用的字段。 +type ipInfoResponse struct { + IP string `json:"ip"` + Hostname string `json:"hostname"` + City string `json:"city"` + Region string `json:"region"` + Country string `json:"country"` + CountryCode string `json:"countryCode"` // ipinfo.io 返回 "country" 的 ISO 代码，这里为了与 GeoInfo 保持一致，额外添加一个 CountryCode + Loc string `json:"loc"` // Latitude,Longitude + Org string `json:"org"` + Postal string `json:"postal"` + Timezone string `json:"timezone"` +} + +// NewIPInfoService 创建并返回一个 IPInfoService 的新实例。 +func NewIPInfoService() (*IPInfoService, error) { + return &IPInfoService{ + Client: &http.Client{ + Timeout: 5 * time.Second, + }, + }, nil +} + +// Name 返回服务的名称。 +func (s *IPInfoService) Name() string { + return "ipinfo.io" +} + +// GetGeoInfo 使用 ipinfo.io 服务检索给定 IP 地址的地理位置信息。 +// 免费额度主要提供国家信息。 +func (s *IPInfoService) GetGeoInfo(ip net.IP) (*GeoInfo, error) { + // IPinfo 免费额度不需要 API token 就可以查询基本的 IP 信息。 + // API URL: https://ipinfo.io/json (查询自身IP) 或 https://ipinfo.io/YOUR_IP/json + apiURL := fmt.Sprintf("https://ipinfo.io/%s/json", ip.String()) + + resp, err := s.Client.Get(apiURL) + if err != nil { + return nil, fmt.Errorf("failed to get geo info from ipinfo.io: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("ipinfo.io returned non-200 status: %d %s", resp.StatusCode, resp.Status) + } + + var apiResp ipInfoResponse + if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil { + return nil, fmt.Errorf("failed to decode ipinfo.io response: %w", err) + } + + // IPinfo 的 "country" 字段直接返回 ISO 2-letter code，例如 "US", "CN" + // 我们需要将 "country" 字段作为 ISOCode，并尝试获取其对应的国家名称。 + // IPinfo 响应中通常不直接提供完整的国家名称，但我们可以通过 CountryCode 映射。 + // 为了简化并符合 GeoInfo 结构，我们直接使用 Country 作为 ISOCode，并尝试从 CountryCode 获取名称。 + // 实际上，IPinfo 的 'country' 字段就是 ISO 2-letter code。 + // 如果需要完整的国家名称，可能需要一个本地的 ISO 代码到名称的映射。 + // 为了与 GetRegionUnicodeEmoji 函数兼容，我们直接使用 country 作为 ISOCode。 + return &GeoInfo{ + ISOCode: apiResp.Country, // IPinfo 的 'country' 字段就是 ISO 2-letter code + Name: apiResp.Country, // 免费额度通常只提供 ISO 编码，这里暂时用 ISO 编码作为名称 + }, nil +} + +// UpdateDatabase 对于 ipinfo.io 是一个空操作，因为它是一个 Web 服务。 +func (s *IPInfoService) UpdateDatabase() error { + // 无需执行任何操作，因为数据由外部服务提供 + return nil +} + +// Close 对于 ipinfo.io 是一个空操作，因为没有需要关闭的持久连接。 +func (s *IPInfoService) Close() error { + // 无需执行任何操作 + return nil +} diff --git a/atsf_server/utils/geoip/mmdb.go b/atsf_server/utils/geoip/mmdb.go new file mode 100644 index 00000000..787d624a --- /dev/null +++ b/atsf_server/utils/geoip/mmdb.go @@ -0,0 +1,151 @@ +package geoip + +import ( + "fmt" + "io" + "log/slog" + "net" + "net/http" + "os" + "path/filepath" + "sync" + + "github.com/oschwald/maxminddb-golang" +) + +var GeoIpUrl = "https://raw.githubusercontent.com/Loyalsoldier/geoip/release/GeoLite2-Country.mmdb" +var GeoIpFilePath = "./data/GeoLite2-Country.mmdb" + +type GeoIpRecord struct { + Country struct { + ISOCode string `maxminddb:"iso_code"` + Names map[string]string `maxminddb:"names"` + } `maxminddb:"country"` +} + +type MaxMindGeoIPService struct { + maxMindDBReader *maxminddb.Reader + dbFilePath string + mu sync.RWMutex +} + +func (s *MaxMindGeoIPService) Name() string { + return "MaxMind" +} + +func NewMaxMindGeoIPService() (*MaxMindGeoIPService, error) { + service := &MaxMindGeoIPService{ + dbFilePath: GeoIpFilePath, + } + + if err := os.MkdirAll(filepath.Dir(service.dbFilePath), os.ModePerm); err != nil { + return nil, fmt.Errorf("failed to create data directory for MaxMind database: %w", err) + } + + if _, err := os.Stat(service.dbFilePath); os.IsNotExist(err) { + if err := service.UpdateDatabase(); err != nil { + return nil, fmt.Errorf("failed to download initial MaxMind database: %w", err) + } + } + + if err := service.initialize(); err != nil { + return nil, fmt.Errorf("failed to initialize MaxMind database: %w", err) + } + + return service, nil +} + +func (s *MaxMindGeoIPService) initialize() error { + s.mu.Lock() + defer s.mu.Unlock() + + if s.maxMindDBReader != nil { + _ = s.maxMindDBReader.Close() + s.maxMindDBReader = nil + } + + reader, err := maxminddb.Open(s.dbFilePath) + if err != nil { + return fmt.Errorf("error opening MaxMind database at %s: %w", s.dbFilePath, err) + } + s.maxMindDBReader = reader + return nil +} + +func (s *MaxMindGeoIPService) GetGeoInfo(ip net.IP) (*GeoInfo, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + if s.maxMindDBReader == nil { + return nil, fmt.Errorf("MaxMind database is not initialized or failed to open") + } + if ip == nil { + return nil, fmt.Errorf("IP address cannot be nil") + } + + var record GeoIpRecord + if err := s.maxMindDBReader.Lookup(ip, &record); err != nil { + return nil, fmt.Errorf("error looking up IP %s in MaxMind database: %w", ip.String(), err) + } + + geoInfo := &GeoInfo{ + ISOCode: record.Country.ISOCode, + Name: record.Country.Names["en"], + } + if geoInfo.Name == "" && geoInfo.ISOCode != "" { + geoInfo.Name = geoInfo.ISOCode + } + + return geoInfo, nil +} + +func (s *MaxMindGeoIPService) UpdateDatabase() error { + resp, err := http.Get(GeoIpUrl) + if err != nil { + return fmt.Errorf("failed to initiate MaxMind database download: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("failed to download MaxMind database: HTTP status %s", resp.Status) + } + + if err := os.MkdirAll(filepath.Dir(s.dbFilePath), os.ModePerm); err != nil { + return fmt.Errorf("failed to create data directory for MaxMind database update: %w", err) + } + + tempPath := s.dbFilePath + ".download" + out, err := os.Create(tempPath) + if err != nil { + return fmt.Errorf("failed to create MaxMind database file at %s: %w", tempPath, err) + } + defer func() { + _ = out.Close() + }() + + if _, err = io.Copy(out, resp.Body); err != nil { + return fmt.Errorf("failed to write MaxMind database file: %w", err) + } + if err = out.Close(); err != nil { + return fmt.Errorf("failed to close MaxMind database file: %w", err) + } + if err = os.Rename(tempPath, s.dbFilePath); err != nil { + return fmt.Errorf("failed to move MaxMind database file into place: %w", err) + } + + return s.initialize() +} + +func (s *MaxMindGeoIPService) Close() error { + s.mu.Lock() + defer s.mu.Unlock() + if s.maxMindDBReader != nil { + err := s.maxMindDBReader.Close() + s.maxMindDBReader = nil + if err != nil { + return fmt.Errorf("error closing MaxMind database: %w", err) + } + } + slog.Info("MaxMind GeoIP service closed.") + return nil +} diff --git a/docs/app-config.md b/docs/app-config.md index 6a1c30f7..f41a5681 100644 --- a/docs/app-config.md +++ b/docs/app-config.md @@ -108,10 +108,11 @@ volumes: | 配置项 | 作用 | 默认值 | | --- | --- | --- | -| `AgentHeartbeatInterval` | Agent 心跳间隔（毫秒） | `10000` | -| `NodeOfflineThreshold` | 节点离线判定阈值（毫秒） | `120000` | -| `AgentUpdateRepo` | Agent 自更新仓库 | `Rain-kl/ATSFlare` | -| `GlobalApiRateLimitNum` / `GlobalApiRateLimitDuration` | 全局 API 限流次数 / 时间窗口（秒） | `300` / `180` | +| `AgentHeartbeatInterval` | Agent 心跳间隔（毫秒） | `10000` | +| `NodeOfflineThreshold` | 节点离线判定阈值（毫秒） | `120000` | +| `AgentUpdateRepo` | Agent 自更新仓库 | `Rain-kl/ATSFlare` | +| `GeoIPProvider` | IP 归属解析方式；支持 `disabled`、`mmdb`、`ip-api`、`geojs`、`ipinfo` | `disabled` | +| `GlobalApiRateLimitNum` / `GlobalApiRateLimitDuration` | 全局 API 限流次数 / 时间窗口（秒） | `300` / `180` | | `GlobalWebRateLimitNum` / `GlobalWebRateLimitDuration` | 全局 Web 限流次数 / 时间窗口（秒） | `300` / `180` | | `UploadRateLimitNum` / `UploadRateLimitDuration` | 上传接口限流次数 / 时间窗口（秒） | `50` / `60` | | `DownloadRateLimitNum` / `DownloadRateLimitDuration` | 下载接口限流次数 / 时间窗口（秒） | `50` / `60` | @@ -119,8 +120,9 @@ volumes: 说明： -* 限流窗口上限不能超过 `RateLimitKeyExpirationDuration`，当前为 20 分钟 -* 限流按来源 IP 统计，若前置了 Nginx/CDN/LB，应正确透传真实客户端 IP +* 限流窗口上限不能超过 `RateLimitKeyExpirationDuration`，当前为 20 分钟 +* 限流按来源 IP 统计，若前置了 Nginx/CDN/LB，应正确透传真实客户端 IP +* `GeoIPProvider=mmdb` 时，Server 会按需下载并使用本地 MaxMind Country 数据库；其余非 `disabled` 选项会直接请求对应外部 GeoIP 服务 ### 1.2.2 第五版当前支持的 OpenResty 优化配置项