mirror of
https://github.com/therootcompany/golib.git
synced 2026-11-04 06:55:00 +00:00
refactor(net/geoip): use released cache APIs
This commit is contained in:
parent
c43fd48b34
commit
87ab75126a
@ -175,10 +175,13 @@ func main() {
|
||||
exitCode := 0
|
||||
for _, edition := range conf.EditionIDs {
|
||||
path := filepath.Join(outDir, geoip.TarGzName(edition))
|
||||
base := geoip.DownloadBase
|
||||
base := conf.BaseURL
|
||||
if cfg.BaseURL != "" {
|
||||
base = cfg.BaseURL
|
||||
}
|
||||
if base == "" {
|
||||
base = geoip.DownloadBase
|
||||
}
|
||||
cacher := httpcache.New(
|
||||
base+"/"+edition+"/download?suffix=tar.gz",
|
||||
path,
|
||||
@ -188,7 +191,7 @@ func main() {
|
||||
cacher.Header = authHeader
|
||||
_, _ = fmt.Fprintf(os.Stderr, "Fetching %s... ", edition)
|
||||
t := time.Now()
|
||||
updated, err := cacher.Fetch(ctx)
|
||||
updated, err := cacher.Update(ctx)
|
||||
if err != nil {
|
||||
_, _ = fmt.Fprintln(os.Stderr)
|
||||
_, _ = fmt.Fprintf(os.Stderr, "error: %s: %v\n", edition, err)
|
||||
|
||||
@ -15,11 +15,13 @@ import (
|
||||
// AccountID 123456
|
||||
// LicenseKey xxxxxxxxxxxxxxxx
|
||||
// EditionIDs GeoLite2-City GeoLite2-ASN
|
||||
// BaseURL https://geoip.example.com/geoip/databases
|
||||
type Conf struct {
|
||||
AccountID string
|
||||
LicenseKey string
|
||||
EditionIDs []string
|
||||
DatabaseDirectory string
|
||||
BaseURL string
|
||||
}
|
||||
|
||||
// ErrMissingCredentials is returned by ParseConf when AccountID or LicenseKey
|
||||
@ -48,6 +50,7 @@ func ParseConf(s string) (*Conf, error) {
|
||||
AccountID: kv["AccountID"],
|
||||
LicenseKey: kv["LicenseKey"],
|
||||
DatabaseDirectory: kv["DatabaseDirectory"],
|
||||
BaseURL: kv["BaseURL"],
|
||||
}
|
||||
if c.AccountID == "" || c.LicenseKey == "" {
|
||||
return nil, ErrMissingCredentials
|
||||
|
||||
@ -98,9 +98,18 @@ func openMMDBTarGz(path string) (*geoip2.Reader, error) {
|
||||
}
|
||||
}
|
||||
|
||||
// Close closes the city and ASN readers.
|
||||
// Close closes the city and ASN readers. It is safe to call on a
|
||||
// partially-initialised Databases (nil readers are skipped), which makes
|
||||
// Set/Clear robust against manually-constructed snapshots.
|
||||
func (d *Databases) Close() error {
|
||||
return errors.Join(d.City.Close(), d.ASN.Close())
|
||||
var errs []error
|
||||
if d.City != nil {
|
||||
errs = append(errs, d.City.Close())
|
||||
}
|
||||
if d.ASN != nil {
|
||||
errs = append(errs, d.ASN.Close())
|
||||
}
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
|
||||
// Info is the structured result of a GeoIP lookup.
|
||||
|
||||
296
net/geoip/fetch.go
Normal file
296
net/geoip/fetch.go
Normal file
@ -0,0 +1,296 @@
|
||||
package geoip
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/therootcompany/golib/https"
|
||||
"github.com/therootcompany/golib/net/httpcache"
|
||||
"github.com/therootcompany/golib/sync/cachable"
|
||||
)
|
||||
|
||||
const (
|
||||
DefaultFailureBackoff = 6 * time.Hour
|
||||
DefaultRefreshTimeout = 24*time.Hour + time.Minute + 7*time.Second
|
||||
)
|
||||
|
||||
// GeoIPDB downloads the GeoLite2 archives named by a GeoIP.conf. It
|
||||
// implements dataset.Upstream and keeps freshness/backoff policy in the
|
||||
// httpcache layer.
|
||||
type GeoIPDB struct {
|
||||
ConfPath string
|
||||
Dir string
|
||||
BaseURL string
|
||||
FreshDays int
|
||||
FailureBackoff time.Duration
|
||||
RefreshTimeout time.Duration
|
||||
HTTPClient *http.Client
|
||||
|
||||
conf *Conf
|
||||
auth http.Header
|
||||
maxAge time.Duration
|
||||
cachers []*httpcache.Cacher
|
||||
current atomic.Pointer[Databases]
|
||||
start sync.Once
|
||||
generation atomic.Uint64
|
||||
cancel context.CancelFunc
|
||||
loadedAt atomic.Pointer[time.Time]
|
||||
lastErr atomic.Pointer[error]
|
||||
refresh cachable.Refresh
|
||||
local bool
|
||||
}
|
||||
|
||||
// NewGeoIPDB parses confPath and creates a fetcher for its editions.
|
||||
func NewGeoIPDB(confPath string) (*GeoIPDB, error) {
|
||||
data, err := os.ReadFile(confPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
conf, err := ParseConf(string(data))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parse %s: %w", confPath, err)
|
||||
}
|
||||
if len(conf.EditionIDs) == 0 {
|
||||
return nil, fmt.Errorf("no EditionIDs found in %s", confPath)
|
||||
}
|
||||
dir := conf.DatabaseDirectory
|
||||
if dir == "" {
|
||||
dir = "."
|
||||
}
|
||||
return &GeoIPDB{ConfPath: confPath, Dir: dir, BaseURL: conf.BaseURL, conf: conf}, nil
|
||||
}
|
||||
|
||||
func (f *GeoIPDB) ready() error {
|
||||
if f.auth != nil {
|
||||
return nil
|
||||
}
|
||||
f.auth = http.Header{"Authorization": []string{
|
||||
"Basic " + base64.StdEncoding.EncodeToString([]byte(f.conf.AccountID+":"+f.conf.LicenseKey)),
|
||||
}}
|
||||
freshDays := f.FreshDays
|
||||
if freshDays <= 0 {
|
||||
freshDays = 3
|
||||
}
|
||||
f.maxAge = time.Duration(freshDays) * 24 * time.Hour
|
||||
backoff := f.FailureBackoff
|
||||
if backoff == 0 {
|
||||
backoff = DefaultFailureBackoff
|
||||
}
|
||||
base := f.BaseURL
|
||||
if base == "" {
|
||||
base = DownloadBase
|
||||
}
|
||||
client := f.HTTPClient
|
||||
if client == nil {
|
||||
client = https.NewSlowClient(5 * time.Minute)
|
||||
}
|
||||
f.cachers = make([]*httpcache.Cacher, 0, len(f.conf.EditionIDs))
|
||||
for _, edition := range f.conf.EditionIDs {
|
||||
path := filepath.Join(f.Dir, TarGzName(edition))
|
||||
cacher := httpcache.NewWith(base+"/"+edition+"/download?suffix=tar.gz", path, client)
|
||||
cacher.MaxAge = f.maxAge
|
||||
cacher.FailureBackoff = backoff
|
||||
cacher.Header = f.auth
|
||||
f.cachers = append(f.cachers, cacher)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Fetch downloads all configured editions and reports whether any changed.
|
||||
var (
|
||||
_ cachable.Cacheable[Databases] = (*GeoIPDB)(nil)
|
||||
_ cachable.Mutable[Databases] = (*GeoIPDB)(nil)
|
||||
_ cachable.Inspectable = (*GeoIPDB)(nil)
|
||||
_ cachable.Tickable = (*GeoIPDB)(nil)
|
||||
)
|
||||
|
||||
// Current returns the newest completed City+ASN snapshot without blocking.
|
||||
func (f *GeoIPDB) Current() *Databases {
|
||||
return f.current.Load()
|
||||
}
|
||||
|
||||
// Due reports whether the source should be revalidated. The fetcher
|
||||
// delegates freshness checks to httpcache, so this is deliberately cheap.
|
||||
func (f *GeoIPDB) Due(ctx context.Context) (bool, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return false, err
|
||||
}
|
||||
status := f.Status()
|
||||
return !status.HasValue || status.StaleAt.IsZero() || !time.Now().Before(status.StaleAt), nil
|
||||
}
|
||||
|
||||
// Revalidate fetches archives when stale and publishes City and ASN together.
|
||||
func (f *GeoIPDB) Revalidate(ctx context.Context) (started bool, err error) {
|
||||
due, err := f.Due(ctx)
|
||||
if err != nil || !due {
|
||||
return false, err
|
||||
}
|
||||
if !f.refresh.Begin() {
|
||||
return false, nil
|
||||
}
|
||||
go func() {
|
||||
timeout := f.RefreshTimeout
|
||||
if timeout <= 0 {
|
||||
timeout = DefaultRefreshTimeout
|
||||
}
|
||||
refreshCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), timeout)
|
||||
defer cancel()
|
||||
f.refresh.Finish(f.update(refreshCtx))
|
||||
}()
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (f *GeoIPDB) update(ctx context.Context) error {
|
||||
generation := f.generation.Load()
|
||||
changed, err := f.Update(ctx)
|
||||
if err != nil {
|
||||
errCopy := err
|
||||
f.lastErr.Store(&errCopy)
|
||||
return err
|
||||
}
|
||||
if !changed && f.current.Load() != nil {
|
||||
return nil
|
||||
}
|
||||
databases, err := Open(f.Dir)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if f.generation.Load() != generation {
|
||||
_ = databases.Close()
|
||||
return nil
|
||||
}
|
||||
previous := f.current.Swap(databases)
|
||||
if previous != nil {
|
||||
_ = previous.Close()
|
||||
}
|
||||
now := time.Now()
|
||||
f.loadedAt.Store(&now)
|
||||
f.lastErr.Store(nil)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *GeoIPDB) Set(databases *Databases) error {
|
||||
f.generation.Add(1)
|
||||
previous := f.current.Swap(databases)
|
||||
if previous != nil && previous != databases {
|
||||
_ = previous.Close()
|
||||
}
|
||||
if databases != nil {
|
||||
now := time.Now()
|
||||
f.loadedAt.Store(&now)
|
||||
} else {
|
||||
f.loadedAt.Store(nil)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *GeoIPDB) Clear() error {
|
||||
return f.Set(nil)
|
||||
}
|
||||
|
||||
func (f *GeoIPDB) Status() cachable.Status {
|
||||
loadedAt := time.Time{}
|
||||
if t := f.loadedAt.Load(); t != nil {
|
||||
loadedAt = *t
|
||||
}
|
||||
var lastErr error
|
||||
if p := f.lastErr.Load(); p != nil {
|
||||
lastErr = *p
|
||||
}
|
||||
status := cachable.Status{LoadedAt: loadedAt, HasValue: f.Current() != nil, Refreshing: f.refresh.Running(), LastError: lastErr}
|
||||
if !loadedAt.IsZero() && f.maxAge > 0 {
|
||||
status.StaleAt = loadedAt.Add(f.maxAge)
|
||||
}
|
||||
return status
|
||||
}
|
||||
|
||||
// Load returns the newest completed snapshot. A failed refresh keeps the last
|
||||
// good snapshot, but is returned to the caller for observability.
|
||||
func (f *GeoIPDB) Load(ctx context.Context, wait bool) (*Databases, error) {
|
||||
if _, err := f.Revalidate(ctx); err != nil {
|
||||
return f.Current(), err
|
||||
}
|
||||
if f.Current() == nil || wait {
|
||||
if err := f.refresh.Wait(ctx); err != nil {
|
||||
return f.Current(), err
|
||||
}
|
||||
}
|
||||
return f.Current(), nil
|
||||
}
|
||||
|
||||
// Start starts at most one optional background revalidation ticker.
|
||||
func (f *GeoIPDB) Start(ctx context.Context, interval time.Duration) {
|
||||
f.start.Do(func() {
|
||||
if interval <= 0 {
|
||||
interval = time.Hour
|
||||
}
|
||||
tickCtx, cancel := context.WithCancel(ctx)
|
||||
f.cancel = cancel
|
||||
go func() {
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-tickCtx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
_, _ = f.Revalidate(tickCtx)
|
||||
}
|
||||
}
|
||||
}()
|
||||
})
|
||||
}
|
||||
|
||||
// Stop stops the optional background ticker and waits for any in-flight
|
||||
// refresh to finish or reach RefreshTimeout. It is safe to call repeatedly.
|
||||
func (f *GeoIPDB) Stop() error {
|
||||
if f == nil {
|
||||
return nil
|
||||
}
|
||||
if f.cancel != nil {
|
||||
f.cancel()
|
||||
}
|
||||
timeout := f.RefreshTimeout
|
||||
if timeout <= 0 {
|
||||
timeout = DefaultRefreshTimeout
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
||||
defer cancel()
|
||||
return f.refresh.Stop(ctx)
|
||||
}
|
||||
|
||||
func NewLocalGeoIPDB(dir string) *GeoIPDB {
|
||||
return &GeoIPDB{Dir: dir, local: true, maxAge: time.Hour}
|
||||
}
|
||||
|
||||
func (f *GeoIPDB) Update(ctx context.Context) (bool, error) {
|
||||
if f.local {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return false, err
|
||||
}
|
||||
if f.Current() != nil {
|
||||
return true, nil
|
||||
}
|
||||
} else if err := f.ready(); err != nil {
|
||||
return false, err
|
||||
}
|
||||
updated := false
|
||||
for _, cacher := range f.cachers {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return false, err
|
||||
}
|
||||
u, err := cacher.Update(ctx)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
updated = updated || u
|
||||
}
|
||||
return updated, nil
|
||||
}
|
||||
117
net/geoip/fetch_concurrency_test.go
Normal file
117
net/geoip/fetch_concurrency_test.go
Normal file
@ -0,0 +1,117 @@
|
||||
package geoip
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestGeoIPDB_SetPublishesManually confirms Set publishes immediately and
|
||||
// reports HasValue, without touching the network.
|
||||
func TestGeoIPDB_SetPublishesManually(t *testing.T) {
|
||||
f := &GeoIPDB{Dir: t.TempDir()}
|
||||
db := &Databases{} // empty but non-nil
|
||||
if err := f.Set(db); err != nil {
|
||||
t.Fatalf("Set: %v", err)
|
||||
}
|
||||
if f.Current() == nil {
|
||||
t.Fatal("Current should be non-nil after Set")
|
||||
}
|
||||
if !f.Status().HasValue {
|
||||
t.Error("Status.HasValue should be true after Set")
|
||||
}
|
||||
if f.Status().LoadedAt.IsZero() {
|
||||
t.Error("LoadedAt should be set after Set")
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeoIPDB_ClearNils confirms Clear removes the published snapshot.
|
||||
func TestGeoIPDB_ClearNils(t *testing.T) {
|
||||
f := &GeoIPDB{Dir: t.TempDir()}
|
||||
if err := f.Set(&Databases{}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := f.Clear(); err != nil {
|
||||
t.Fatalf("Clear: %v", err)
|
||||
}
|
||||
if f.Current() != nil {
|
||||
t.Error("Current should be nil after Clear")
|
||||
}
|
||||
if f.Status().HasValue {
|
||||
t.Error("HasValue should be false after Clear")
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeoIPDB_ClearIsIdempotent confirms repeated Clear is safe.
|
||||
func TestGeoIPDB_ClearIsIdempotent(t *testing.T) {
|
||||
f := &GeoIPDB{Dir: t.TempDir()}
|
||||
for range 3 {
|
||||
if err := f.Clear(); err != nil {
|
||||
t.Fatalf("Clear: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeoIPDB_DueWhenEmpty confirms an unloaded DB is due.
|
||||
func TestGeoIPDB_DueWhenEmpty(t *testing.T) {
|
||||
f := &GeoIPDB{Dir: t.TempDir(), maxAge: 3 * 24 * time.Hour}
|
||||
due, err := f.Due(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Due: %v", err)
|
||||
}
|
||||
if !due {
|
||||
t.Error("empty GeoIPDB should be due")
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeoIPDB_DueGatedByMaxAge confirms that after a Set, with a future
|
||||
// StaleAt (loadedAt + maxAge), Due returns false.
|
||||
func TestGeoIPDB_DueGatedByMaxAge(t *testing.T) {
|
||||
f := &GeoIPDB{Dir: t.TempDir(), maxAge: 3 * 24 * time.Hour}
|
||||
if err := f.Set(&Databases{}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
due, err := f.Due(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Due: %v", err)
|
||||
}
|
||||
if due {
|
||||
t.Error("Due should be false immediately after Set with a future StaleAt")
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeoIPDB_KeepsLastGoodOnSetNil confirms Set(nil) clears but does not
|
||||
// panic on a previously-nil snapshot.
|
||||
func TestGeoIPDB_SetNilOnEmpty(t *testing.T) {
|
||||
f := &GeoIPDB{Dir: t.TempDir()}
|
||||
if err := f.Set(nil); err != nil {
|
||||
t.Fatalf("Set(nil) on empty: %v", err)
|
||||
}
|
||||
if f.Current() != nil {
|
||||
t.Error("Current should be nil after Set(nil)")
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeoIPDB_StatusLastErrorStored confirms a failed refresh surfaces its
|
||||
// error in Status.LastError while keeping the last-good snapshot.
|
||||
func TestGeoIPDB_StatusLastErrorStored(t *testing.T) {
|
||||
f := &GeoIPDB{Dir: t.TempDir()}
|
||||
// Seed a last-good snapshot.
|
||||
if err := f.Set(&Databases{}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Simulate a refresh failure by storing an error directly (update() does
|
||||
// this on fetch error).
|
||||
want := errors.New("fetch failed")
|
||||
errCopy := want
|
||||
f.lastErr.Store(&errCopy)
|
||||
|
||||
status := f.Status()
|
||||
if !errors.Is(status.LastError, want) {
|
||||
t.Errorf("LastError = %v, want %v", status.LastError, want)
|
||||
}
|
||||
if !status.HasValue {
|
||||
t.Error("last-good snapshot should still be present after failure")
|
||||
}
|
||||
}
|
||||
220
net/geoip/fetch_test.go
Normal file
220
net/geoip/fetch_test.go
Normal file
@ -0,0 +1,220 @@
|
||||
package geoip
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// writeConf writes a minimal GeoIP.conf with the given edition IDs to dir.
|
||||
func writeConf(t *testing.T, dir string, editions ...string) string {
|
||||
t.Helper()
|
||||
conf := "AccountID 123456\nLicenseKey testkey\nEditionIDs " +
|
||||
joinSpace(editions) + "\n"
|
||||
p := filepath.Join(dir, "GeoIP.conf")
|
||||
if err := os.WriteFile(p, []byte(conf), 0o600); err != nil {
|
||||
t.Fatalf("write conf: %v", err)
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
func joinSpace(ss []string) string {
|
||||
var s strings.Builder
|
||||
for i, x := range ss {
|
||||
if i > 0 {
|
||||
s.WriteString(" ")
|
||||
}
|
||||
s.WriteString(x)
|
||||
}
|
||||
return s.String()
|
||||
}
|
||||
|
||||
// TestGeoIPDBDownloadAuth verifies the first fetch issues a Basic-auth GET,
|
||||
// writes the tarball to disk, and reports updated=true. A second fetch within
|
||||
// MaxAge skips HTTP entirely (server hit count stays at 1).
|
||||
func TestGeoIPDBDownloadAuth(t *testing.T) {
|
||||
const (
|
||||
user = "123456"
|
||||
key = "testkey"
|
||||
)
|
||||
var hits int32
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
atomic.AddInt32(&hits, 1)
|
||||
got := r.Header.Get("Authorization")
|
||||
want := "Basic " + base64.StdEncoding.EncodeToString([]byte(user+":"+key))
|
||||
if got != want {
|
||||
http.Error(w, "bad auth", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
if r.URL.Path != "/GeoLite2-City/download" {
|
||||
http.Error(w, "bad path "+r.URL.Path, http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
w.Header().Set("ETag", `"v1"`)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte("TARGZ-BODY-CITY"))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
dir := t.TempDir()
|
||||
confPath := writeConf(t, dir, "GeoLite2-City")
|
||||
|
||||
f, err := NewGeoIPDB(confPath)
|
||||
if err != nil {
|
||||
t.Fatalf("NewGeoIPDB: %v", err)
|
||||
}
|
||||
f.Dir = dir
|
||||
f.BaseURL = srv.URL
|
||||
f.HTTPClient = srv.Client()
|
||||
|
||||
updated, err := f.Update(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Fetch #1: %v", err)
|
||||
}
|
||||
if !updated {
|
||||
t.Fatal("Fetch #1: expected updated=true")
|
||||
}
|
||||
if got := atomic.LoadInt32(&hits); got != 1 {
|
||||
t.Fatalf("Fetch #1: server hits = %d, want 1", got)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(dir, TarGzName("GeoLite2-City"))); err != nil {
|
||||
t.Fatalf("tarball not written: %v", err)
|
||||
}
|
||||
|
||||
// Second fetch within MaxAge (default 3 days) should skip HTTP.
|
||||
updated, err = f.Update(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Fetch #2: %v", err)
|
||||
}
|
||||
if updated {
|
||||
t.Fatal("Fetch #2: expected updated=false within MaxAge")
|
||||
}
|
||||
if got := atomic.LoadInt32(&hits); got != 1 {
|
||||
t.Fatalf("Fetch #2: server hits = %d, want 1 (MaxAge gate)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeoIPDBConditionalGET verifies the conditional-GET flow: a stale local
|
||||
// file triggers an If-None-Match GET. A 304 reports updated=false; a 200 with
|
||||
// a new ETag reports updated=true and rewrites the file.
|
||||
func TestGeoIPDBConditionalGET(t *testing.T) {
|
||||
var hits, condHits int32
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
atomic.AddInt32(&hits, 1)
|
||||
switch r.Header.Get("If-None-Match") {
|
||||
case `"v1"`:
|
||||
// First conditional request: not modified.
|
||||
if atomic.AddInt32(&condHits, 1) == 1 {
|
||||
w.WriteHeader(http.StatusNotModified)
|
||||
return
|
||||
}
|
||||
// Server changed its mind: ship v2.
|
||||
w.Header().Set("ETag", `"v2"`)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte("TARGZ-CITY-V2"))
|
||||
default:
|
||||
// Prime: offer v1.
|
||||
w.Header().Set("ETag", `"v1"`)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte("TARGZ-CITY-V1"))
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
dir := t.TempDir()
|
||||
confPath := writeConf(t, dir, "GeoLite2-City")
|
||||
|
||||
f, err := NewGeoIPDB(confPath)
|
||||
if err != nil {
|
||||
t.Fatalf("NewGeoIPDB: %v", err)
|
||||
}
|
||||
f.Dir = dir
|
||||
f.BaseURL = srv.URL
|
||||
f.HTTPClient = srv.Client()
|
||||
|
||||
// Prime with v1.
|
||||
if updated, err := f.Update(context.Background()); err != nil || !updated {
|
||||
t.Fatalf("prime: updated=%v err=%v", updated, err)
|
||||
}
|
||||
if got := atomic.LoadInt32(&hits); got != 1 {
|
||||
t.Fatalf("prime hits = %d, want 1", got)
|
||||
}
|
||||
|
||||
// Backdate the local file so MaxAge no longer gates the next fetch.
|
||||
p := filepath.Join(dir, TarGzName("GeoLite2-City"))
|
||||
old := time.Now().Add(-4 * 24 * time.Hour)
|
||||
if err := os.Chtimes(p, old, old); err != nil {
|
||||
t.Fatalf("Chtimes: %v", err)
|
||||
}
|
||||
|
||||
// Server now returns 304 (not modified) for the conditional request.
|
||||
if updated, err := f.Update(context.Background()); err != nil || updated {
|
||||
t.Fatalf("304 path: updated=%v err=%v", updated, err)
|
||||
}
|
||||
if got := atomic.LoadInt32(&hits); got != 2 {
|
||||
t.Fatalf("after 304 hits = %d, want 2", got)
|
||||
}
|
||||
|
||||
// Server now returns 200 with the new body/ETag.
|
||||
if updated, err := f.Update(context.Background()); err != nil || !updated {
|
||||
t.Fatalf("200 path: updated=%v err=%v", updated, err)
|
||||
}
|
||||
if got := atomic.LoadInt32(&hits); got != 3 {
|
||||
t.Fatalf("after 200 hits = %d, want 3", got)
|
||||
}
|
||||
b, _ := os.ReadFile(p)
|
||||
if string(b) != "TARGZ-CITY-V2" {
|
||||
t.Fatalf("body = %q, want TARGZ-CITY-V2", string(b))
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeoIPDBDefaultFreshDays ensures a zero FreshDays defaults to a 3-day
|
||||
// MaxAge gate.
|
||||
func TestGeoIPDBDefaultFreshDays(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
confPath := writeConf(t, dir, "GeoLite2-City")
|
||||
f, err := NewGeoIPDB(confPath)
|
||||
if err != nil {
|
||||
t.Fatalf("NewGeoIPDB: %v", err)
|
||||
}
|
||||
if err := f.ready(); err != nil {
|
||||
t.Fatalf("ready: %v", err)
|
||||
}
|
||||
if want := 3 * 24 * time.Hour; f.maxAge != want {
|
||||
t.Fatalf("maxAge = %v, want %v", f.maxAge, want)
|
||||
}
|
||||
if len(f.cachers) != 1 {
|
||||
t.Fatalf("cachers = %d, want 1", len(f.cachers))
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeoIPDBExplicitFreshDays checks a non-default FreshDays is applied.
|
||||
func TestGeoIPDBExplicitFreshDays(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
confPath := writeConf(t, dir, "GeoLite2-City")
|
||||
f, err := NewGeoIPDB(confPath)
|
||||
if err != nil {
|
||||
t.Fatalf("NewGeoIPDB: %v", err)
|
||||
}
|
||||
f.FreshDays = 10
|
||||
if err := f.ready(); err != nil {
|
||||
t.Fatalf("ready: %v", err)
|
||||
}
|
||||
if want := 10 * 24 * time.Hour; f.maxAge != want {
|
||||
t.Fatalf("maxAge = %v, want %v", f.maxAge, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestGeoIPDBMissingConf ensures a missing conf file returns an error.
|
||||
func TestGeoIPDBMissingConf(t *testing.T) {
|
||||
if _, err := NewGeoIPDB(filepath.Join(t.TempDir(), "nope.conf")); err == nil {
|
||||
t.Fatal("expected error for missing conf")
|
||||
}
|
||||
}
|
||||
@ -69,7 +69,7 @@ func TestDownload_CityAndASN(t *testing.T) {
|
||||
os.Remove(path)
|
||||
os.Remove(path + ".meta")
|
||||
|
||||
updated, err := newCacher(cfg, edition, path).Fetch(t.Context())
|
||||
updated, err := newCacher(cfg, edition, path).Update(t.Context())
|
||||
if err != nil {
|
||||
t.Fatalf("%s Fetch: %v", edition, err)
|
||||
}
|
||||
@ -99,11 +99,11 @@ func TestDownload_ConditionalGet_FreshCacher(t *testing.T) {
|
||||
for _, edition := range []string{geoip.CityEdition, geoip.ASNEdition} {
|
||||
path := filepath.Join(td, geoip.TarGzName(edition))
|
||||
|
||||
if _, err := newCacher(cfg, edition, path).Fetch(t.Context()); err != nil {
|
||||
if _, err := newCacher(cfg, edition, path).Update(t.Context()); err != nil {
|
||||
t.Fatalf("%s initial Fetch: %v", edition, err)
|
||||
}
|
||||
|
||||
updated, err := newCacher(cfg, edition, path).Fetch(t.Context())
|
||||
updated, err := newCacher(cfg, edition, path).Update(t.Context())
|
||||
if err != nil {
|
||||
t.Fatalf("%s fresh Fetch: %v", edition, err)
|
||||
}
|
||||
|
||||
@ -1,10 +1,12 @@
|
||||
module github.com/therootcompany/golib/net/geoip
|
||||
|
||||
go 1.26.0
|
||||
go 1.26.5
|
||||
|
||||
require (
|
||||
github.com/oschwald/geoip2-golang v1.13.0
|
||||
github.com/therootcompany/golib/net/httpcache v0.5.1
|
||||
github.com/therootcompany/golib/https v0.9.0
|
||||
github.com/therootcompany/golib/net/httpcache v0.6.0
|
||||
github.com/therootcompany/golib/sync/cachable v0.1.0
|
||||
)
|
||||
|
||||
require (
|
||||
|
||||
@ -8,8 +8,12 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
|
||||
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
|
||||
github.com/therootcompany/golib/net/httpcache v0.5.1 h1:vQM758tBwPEqCo7ojpSeL3fHqgq29LYmvD2hvDTBfKY=
|
||||
github.com/therootcompany/golib/net/httpcache v0.5.1/go.mod h1:mKSQJWpGSCDBzWkTCwT7dO1SCVCsFmdkRgKr39fNS1U=
|
||||
github.com/therootcompany/golib/https v0.9.0 h1:D9gh0K6ljkvF805JUbylITn8SHPMo08q5PglQGXsK8c=
|
||||
github.com/therootcompany/golib/https v0.9.0/go.mod h1:ilExqR/Uh7oB36oy9mwuNhANcFh+zdCfD3uJiZ11YAY=
|
||||
github.com/therootcompany/golib/net/httpcache v0.6.0 h1:qFBnSs6rZrtUQ+fweP4bHF2FQd/Syvs7mRN5nO5YtkY=
|
||||
github.com/therootcompany/golib/net/httpcache v0.6.0/go.mod h1:mKSQJWpGSCDBzWkTCwT7dO1SCVCsFmdkRgKr39fNS1U=
|
||||
github.com/therootcompany/golib/sync/cachable v0.1.0 h1:swCXtXB116jhSmla81HaXi/9cPsqFI+Oq/6OcSfdU5Q=
|
||||
github.com/therootcompany/golib/sync/cachable v0.1.0/go.mod h1:l/NZdGiSbW5CkLlkkRa0yObkqak2sj/jQTXXqXwvCl0=
|
||||
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
|
||||
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
|
||||
golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY=
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user