tcjlj/OpenFlare
1
0
Fork 0
mirror of https://github.com/Rain-kl/OpenFlare.git synced 2026-10-02 23:06:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(agent): embed GeoLite2 City database

Browse Source 
Initialize missing Country and City databases from embedded assets and use FyraLabs releases for periodic updates.
This commit is contained in:
ryan
2026-07-13 14:38:36 +08:00
parent a1a997bcda
commit 74ddf97b36
12 changed files with 84 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/apps/agent/config/config.go
+2 -2
View File
@@ -31,8 +31,8 @@ const (
defaultOpenRestyObservabilityPort = 18081
defaultObservabilityReplayMinutes = 15
defaultMMDBUpdateInterval = 24 * time.Hour
defaultMMDBDownloadURL = "https://raw.githubusercontent.com/Loyalsoldier/geoip/release/GeoLite2-Country.mmdb"
defaultCityMMDBDownloadURL = "https://raw.githubusercontent.com/Loyalsoldier/geoip/release/GeoLite2-City.mmdb"
defaultMMDBDownloadURL = "https://github.com/FyraLabs/geolite2/releases/latest/download/GeoLite2-Country.mmdb"
defaultCityMMDBDownloadURL = "https://github.com/FyraLabs/geolite2/releases/latest/download/GeoLite2-City.mmdb"
defaultHeartbeatInterval = 10 * time.Second
defaultRequestTimeout = 10 * time.Second
configFilePerm = 0o600
internal/apps/agent/geoipdata/GeoLite2-City.mmdb
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 63 MiB

internal/apps/agent/geoipdata/data.go
+9 -5
View File
@@ -1,12 +1,16 @@
// Package geoipdata embeds the default MaxMind GeoLite2 country database.
// Package geoipdata embeds the default MaxMind GeoLite2 databases.
package geoipdata
import "embed"
// FS holds the embedded GeoLite2-Country.mmdb database.
// FS holds the embedded GeoLite2 Country and City databases.
//
//go:embed GeoLite2-Country.mmdb
//go:embed GeoLite2-Country.mmdb GeoLite2-City.mmdb
var FS embed.FS
// DefaultMMDBName is the filename of the embedded MaxMind country database.
const DefaultMMDBName = "GeoLite2-Country.mmdb"
const (
// DefaultMMDBName is the filename of the embedded MaxMind Country database.
DefaultMMDBName = "GeoLite2-Country.mmdb"
// DefaultCityMMDBName is the filename of the embedded MaxMind City database.
DefaultCityMMDBName = "GeoLite2-City.mmdb"
)
internal/apps/agent/geoipdata/data_test.go
+38
View File
@@ -0,0 +1,38 @@
package geoipdata
import (
"io/fs"
"strings"
"testing"
"github.com/oschwald/maxminddb-golang"
)
func TestEmbeddedDatabasesAreValid(t *testing.T) {
tests := []struct {
name string
filename string
databaseTypePart string
}{
{name: "Country", filename: DefaultMMDBName, databaseTypePart: "Country"},
{name: "City", filename: DefaultCityMMDBName, databaseTypePart: "City"},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
data, err := fs.ReadFile(FS, test.filename)
if err != nil {
t.Fatalf("read embedded database: %v", err)
}
reader, err := maxminddb.FromBytes(data)
if err != nil {
t.Fatalf("open embedded database: %v", err)
}
defer reader.Close()
if !strings.Contains(reader.Metadata.DatabaseType, test.databaseTypePart) {
t.Fatalf("unexpected database type %q", reader.Metadata.DatabaseType)
}
})
}
}
internal/apps/agent/geoipupdate/updater.go
+14 -22
View File
@@ -31,9 +31,13 @@ type Updater struct {
downloadDatabase func(context.Context, string, string) error
}
// EnsureInitialDatabase seeds the MMDB file from the embedded database if it does not exist on disk.
// EnsureInitialDatabase seeds the Country MMDB file from the embedded database if it does not exist on disk.
func (u *Updater) EnsureInitialDatabase() error {
path := filepath.Clean(u.MMDBPath)
return ensureEmbeddedDatabase(u.MMDBPath, geoipdata.DefaultMMDBName, "Country")
}
func ensureEmbeddedDatabase(targetPath string, embeddedName string, databaseName string) error {
path := filepath.Clean(targetPath)
if path == "" || path == "." {
return nil
}
@@ -42,9 +46,9 @@ func (u *Updater) EnsureInitialDatabase() error {
} else if !os.IsNotExist(err) {
return fmt.Errorf("stat mmdb file failed: %w", err)
}
data, err := fs.ReadFile(geoipdata.FS, geoipdata.DefaultMMDBName)
data, err := fs.ReadFile(geoipdata.FS, embeddedName)
if err != nil {
return fmt.Errorf("read embedded mmdb failed: %w", err)
return fmt.Errorf("read embedded %s mmdb failed: %w", databaseName, err)
}
if err := os.MkdirAll(filepath.Dir(path), mmdbDirPerm); err != nil {
return fmt.Errorf("create mmdb directory failed: %w", err)
@@ -52,31 +56,19 @@ func (u *Updater) EnsureInitialDatabase() error {
if err := os.WriteFile(path, data, mmdbFilePerm); err != nil {
return fmt.Errorf("write initial mmdb failed: %w", err)
}
slog.Info("initialized GeoIP mmdb from embedded database", "path", path, "size", len(data))
slog.Info("initialized GeoIP mmdb from embedded database", "database", databaseName, "path", path, "size", len(data))
return nil
}
// EnsureInitialDatabases retains the embedded Country seed and immediately
// downloads City when it is absent so subdivision rules work before the first ticker interval.
func (u *Updater) EnsureInitialDatabases(ctx context.Context) error {
// EnsureInitialDatabases seeds both Country and City from embedded databases
// when either managed file is absent. Network downloads are reserved for the periodic updater.
func (u *Updater) EnsureInitialDatabases(_ context.Context) error {
var errs []error
if err := u.EnsureInitialDatabase(); err != nil {
errs = append(errs, err)
}
cityPath := filepath.Clean(u.CityMMDBPath)
if cityPath == "" || cityPath == "." || u.CityDownloadURL == "" {
return errors.Join(errs...)
}
if _, err := os.Stat(cityPath); err == nil {
return errors.Join(errs...)
} else if !os.IsNotExist(err) {
errs = append(errs, fmt.Errorf("stat City mmdb file failed: %w", err))
return errors.Join(errs...)
}
if err := u.download(ctx, cityPath, u.CityDownloadURL); err != nil {
errs = append(errs, fmt.Errorf("download initial City mmdb failed: %w", err))
} else {
slog.Info("initialized GeoIP City mmdb from provider", "path", cityPath)
if err := ensureEmbeddedDatabase(u.CityMMDBPath, geoipdata.DefaultCityMMDBName, "City"); err != nil {
errs = append(errs, err)
}
return errors.Join(errs...)
}
internal/apps/agent/geoipupdate/updater_test.go
+14 -12
View File
@@ -26,7 +26,7 @@ func TestEnsureInitialDatabaseCopiesEmbeddedMMDB(t *testing.T) {
}
}
func TestEnsureInitialDatabasesDownloadsMissingCity(t *testing.T) {
func TestEnsureInitialDatabasesCopiesEmbeddedCityWithoutDownload(t *testing.T) {
tempDir := t.TempDir()
countryPath := filepath.Join(tempDir, "GeoLite2-Country.mmdb")
cityPath := filepath.Join(tempDir, "GeoLite2-City.mmdb")
@@ -35,10 +35,8 @@ func TestEnsureInitialDatabasesDownloadsMissingCity(t *testing.T) {
CityMMDBPath: cityPath,
CityDownloadURL: "https://geo.example/GeoLite2-City.mmdb",
downloadDatabase: func(_ context.Context, path, downloadURL string) error {
if path != cityPath || downloadURL != "https://geo.example/GeoLite2-City.mmdb" {
t.Fatalf("unexpected initial download: %s / %s", path, downloadURL)
}
return os.WriteFile(path, []byte("city-mmdb"), 0o600)
t.Fatalf("initial embedded seed must not download %s from %s", path, downloadURL)
return nil
},
}
@@ -49,15 +47,18 @@ func TestEnsureInitialDatabasesDownloadsMissingCity(t *testing.T) {
t.Fatalf("expected embedded Country database: %v", err)
}
data, err := os.ReadFile(cityPath)
if err != nil || string(data) != "city-mmdb" {
t.Fatalf("expected downloaded City database, data=%q err=%v", data, err)
if err != nil || len(data) == 0 {
t.Fatalf("expected embedded City database, size=%d err=%v", len(data), err)
}
}
func TestEnsureInitialDatabasesKeepsCountryFallbackWhenCityDownloadFails(t *testing.T) {
func TestEnsureInitialDatabasesKeepsExistingCityWithoutDownload(t *testing.T) {
tempDir := t.TempDir()
countryPath := filepath.Join(tempDir, "GeoLite2-Country.mmdb")
cityPath := filepath.Join(tempDir, "GeoLite2-City.mmdb")
if err := os.WriteFile(cityPath, []byte("existing-city"), 0o600); err != nil {
t.Fatal(err)
}
updater := &Updater{
MMDBPath: countryPath,
CityMMDBPath: cityPath,
@@ -67,14 +68,15 @@ func TestEnsureInitialDatabasesKeepsCountryFallbackWhenCityDownloadFails(t *test
},
}
if err := updater.EnsureInitialDatabases(context.Background()); err == nil {
t.Fatal("expected City download error to be reported")
if err := updater.EnsureInitialDatabases(context.Background()); err != nil {
t.Fatalf("EnsureInitialDatabases failed: %v", err)
}
if _, err := os.Stat(countryPath); err != nil {
t.Fatalf("expected Country fallback to remain available: %v", err)
}
if _, err := os.Stat(cityPath); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("expected failed City download not to create a database, err=%v", err)
data, err := os.ReadFile(cityPath)
if err != nil || string(data) != "existing-city" {
t.Fatalf("expected existing City database to remain untouched, data=%q err=%v", data, err)
}
}