Files
OpenFlare/internal/apps/openflare/geoip/runtime.go
T
ryan 3328d3d121 refactor(agent): stop embedding GeoIP MMDB in agent binary
Agent ships without City/Country MMDB in the binary; Docker images COPY
databases into data_dir, bare installs seed via download on first start.
Server keeps Country-only embed for optional MaxMind control-plane use.
Also harden fetch script nonempty check and reject non-file MMDB paths.
2026-08-06 16:24:57 +08:00

123 lines
3.0 KiB
Go

// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package geoip
import (
"context"
"io/fs"
"os"
"path/filepath"
"strings"
"sync"
geodata "github.com/Rain-kl/Wavelet/internal/apps/openflare/geoip/data"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/repository"
pkggeoip "github.com/Rain-kl/Wavelet/pkg/geoip"
"github.com/Rain-kl/Wavelet/pkg/logger"
)
const (
serverMMDBRelativePath = "data/GeoLite2-Country.mmdb"
serverMMDBDirPerm = 0o750
serverMMDBFilePerm = 0o644
)
var (
runtimeOnce sync.Once
runtimeInitErr error
currentProviderMu sync.RWMutex
currentProvider string
)
// EnsureRuntimeProvider loads GeoIP provider config from SystemConfig.
func EnsureRuntimeProvider(ctx context.Context) error {
runtimeOnce.Do(func() {
runtimeInitErr = applyProviderFromSystemConfig(ctx)
})
return runtimeInitErr
}
// RefreshRuntimeProvider reapplies GeoIPProvider after config updates.
func RefreshRuntimeProvider(ctx context.Context) error {
return applyProviderFromSystemConfig(ctx)
}
func applyProviderFromSystemConfig(ctx context.Context) error {
config, err := repository.GetSystemConfigByKey(ctx, model.ConfigKeyGeoIPProvider)
if err != nil {
// 降级到默认值
return ApplyProvider(ctx, "ipinfo")
}
provider := strings.TrimSpace(config.Value)
if provider == "" {
provider = "ipinfo"
}
return ApplyProvider(ctx, provider)
}
// ApplyProvider switches the process-wide GeoIP backend.
func ApplyProvider(ctx context.Context, provider string) error {
normalized := strings.TrimSpace(strings.ToLower(provider))
if normalized == "" {
normalized = pkggeoip.ProviderDisabled
}
currentProviderMu.Lock()
if currentProvider == normalized {
currentProviderMu.Unlock()
return nil
}
currentProvider = normalized
currentProviderMu.Unlock()
if normalized == pkggeoip.ProviderMaxMind {
path, err := ensureServerMMDB()
if err != nil {
logger.WarnF(ctx, "[GeoIP] seed MaxMind database failed: %v", err)
}
if path != "" {
pkggeoip.GeoIPFilePath = path
}
}
pkggeoip.InitGeoIP(normalized)
return nil
}
func ensureServerMMDB() (string, error) {
path, err := filepath.Abs(serverMMDBRelativePath)
if err != nil {
return "", err
}
if _, err := os.Stat(path); err == nil {
return path, nil
}
if err != nil && !os.IsNotExist(err) {
return "", err
}
// Control plane: seed Country from embedded asset (no City; Agent uses disk/image).
data, err := fs.ReadFile(geodata.FS, geodata.DefaultMMDBName)
if err != nil {
return "", err
}
if err := os.MkdirAll(filepath.Dir(path), serverMMDBDirPerm); err != nil {
return "", err
}
if err := os.WriteFile(path, data, serverMMDBFilePerm); err != nil { //nolint:gosec // world-readable mmdb
return "", err
}
return path, nil
}
// ResetRuntimeForTest clears lazy-init state for unit tests.
func ResetRuntimeForTest() {
runtimeOnce = sync.Once{}
runtimeInitErr = nil
currentProviderMu.Lock()
currentProvider = ""
currentProviderMu.Unlock()
}