diff --git a/geoblock.go b/geoblock.go index 6b7c0c1..0eabcd1 100755 --- a/geoblock.go +++ b/geoblock.go @@ -21,7 +21,7 @@ const ( xForwardedFor = "X-Forwarded-For" xRealIP = "X-Real-IP" countryHeader = "X-IPCountry" - numberOfHoursInMonth = 30 * 24 + defaultCacheTTL = 30 * 24 * time.Hour // legacy forceMonthlyUpdate interval unknownCountryCode = "AA" countryCodeLength = 2 defaultDeniedRequestHTTPStatusCode = 403 @@ -42,6 +42,7 @@ type Config struct { IPGeolocationHTTPHeaderField string `yaml:"ipGeolocationHttpHeaderField"` XForwardedForReverseProxy bool `yaml:"xForwardedForReverseProxy"` CacheSize int `yaml:"cacheSize"` + CacheTTLSeconds int `yaml:"cacheTtlSeconds"` ForceMonthlyUpdate bool `yaml:"forceMonthlyUpdate"` AllowUnknownCountries bool `yaml:"allowUnknownCountries"` UnknownCountryAPIResponse string `yaml:"unknownCountryApiResponse"` @@ -76,6 +77,7 @@ type GeoBlock struct { logAPIRequests bool apiURI string apiTimeoutMs int + cacheTTL time.Duration ignoreAPITimeout bool ignoreAPIFailures bool iPGeolocationHTTPHeaderField string @@ -177,6 +179,7 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h logAPIRequests: config.LogAPIRequests, apiURI: config.API, apiTimeoutMs: config.APITimeoutMs, + cacheTTL: time.Duration(config.CacheTTLSeconds) * time.Second, ignoreAPITimeout: config.IgnoreAPITimeout, ignoreAPIFailures: config.IgnoreAPIFailures, iPGeolocationHTTPHeaderField: config.IPGeolocationHTTPHeaderField, @@ -306,6 +309,24 @@ func (a *GeoBlock) allowDenyIPAddress(requestIPAddr *net.IP, req *http.Request) return allowed } +// shouldRefreshEntry reports whether a cached entry has outlived its TTL and +// must be re-fetched from the API. +// +// An explicit cacheTtlSeconds takes effect on its own. When it is unset (<= 0), +// forceMonthlyUpdate preserves the legacy behavior of refreshing entries once +// they are ~30 days old otherwise entries never expire by age. +func (a *GeoBlock) shouldRefreshEntry(entry ipEntry) bool { + ttl := a.cacheTTL + if ttl <= 0 { + if !a.forceMonthlyUpdate { + return false + } + ttl = defaultCacheTTL + } + + return time.Since(entry.Timestamp) >= ttl +} + func (a *GeoBlock) allowDenyCachedRequestIP(requestIPAddr *net.IP, req *http.Request) (bool, string) { ipAddressString := requestIPAddr.String() cacheEntry, cacheHit := a.database.Get(ipAddressString) @@ -339,8 +360,8 @@ func (a *GeoBlock) allowDenyCachedRequestIP(requestIPAddr *net.IP, req *http.Req a.infoLogger.Printf("%s: [%s] loaded from database: %s", a.name, requestIPAddr, entry) } - // check if existing entry was made more than a month ago, if so update the entry - if time.Since(entry.Timestamp).Hours() >= numberOfHoursInMonth && a.forceMonthlyUpdate { + // check if existing entry is older than the configured cache TTL, if so update the entry + if a.shouldRefreshEntry(entry) { entry, err = a.createNewIPEntry(req, ipAddressString) if err != nil { if a.ignoreAPIFailures { @@ -410,8 +431,8 @@ func (a *GeoBlock) cachedRequestIP(requestIPAddr *net.IP, req *http.Request) (bo a.infoLogger.Printf("%s: [%s] Loaded from database: %s", a.name, ipAddressString, entry) } - // check if existing entry was made more than a month ago, if so update the entry - if time.Since(entry.Timestamp).Hours() >= numberOfHoursInMonth && a.forceMonthlyUpdate { + // check if existing entry is older than the configured cache TTL, if so update the entry + if a.shouldRefreshEntry(entry) { entry, err = a.createNewIPEntry(req, ipAddressString) if err != nil { return false, "" @@ -699,6 +720,7 @@ func printConfiguration(name string, config *Config, logger *log.Logger) { logger.Printf("%s: API timeout: %d", name, config.APITimeoutMs) logger.Printf("%s: ignore API timeout: %t", name, config.IgnoreAPITimeout) logger.Printf("%s: cache size: %d", name, config.CacheSize) + logger.Printf("%s: cache ttl seconds: %d", name, config.CacheTTLSeconds) logger.Printf("%s: force monthly update: %t", name, config.ForceMonthlyUpdate) logger.Printf("%s: allow unknown countries: %t", name, config.AllowUnknownCountries) logger.Printf("%s: unknown country api response: %s", name, config.UnknownCountryAPIResponse) diff --git a/geoblock_test.go b/geoblock_test.go index c360114..87df9cc 100755 --- a/geoblock_test.go +++ b/geoblock_test.go @@ -11,6 +11,7 @@ import ( "os" "path/filepath" "strings" + "sync/atomic" "testing" "time" @@ -1742,6 +1743,64 @@ func waitForFileNonEmpty(t *testing.T, path string, timeout time.Duration) { t.Fatalf("persistence file %q remained empty within %s (size=%d)", path, timeout, fi.Size()) } +func TestCacheTTLExpiresWithoutForceMonthlyUpdate(t *testing.T) { + // cacheTtlSeconds must drive cache refresh on its own, independent of the + // legacy forceMonthlyUpdate flag. + // + // Use a unique middleware name and a dedicated IP so the test is isolated + // from any process-wide cache shared by name: it must observe its own + // mock's call count, not a cache entry warmed by another test. + const ttlTestIP = "203.0.113.7" // TEST-NET-3, not used by other tests + + var apiCalls int32 + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, _ *http.Request) { + atomic.AddInt32(&apiCalls, 1) + rw.WriteHeader(http.StatusOK) + _, _ = rw.Write([]byte("CH")) + })) + defer server.Close() + + cfg := createTesterConfig() + cfg.Countries = append(cfg.Countries, "CH") + cfg.API = server.URL + "/{ip}" + cfg.CacheTTLSeconds = 1 + cfg.ForceMonthlyUpdate = false + + ctx := context.Background() + next := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {}) + + handler, err := geoblock.New(ctx, next, cfg, t.Name()) + if err != nil { + t.Fatal(err) + } + + doRequest := func() { + recorder := httptest.NewRecorder() + req, reqErr := http.NewRequestWithContext(ctx, http.MethodGet, "http://localhost", nil) + if reqErr != nil { + t.Fatal(reqErr) + } + req.Header.Add(xForwardedFor, ttlTestIP) + handler.ServeHTTP(recorder, req) + assertStatusCode(t, recorder.Result(), http.StatusOK) + } + + doRequest() // cache miss -> 1 API call + doRequest() // within TTL -> served from cache, no new call + + if got := atomic.LoadInt32(&apiCalls); got != 1 { + t.Fatalf("expected 1 API call while entry is within TTL, got %d", got) + } + + time.Sleep(1200 * time.Millisecond) // exceed the 1s TTL + + doRequest() // entry older than TTL -> refresh -> 2nd API call + + if got := atomic.LoadInt32(&apiCalls); got != 2 { + t.Fatalf("expected refresh after TTL expiry (2 API calls), got %d", got) + } +} + func createTesterConfig() *geoblock.Config { cfg := geoblock.CreateConfig()