cmd/generate: Allow specifying HTTP headers

1 files changed, 30 insertions, 12 deletions

diff --git a/cmd/generate/httpcache.go b/cmd/generate/httpcache.go
index 04762e3..08153d1 100644
--- a/cmd/generate/httpcache.go
+++ b/cmd/generate/httpcache.go
@@ -8,29 +8,47 @@ import (
 	"net/url"
 	"os"
 	"path/filepath"
+	"sort"
 )
 
 var httpCache = map[string]string{}
 
-func httpGet(u string) (string, error) {
-	if cache, ok := httpCache[u]; ok {
+func httpGet(u string, hdr map[string]string) (string, error) {
+	cacheKey := url.QueryEscape(u)
+	hdrKeys := make([]string, 0, len(hdr))
+	for k := range hdr {
+		hdrKeys = append(hdrKeys, http.CanonicalHeaderKey(k))
+	}
+	sort.Strings(hdrKeys)
+	for _, k := range hdrKeys {
+		cacheKey += "|" + url.QueryEscape(k) + ":" + url.QueryEscape(hdr[k])
+	}
+
+	if cache, ok := httpCache[cacheKey]; ok {
 		fmt.Printf("CACHE-GET %q\n", u)
 		return cache, nil
 	}
 	if err := os.Mkdir(".http-cache", 0777); err != nil && !os.IsExist(err) {
 		return "", err
 	}
-	cacheFile := filepath.Join(".http-cache", url.QueryEscape(u))
+	cacheFile := filepath.Join(".http-cache", cacheKey)
 	if bs, err := os.ReadFile(cacheFile); err == nil {
-		httpCache[u] = string(bs)
+		httpCache[cacheKey] = string(bs)
 		fmt.Printf("CACHE-GET %q\n", u)
-		return httpCache[u], nil
+		return httpCache[cacheKey], nil
 	} else if !os.IsNotExist(err) {
 		return "", err
 	}
 
 	fmt.Printf("GET %q...", u)
-	resp, err := http.Get(u)
+	req, err := http.NewRequest(http.MethodGet, u, nil)
+	if err != nil {
+		return "", err
+	}
+	for k, v := range hdr {
+		req.Header.Add(k, v)
+	}
+	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
 		fmt.Printf(" err\n")
 		return "", err
@@ -48,19 +66,19 @@ func httpGet(u string) (string, error) {
 	if err := os.WriteFile(cacheFile, bs, 0666); err != nil {
 		return "", err
 	}
-	httpCache[u] = string(bs)
-	return httpCache[u], nil
+	httpCache[cacheKey] = string(bs)
+	return httpCache[cacheKey], nil
 }
 
-func httpGetJSON(u string, out any) error {
-	str, err := httpGet(u)
+func httpGetJSON(u string, hdr map[string]string, out any) error {
+	str, err := httpGet(u, hdr)
 	if err != nil {
 		return err
 	}
 	return json.Unmarshal([]byte(str), out)
 }
 
-func httpGetPaginatedJSON[T any](uStr string, out *[]T, pageFn func(i int) url.Values) error {
+func httpGetPaginatedJSON[T any](uStr string, hdr map[string]string, out *[]T, pageFn func(i int) url.Values) error {
 	u, err := url.Parse(uStr)
 	if err != nil {
 		return err
@@ -75,7 +93,7 @@ func httpGetPaginatedJSON[T any](uStr string, out *[]T, pageFn func(i int) url.V
 
 		u.RawQuery = query.Encode()
 		var resp []T
-		if err := httpGetJSON(u.String(), &resp); err != nil {
+		if err := httpGetJSON(u.String(), hdr, &resp); err != nil {
 			return err
 		}
 		fmt.Printf(" -> %d records\n", len(resp))