diff --git a/.gitignore b/.gitignore index aaadf73..c5dfab1 100644 --- a/.gitignore +++ b/.gitignore @@ -28,5 +28,5 @@ go.work.sum .env # Editor/IDE -# .idea/ # .vscode/ +.idea/ diff --git a/api/geolite_asn.go b/api/geolite_asn.go index 9f6990c..07b2bce 100644 --- a/api/geolite_asn.go +++ b/api/geolite_asn.go @@ -1,11 +1,11 @@ package api import ( + "context" "fmt" "log" "net" - - "github.com/oschwald/maxminddb-golang/v2" + "time" ) type AsnRecord struct { @@ -14,20 +14,26 @@ type AsnRecord struct { } type AsnReader struct { - db *maxminddb.Reader + db *ReloadableGeoIPDB } func NewAsnReader(path string) (*AsnReader, error) { - db, err := maxminddb.Open(path) + db, err := NewReloadableGeoIPDB(path) if err != nil { return nil, fmt.Errorf("open asn mmdb: %w", err) } - log.Printf("asn mmdb type: %s", db.Metadata.DatabaseType) + log.Printf("asn mmdb type: %s", db.DatabaseType()) return &AsnReader{db: db}, nil } +// StartWatcher polls the mmdb file and hot reloads it on change. Blocks until +// ctx is cancelled. +func (a *AsnReader) StartWatcher(ctx context.Context, interval time.Duration) { + a.db.StartWatcher(ctx, interval) +} + func (a *AsnReader) Close() error { return a.db.Close() } func (a *AsnReader) Enrich(ip net.IP, out *LookupResult) error { @@ -37,7 +43,7 @@ func (a *AsnReader) Enrich(ip net.IP, out *LookupResult) error { } var rec AsnRecord - if err := a.db.Lookup(addr).Decode(&rec); err != nil { + if err := a.db.Lookup(addr, &rec); err != nil { return err } diff --git a/api/geolite_city.go b/api/geolite_city.go index 4822c7f..43b4fad 100644 --- a/api/geolite_city.go +++ b/api/geolite_city.go @@ -1,29 +1,34 @@ package api import ( + "context" "fmt" "log" "net" "time" - - "github.com/oschwald/maxminddb-golang/v2" ) type CityReader struct { - db *maxminddb.Reader + db *ReloadableGeoIPDB } func NewCityReader(path string) (*CityReader, error) { - db, err := maxminddb.Open(path) + db, err := NewReloadableGeoIPDB(path) if err != nil { return nil, fmt.Errorf("open city mmdb: %w", err) } - log.Printf("city mmdb type: %s", db.Metadata.DatabaseType) + log.Printf("city mmdb type: %s", db.DatabaseType()) return &CityReader{db: db}, nil } +// StartWatcher polls the mmdb file and hot reloads it on change. Blocks until +// ctx is cancelled. +func (c *CityReader) StartWatcher(ctx context.Context, interval time.Duration) { + c.db.StartWatcher(ctx, interval) +} + func (c *CityReader) Close() error { return c.db.Close() } type cityRecord struct { @@ -59,7 +64,7 @@ func (c *CityReader) Enrich(ip net.IP, out *LookupResult) error { } var rec cityRecord - if err := c.db.Lookup(addr).Decode(&rec); err != nil { + if err := c.db.Lookup(addr, &rec); err != nil { return err } diff --git a/api/geolite_reloadable.go b/api/geolite_reloadable.go new file mode 100644 index 0000000..d3d905e --- /dev/null +++ b/api/geolite_reloadable.go @@ -0,0 +1,173 @@ +package api + +import ( + "context" + "fmt" + "log" + "net/netip" + "os" + "sync" + "time" + + "github.com/oschwald/maxminddb-golang/v2" +) + +// DefaultReloadInterval is used when no explicit interval is configured. +const DefaultReloadInterval = 60 * time.Second + +// fileSignature is the cheap, portable fingerprint we use to detect that an +// mmdb file on disk has been replaced. geoipupdate writes a temporary file and +// renames it over the target, so the path must be re-stat'ed instead of +// relying on the already open file descriptor. +type fileSignature struct { + modTime time.Time + size int64 +} + +func statSignature(path string) (fileSignature, error) { + fi, err := os.Stat(path) + if err != nil { + return fileSignature{}, err + } + return fileSignature{modTime: fi.ModTime(), size: fi.Size()}, nil +} + +// ReloadableGeoIPDB keeps a single open maxminddb.Reader that is shared by all +// lookups and swaps it for a freshly opened one whenever the underlying file +// changes on disk. It is safe for concurrent use. +type ReloadableGeoIPDB struct { + path string + + mu sync.RWMutex + db *maxminddb.Reader + + sig fileSignature +} + +// NewReloadableGeoIPDB opens path and returns a database handle. The watcher is +// not started automatically; call StartWatcher. +func NewReloadableGeoIPDB(path string) (*ReloadableGeoIPDB, error) { + db, err := maxminddb.Open(path) + if err != nil { + return nil, err + } + + // A failing stat here is not fatal: the next check cycle will simply see a + // different signature and attempt a reload. + sig, _ := statSignature(path) + + return &ReloadableGeoIPDB{path: path, db: db, sig: sig}, nil +} + +// Path returns the file this database was opened from. +func (d *ReloadableGeoIPDB) Path() string { return d.path } + +// DatabaseType returns the mmdb metadata database type of the current reader. +func (d *ReloadableGeoIPDB) DatabaseType() string { + d.mu.RLock() + defer d.mu.RUnlock() + if d.db == nil { + return "" + } + return d.db.Metadata.DatabaseType +} + +// Lookup resolves addr against the currently active reader and decodes the +// record into out. The reader is held under a read lock for the whole decode, +// which is what makes it safe to close a replaced reader under the write lock. +func (d *ReloadableGeoIPDB) Lookup(addr netip.Addr, out any) error { + d.mu.RLock() + defer d.mu.RUnlock() + + if d.db == nil { + return fmt.Errorf("geoip database %s is closed", d.path) + } + return d.db.Lookup(addr).Decode(out) +} + +// StartWatcher polls the file for changes until ctx is cancelled. It blocks, so +// callers normally run it in its own goroutine. +func (d *ReloadableGeoIPDB) StartWatcher(ctx context.Context, interval time.Duration) { + if interval <= 0 { + interval = DefaultReloadInterval + } + + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + if err := d.reloadIfChanged(); err != nil { + log.Printf("Failed to reload GeoLite2 database %s: %v; continuing with previous database", d.path, err) + } + } + } +} + +// reloadIfChanged is a single watcher iteration, exposed for tests. It is a +// no-op (and silent) while the file signature is unchanged. +func (d *ReloadableGeoIPDB) reloadIfChanged() error { + sig, err := statSignature(d.path) + if err != nil { + return fmt.Errorf("stat: %w", err) + } + + d.mu.RLock() + unchanged := sig == d.sig + closed := d.db == nil + d.mu.RUnlock() + + if closed || unchanged { + return nil + } + + // Open the replacement first: if it is corrupt we keep serving from the + // current reader and retry on the next cycle. + next, err := maxminddb.Open(d.path) + if err != nil { + return fmt.Errorf("open: %w", err) + } + + d.mu.Lock() + // Re-check under the write lock: the Open above runs unlocked, so a + // concurrent Close may have set d.db to nil meanwhile. Swapping next in + // anyway would revive the database after Close reported success, leaving a + // reader nobody closes. Harmless if the process exits right away, but not + // when Close is followed by more work in the same process. + if d.db == nil { + d.mu.Unlock() + _ = next.Close() + return nil + } + prev := d.db + d.db = next + d.sig = sig + // Holding the write lock guarantees no lookup is inside prev, and no new + // lookup can enter it, so closing here is safe. + err = prev.Close() + d.mu.Unlock() + + if err != nil { + log.Printf("GeoLite2 database reloaded: %s (closing previous reader failed: %v)", d.path, err) + return nil + } + + log.Printf("GeoLite2 database reloaded: %s", d.path) + return nil +} + +// Close releases the active reader. It is idempotent. +func (d *ReloadableGeoIPDB) Close() error { + d.mu.Lock() + defer d.mu.Unlock() + + if d.db == nil { + return nil + } + db := d.db + d.db = nil + return db.Close() +} diff --git a/api/geolite_reloadable_test.go b/api/geolite_reloadable_test.go new file mode 100644 index 0000000..635e5b3 --- /dev/null +++ b/api/geolite_reloadable_test.go @@ -0,0 +1,218 @@ +package api + +import ( + "context" + "net/netip" + "os" + "path/filepath" + "sync" + "testing" + "time" +) + +var testAddr = netip.MustParseAddr("1.2.3.4") + +type testRecord struct { + Test string `maxminddb:"test"` +} + +// writeDB atomically installs content at path the same way geoipupdate does: +// write a temp file, then rename over the target. +func writeDB(t *testing.T, path string, content []byte) { + t.Helper() + + tmp := path + ".tmp" + if err := os.WriteFile(tmp, content, 0o644); err != nil { + t.Fatalf("write temp db: %v", err) + } + if err := os.Rename(tmp, path); err != nil { + t.Fatalf("rename db: %v", err) + } + // Make sure the new file has a distinguishable ModTime even on filesystems + // with coarse timestamps. + future := time.Now().Add(2 * time.Second) + if err := os.Chtimes(path, future, future); err != nil { + t.Fatalf("chtimes: %v", err) + } +} + +func newTestDB(t *testing.T, value string) (*ReloadableGeoIPDB, string) { + t.Helper() + + path := filepath.Join(t.TempDir(), "Test.mmdb") + writeDB(t, path, buildTestMMDB(value)) + + db, err := NewReloadableGeoIPDB(path) + if err != nil { + t.Fatalf("open db: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + + return db, path +} + +func mustLookup(t *testing.T, db *ReloadableGeoIPDB) string { + t.Helper() + + var rec testRecord + if err := db.Lookup(testAddr, &rec); err != nil { + t.Fatalf("lookup: %v", err) + } + return rec.Test +} + +func TestOpensAtStartup(t *testing.T) { + db, _ := newTestDB(t, "v1") + + if got := db.DatabaseType(); got != "Test" { + t.Errorf("DatabaseType = %q, want %q", got, "Test") + } + if got := mustLookup(t, db); got != "v1" { + t.Errorf("lookup = %q, want %q", got, "v1") + } +} + +func TestNoReloadWhenUnchanged(t *testing.T) { + db, _ := newTestDB(t, "v1") + + before := db.db + for i := 0; i < 3; i++ { + if err := db.reloadIfChanged(); err != nil { + t.Fatalf("reloadIfChanged: %v", err) + } + } + if db.db != before { + t.Error("reader was swapped even though the file did not change") + } +} + +func TestReloadPicksUpReplacedFile(t *testing.T) { + db, path := newTestDB(t, "v1") + + writeDB(t, path, buildTestMMDB("v2")) + if err := db.reloadIfChanged(); err != nil { + t.Fatalf("reloadIfChanged: %v", err) + } + + if got := mustLookup(t, db); got != "v2" { + t.Errorf("lookup after reload = %q, want %q", got, "v2") + } +} + +func TestCorruptReplacementKeepsPreviousDatabase(t *testing.T) { + db, path := newTestDB(t, "v1") + + writeDB(t, path, []byte("this is not an mmdb file")) + if err := db.reloadIfChanged(); err == nil { + t.Fatal("expected an error for a corrupt database") + } + + if got := mustLookup(t, db); got != "v1" { + t.Errorf("lookup after failed reload = %q, want %q", got, "v1") + } + + // A subsequent valid replacement must reload again. + writeDB(t, path, buildTestMMDB("v3")) + if err := db.reloadIfChanged(); err != nil { + t.Fatalf("reloadIfChanged after recovery: %v", err) + } + if got := mustLookup(t, db); got != "v3" { + t.Errorf("lookup after recovery = %q, want %q", got, "v3") + } +} + +func TestConcurrentLookupsDuringReload(t *testing.T) { + db, path := newTestDB(t, "v1") + + stop := make(chan struct{}) + var wg sync.WaitGroup + + for i := 0; i < 8; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case <-stop: + return + default: + } + var rec testRecord + if err := db.Lookup(testAddr, &rec); err != nil { + t.Errorf("lookup during reload: %v", err) + return + } + if rec.Test == "" { + t.Error("empty record during reload") + return + } + } + }() + } + + for i := 0; i < 20; i++ { + writeDB(t, path, buildTestMMDB("v"+string(rune('A'+i)))) + if err := db.reloadIfChanged(); err != nil { + t.Errorf("reloadIfChanged: %v", err) + } + } + + close(stop) + wg.Wait() +} + +func TestWatcherStopsOnContextCancel(t *testing.T) { + db, path := newTestDB(t, "v1") + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + db.StartWatcher(ctx, 10*time.Millisecond) + close(done) + }() + + writeDB(t, path, buildTestMMDB("v2")) + + deadline := time.After(2 * time.Second) + for { + var rec testRecord + if err := db.Lookup(testAddr, &rec); err != nil { + t.Fatalf("lookup: %v", err) + } + if rec.Test == "v2" { + break + } + select { + case <-deadline: + t.Fatal("watcher did not pick up the new database") + case <-time.After(10 * time.Millisecond): + } + } + + cancel() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("watcher did not stop on context cancellation") + } +} + +func TestCloseReleasesResources(t *testing.T) { + db, _ := newTestDB(t, "v1") + + if err := db.Close(); err != nil { + t.Fatalf("close: %v", err) + } + if err := db.Close(); err != nil { + t.Fatalf("second close should be a no-op: %v", err) + } + + var rec testRecord + if err := db.Lookup(testAddr, &rec); err == nil { + t.Error("lookup on a closed database should fail, not panic or succeed") + } + // Reload on a closed database must stay a no-op. + if err := db.reloadIfChanged(); err != nil { + t.Errorf("reloadIfChanged after close: %v", err) + } +} diff --git a/api/geolite_testdata_test.go b/api/geolite_testdata_test.go new file mode 100644 index 0000000..d4e2222 --- /dev/null +++ b/api/geolite_testdata_test.go @@ -0,0 +1,124 @@ +package api + +import ( + "bytes" + "encoding/binary" +) + +// This file builds a minimal but valid MMDB file in memory so the reload tests +// do not need a multi-megabyte binary fixture in the repository. +// +// Layout: a single search-tree node whose left record points into the data +// section and whose right record is "not found". Any IPv4 address whose first +// bit is 0 (e.g. 1.2.3.4) therefore resolves to the single record. + +const mmdbMetadataMarker = "\xab\xcd\xefMaxMind.com" + +// buildTestMMDB returns the bytes of an mmdb whose only record is +// {"test": value}. +func buildTestMMDB(value string) []byte { + const nodeCount = 1 + + var tree bytes.Buffer + // record_size 32: two 4-byte records. Values >= nodeCount+16 are data + // section offsets (value - (nodeCount + 16)); a value == nodeCount means + // "not found". + _ = binary.Write(&tree, binary.BigEndian, uint32(nodeCount+16)) // left -> data offset 0 + _ = binary.Write(&tree, binary.BigEndian, uint32(nodeCount)) // right -> not found + + data := encMap(map[string]any{"test": value}) + + metadata := encMap(map[string]any{ + "binary_format_major_version": encUint16(2), + "binary_format_minor_version": encUint16(0), + "build_epoch": encUint64(1700000000), + "database_type": "Test", + "description": encMap(map[string]any{"en": "test database"}), + "ip_version": encUint16(4), + "languages": encArray("en"), + "node_count": encUint32(nodeCount), + "record_size": encUint16(32), + }) + + var out bytes.Buffer + out.Write(tree.Bytes()) + out.Write(make([]byte, 16)) // data section separator + out.Write(data) + out.WriteString(mmdbMetadataMarker) + out.Write(metadata) + return out.Bytes() +} + +// raw wraps already-encoded bytes so they pass through enc unchanged. +type raw []byte + +func enc(v any) []byte { + switch t := v.(type) { + case raw: + return t + case string: + return append(ctrl(2, len(t)), t...) + } + panic("unsupported test value") +} + +// ctrl builds a control byte (plus extended size bytes) for the given type and +// payload size. Only sizes < 29 are needed here. +func ctrl(typ, size int) []byte { + if size >= 29 { + panic("test encoder only supports small payloads") + } + if typ < 8 { + return []byte{byte(typ<<5 | size)} + } + // Extended type: type field 0, followed by (type - 7). + return []byte{byte(size), byte(typ - 7)} +} + +func encMap(m map[string]any) raw { + // Key order is irrelevant to the reader, but keep it deterministic. + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + sortStrings(keys) + + out := ctrl(7, len(m)) + for _, k := range keys { + out = append(out, enc(k)...) + out = append(out, enc(m[k])...) + } + return out +} + +func encArray(items ...string) raw { + out := ctrl(11, len(items)) + for _, it := range items { + out = append(out, enc(it)...) + } + return out +} + +func encUint16(v uint64) raw { return encUintType(5, v) } +func encUint32(v uint64) raw { return encUintType(6, v) } +func encUint64(v uint64) raw { return encUintType(9, v) } + +func encUintType(typ int, v uint64) raw { + var b []byte + for i := 7; i >= 0; i-- { + by := byte(v >> (8 * i)) + if len(b) == 0 && by == 0 { + continue + } + b = append(b, by) + } + return append(ctrl(typ, len(b)), b...) +} + +func sortStrings(s []string) { + for i := 1; i < len(s); i++ { + for j := i; j > 0 && s[j] < s[j-1]; j-- { + s[j], s[j-1] = s[j-1], s[j] + } + } +} diff --git a/cmd/main.go b/cmd/main.go index 845d746..6530455 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -1,11 +1,16 @@ package main import ( + "context" + "errors" "fmt" "log" "net" "net/http" + "os/signal" "strings" + "syscall" + "time" ipqapi "github.com/akyriako/ipquery/api" "github.com/caarlos0/env/v11" @@ -19,6 +24,9 @@ type Config struct { GeoLiteAsn string `env:"GEOLITE2_ASN" envDefault:"./geolite/GeoLite2-ASN.mmdb"` GeoLiteCity string `env:"GEOLITE2_CITY" envDefault:"./geolite/GeoLite2-City.mmdb"` AbuseIpDbApiKey *string `env:"ABUSEIPDB_API_KEY"` + // GeoIpReloadInterval controls how often the mmdb files are checked for + // replacement by geoipupdate. + GeoIpReloadInterval time.Duration `env:"GEOIP_RELOAD_INTERVAL" envDefault:"60s"` } func main() { @@ -36,6 +44,9 @@ func main() { log.Printf("trustedProxies: %v", trusted) + ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + asn, err := ipqapi.NewAsnReader(cfg.GeoLiteAsn) if err != nil { log.Fatalf("asn reader error: %v", err) @@ -48,6 +59,10 @@ func main() { } defer city.Close() + log.Printf("geoip reload interval: %s", cfg.GeoIpReloadInterval) + go asn.StartWatcher(ctx, cfg.GeoIpReloadInterval) + go city.StartWatcher(ctx, cfg.GeoIpReloadInterval) + lc := &ipqapi.LookupClient{TrustedProxies: trusted, AsnReader: asn, CityReader: city} if cfg.AbuseIpDbApiKey != nil { risk := ipqapi.NewAbuseIpDbChecker(*cfg.AbuseIpDbApiKey) @@ -67,8 +82,24 @@ func main() { r.Get("/lookup/{ip}", apis.LookupIPAll) r.Get("/health", apis.GetHealth) - log.Printf("listening on %s", cfg.ListenAddr) - log.Fatal(http.ListenAndServe(cfg.ListenAddr, r)) + srv := &http.Server{Addr: cfg.ListenAddr, Handler: r} + + go func() { + log.Printf("listening on %s", cfg.ListenAddr) + if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { + log.Fatalf("listen: %v", err) + } + }() + + <-ctx.Done() + stop() + log.Print("shutting down ipquery server") + + shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + if err := srv.Shutdown(shutdownCtx); err != nil { + log.Printf("graceful shutdown: %v", err) + } } func parseCIDRs(items []string) ([]*net.IPNet, error) {