cache azure hosted fetch under the azure key

Author: kaumnen

internal/azure/logic.go

@@ -119,8 +119,11 @@ func printListedValues(prefixes []Prefix, dim string) error {
 // fetchRawData mirrors utils.GetRawData's source dispatch but adds a
 // browser-UA scrape step when the resolved endpoint is the Microsoft
 // details.aspx page (Microsoft serves an anti-bot stub to non-browser UAs).
+// Hosted runs go through utils.GetCached so the resolved JSON gets cached
+// under the "azure" key — without that wrapper the scrape+download URL
+// path would refetch on every invocation (raw URLs aren't cached).
 func fetchRawData(ctx context.Context, source string) (string, error) {
-	var endpointURL string
+	var endpointURL, cacheKey string
 	switch {
 	case strings.HasPrefix(source, "https://") || strings.HasPrefix(source, "http://"):
 		endpointURL = source
@@ -134,21 +137,23 @@ func fetchRawData(ctx context.Context, source string) (string, error) {
 		if endpointURL == "" {
 			endpointURL = utils.DefaultEndpoints[source]
 		}
+		cacheKey = source
 	}
 
 	if endpointURL == "" {
 		return utils.GetRawData(ctx, source)
 	}
 
-	if strings.HasSuffix(strings.ToLower(endpointURL), ".json") {
-		return utils.GetRawData(ctx, endpointURL)
-	}
-
-	jsonURL, err := scrapeJSONURL(ctx, endpointURL)
-	if err != nil {
-		return "", err
-	}
-	return utils.GetRawData(ctx, jsonURL)
+	return utils.GetCached(ctx, cacheKey, func(ctx context.Context) (string, error) {
+		if strings.HasSuffix(strings.ToLower(endpointURL), ".json") {
+			return utils.GetRawData(ctx, endpointURL)
+		}
+		jsonURL, err := scrapeJSONURL(ctx, endpointURL)
+		if err != nil {
+			return "", err
+		}
+		return utils.GetRawData(ctx, jsonURL)
+	})
 }
 
 func scrapeJSONURL(ctx context.Context, pageURL string) (string, error) {

internal/azure/logic_test.go

@@ -4,10 +4,13 @@ import (
 	"bytes"
 	"context"
 	"io"
+	"net/http"
+	"net/http/httptest"
 	"os"
 	"path/filepath"
 	"testing"
 
+	"github.com/spf13/viper"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -300,3 +303,57 @@ func TestGetIPRanges_List(t *testing.T) {
 	})
 	assert.Equal(t, "ActionGroup\nAzureStorage\n", out)
 }
+
+// Regression: azure's two-stage hosted fetch (HTML scrape -> JSON download)
+// previously bypassed the cache because both steps reduced to raw URL calls
+// that GetRawData doesn't cache. Verify the cache wrap now keys the resolved
+// JSON under "azure" so a second call within TTL returns from disk.
+func TestFetchRawData_HostedCachesUnderAzureKey(t *testing.T) {
+	t.Setenv("XDG_CACHE_HOME", t.TempDir())
+	t.Cleanup(func() { viper.Reset() })
+
+	jsonBody := loadFixture(t)
+
+	var hits int
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		hits++
+		_, _ = w.Write([]byte(jsonBody))
+	}))
+	defer srv.Close()
+
+	// Point at a .json endpoint to skip the page-scrape branch (HTTP-redirecting
+	// download.microsoft.com URLs into a test server is more trouble than it's
+	// worth here). What we're proving is that fetchRawData wraps its inner
+	// work in GetCached at all — the scrape branch shares the same wrapper.
+	viper.Set("azure_endpoint", srv.URL+"/ServiceTags_Public_20260101.json")
+
+	for i := 0; i < 2; i++ {
+		got, err := fetchRawData(context.Background(), "azure")
+		require.NoError(t, err)
+		assert.Equal(t, jsonBody, got)
+	}
+	assert.Equal(t, 1, hits, "second call should hit cache, not refetch")
+}
+
+func TestFetchRawData_NoCacheBypass(t *testing.T) {
+	t.Setenv("XDG_CACHE_HOME", t.TempDir())
+	t.Cleanup(func() { viper.Reset() })
+
+	jsonBody := loadFixture(t)
+
+	var hits int
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		hits++
+		_, _ = w.Write([]byte(jsonBody))
+	}))
+	defer srv.Close()
+
+	viper.Set("azure_endpoint", srv.URL+"/x.json")
+	viper.Set("no_cache", true)
+
+	for i := 0; i < 2; i++ {
+		_, err := fetchRawData(context.Background(), "azure")
+		require.NoError(t, err)
+	}
+	assert.Equal(t, 2, hits, "--no-cache should refetch every call")
+}

