From 0ae08fe952b5790a89319cf9f446022e02ce0976 Mon Sep 17 00:00:00 2001 From: AJ ONeal Date: Wed, 9 Sep 2026 11:11:31 -0600 Subject: [PATCH] feat(net/geoip): add reusable archive fetcher --- net/geoip/cmd/geoip-update/main.go | 5 +- net/geoip/conf.go | 3 + net/geoip/fetch.go | 108 ++++++++++++++ net/geoip/fetch_test.go | 220 +++++++++++++++++++++++++++++ net/geoip/go.mod | 3 +- net/geoip/go.sum | 2 + 6 files changed, 339 insertions(+), 2 deletions(-) create mode 100644 net/geoip/fetch.go create mode 100644 net/geoip/fetch_test.go diff --git a/net/geoip/cmd/geoip-update/main.go b/net/geoip/cmd/geoip-update/main.go index 7252b22..eb3062d 100644 --- a/net/geoip/cmd/geoip-update/main.go +++ b/net/geoip/cmd/geoip-update/main.go @@ -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, diff --git a/net/geoip/conf.go b/net/geoip/conf.go index 3f984db..0203333 100644 --- a/net/geoip/conf.go +++ b/net/geoip/conf.go @@ -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 diff --git a/net/geoip/fetch.go b/net/geoip/fetch.go new file mode 100644 index 0000000..841b81d --- /dev/null +++ b/net/geoip/fetch.go @@ -0,0 +1,108 @@ +package geoip + +import ( + "context" + "encoding/base64" + "fmt" + "net/http" + "os" + "path/filepath" + "time" + + "github.com/therootcompany/golib/https" + "github.com/therootcompany/golib/net/httpcache" +) + +const DefaultFailureBackoff = 6 * time.Hour + +// Fetcher downloads the GeoLite2 archives named by a GeoIP.conf. It +// implements dataset.Fetcher and keeps freshness/backoff policy in the +// httpcache layer. +type Fetcher struct { + ConfPath string + Dir string + BaseURL string + FreshDays int + FailureBackoff time.Duration + HTTPClient *http.Client + + conf *Conf + auth http.Header + maxAge time.Duration + cachers []*httpcache.Cacher +} + +// NewFetcher parses confPath and creates a fetcher for its editions. +func NewFetcher(confPath string) (*Fetcher, 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 &Fetcher{ConfPath: confPath, Dir: dir, BaseURL: conf.BaseURL, conf: conf}, nil +} + +func (f *Fetcher) 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. +func (f *Fetcher) Fetch(ctx context.Context) (bool, error) { + 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.Fetch(ctx) + if err != nil { + return false, err + } + updated = updated || u + } + return updated, nil +} diff --git a/net/geoip/fetch_test.go b/net/geoip/fetch_test.go new file mode 100644 index 0000000..bf9a4b0 --- /dev/null +++ b/net/geoip/fetch_test.go @@ -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() +} + +// TestFetcherDownloadAuth 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 TestFetcherDownloadAuth(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 := NewFetcher(confPath) + if err != nil { + t.Fatalf("NewFetcher: %v", err) + } + f.Dir = dir + f.BaseURL = srv.URL + f.HTTPClient = srv.Client() + + updated, err := f.Fetch(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.Fetch(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) + } +} + +// TestFetcherConditionalGET 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 TestFetcherConditionalGET(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 := NewFetcher(confPath) + if err != nil { + t.Fatalf("NewFetcher: %v", err) + } + f.Dir = dir + f.BaseURL = srv.URL + f.HTTPClient = srv.Client() + + // Prime with v1. + if updated, err := f.Fetch(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.Fetch(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.Fetch(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)) + } +} + +// TestFetcherDefaultFreshDays ensures a zero FreshDays defaults to a 3-day +// MaxAge gate. +func TestFetcherDefaultFreshDays(t *testing.T) { + dir := t.TempDir() + confPath := writeConf(t, dir, "GeoLite2-City") + f, err := NewFetcher(confPath) + if err != nil { + t.Fatalf("NewFetcher: %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)) + } +} + +// TestFetcherExplicitFreshDays checks a non-default FreshDays is applied. +func TestFetcherExplicitFreshDays(t *testing.T) { + dir := t.TempDir() + confPath := writeConf(t, dir, "GeoLite2-City") + f, err := NewFetcher(confPath) + if err != nil { + t.Fatalf("NewFetcher: %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) + } +} + +// TestFetcherMissingConf ensures a missing conf file returns an error. +func TestFetcherMissingConf(t *testing.T) { + if _, err := NewFetcher(filepath.Join(t.TempDir(), "nope.conf")); err == nil { + t.Fatal("expected error for missing conf") + } +} diff --git a/net/geoip/go.mod b/net/geoip/go.mod index 9c41757..f0d8f75 100644 --- a/net/geoip/go.mod +++ b/net/geoip/go.mod @@ -1,6 +1,6 @@ module github.com/therootcompany/golib/net/geoip -go 1.26.0 +go 1.26.5 require ( github.com/oschwald/geoip2-golang v1.13.0 @@ -9,6 +9,7 @@ require ( require ( github.com/oschwald/maxminddb-golang v1.13.0 // indirect + github.com/therootcompany/golib/https v0.9.0 // indirect golang.org/x/sync v0.20.0 // indirect golang.org/x/sys v0.45.0 // indirect ) diff --git a/net/geoip/go.sum b/net/geoip/go.sum index 8f4ace5..443492e 100644 --- a/net/geoip/go.sum +++ b/net/geoip/go.sum @@ -8,6 +8,8 @@ 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/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.5.1 h1:vQM758tBwPEqCo7ojpSeL3fHqgq29LYmvD2hvDTBfKY= github.com/therootcompany/golib/net/httpcache v0.5.1/go.mod h1:mKSQJWpGSCDBzWkTCwT7dO1SCVCsFmdkRgKr39fNS1U= golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=