internal/utils/cache.go

@@ -1,13 +1,55 @@
 package utils
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
 	"time"
+
+	"github.com/spf13/viper"
 )
 
+// GetCached wraps fetch with the standard cache-aside policy: read on hit,
+// run fetch on miss, write the result. If key is "" or --no-cache is set,
+// fetch runs unconditionally and the cache is left untouched. Cache write
+// failures are logged but never fatal. Suitable for providers (e.g. azure)
+// whose fetch path doesn't reduce to a single utils.GetRawData call.
+func GetCached(ctx context.Context, key string, fetch func(context.Context) (string, error)) (string, error) {
+	if key == "" || viper.GetBool("no_cache") {
+		return fetch(ctx)
+	}
+	ttl := resolveCacheTTL(key)
+	if ttl > 0 {
+		if data, ok, _ := readCache(key, ttl); ok {
+			fmt.Fprintf(os.Stderr, "Using cached IP ranges for %s (age %s)\n", key, cacheAge(key))
+			return string(data), nil
+		}
+	}
+	body, err := fetch(ctx)
+	if err != nil {
+		return "", err
+	}
+	if werr := writeCache(key, []byte(body)); werr != nil {
+		fmt.Fprintln(os.Stderr, "Warning: cache write failed:", werr)
+	}
+	return body, nil
+}
+
+func resolveCacheTTL(key string) time.Duration {
+	raw := viper.GetString(key + "_cache_ttl")
+	if raw == "" {
+		return defaultCacheTTL
+	}
+	ttl, err := time.ParseDuration(raw)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Warning: invalid %s_cache_ttl=%q, using %s\n", key, raw, defaultCacheTTL)
+		return defaultCacheTTL
+	}
+	return ttl
+}
+
 func cacheDir() (string, error) {
 	if dir := os.Getenv("XDG_CACHE_HOME"); dir != "" {
 		return filepath.Join(dir, "cipr"), nil

internal/utils/raw_data_handler.go

@@ -37,7 +37,7 @@ func GetRawData(ctx context.Context, source string) (string, error) {
 		if endpointURL == "" && localFile == "" {
 			endpointURL = DefaultEndpoints[source]
 		}
-		if localFile == "" && !viper.GetBool("no_cache") {
+		if localFile == "" {
 			cacheKey = source
 		}
 	}
@@ -51,41 +51,10 @@ func GetRawData(ctx context.Context, source string) (string, error) {
 		return "", fmt.Errorf("no endpoint URL or local file specified for source %q", source)
 	}
 
-	if cacheKey != "" {
-		ttl := resolveCacheTTL(cacheKey)
-		if ttl > 0 {
-			if data, ok, _ := readCache(cacheKey, ttl); ok {
-				fmt.Fprintf(os.Stderr, "Using cached IP ranges for %s (age %s)\n", cacheKey, cacheAge(cacheKey))
-				return string(data), nil
-			}
-		}
-	}
-
-	fmt.Fprintln(os.Stderr, "Fetching IP ranges from endpoint:", endpointURL)
-	body, err := loadFromEndpoint(ctx, endpointURL)
-	if err != nil {
-		return "", err
-	}
-
-	if cacheKey != "" {
-		if werr := writeCache(cacheKey, []byte(body)); werr != nil {
-			fmt.Fprintln(os.Stderr, "Warning: cache write failed:", werr)
-		}
-	}
-	return body, nil
-}
-
-func resolveCacheTTL(key string) time.Duration {
-	raw := viper.GetString(key + "_cache_ttl")
-	if raw == "" {
-		return defaultCacheTTL
-	}
-	ttl, err := time.ParseDuration(raw)
-	if err != nil {
-		fmt.Fprintf(os.Stderr, "Warning: invalid %s_cache_ttl=%q, using %s\n", key, raw, defaultCacheTTL)
-		return defaultCacheTTL
-	}
-	return ttl
+	return GetCached(ctx, cacheKey, func(ctx context.Context) (string, error) {
+		fmt.Fprintln(os.Stderr, "Fetching IP ranges from endpoint:", endpointURL)
+		return loadFromEndpoint(ctx, endpointURL)
+	})
 }
 
 func loadFromFile(filePath string) (string, error) {