feat(iplist): allow custom HTTP clients

This commit is contained in:
AJ ONeal 2026-09-09 01:10:34 -06:00
parent c41ac3732d
commit 03722a2934
No known key found for this signature in database
2 changed files with 45 additions and 14 deletions

View File

@ -52,25 +52,30 @@ const (
//
// Each non-URL entry must be an IP address, CIDR range, or domain name.
// Invalid entries cause an error.
func Load(ctx context.Context, source, cacheDir string) ([]string, error) {
return load(ctx, strings.TrimSpace(source), cacheDir, make(map[string]struct{}), 0)
// Load reads source using client. A nil client uses the internal client with
// the package's standard timeout and transport policy.
func Load(ctx context.Context, source, cacheDir string, client *http.Client) ([]string, error) {
if client == nil {
client = https.NewInternalClient()
}
return load(ctx, strings.TrimSpace(source), cacheDir, make(map[string]struct{}), 0, client)
}
func load(ctx context.Context, source, cacheDir string, seen map[string]struct{}, depth int) ([]string, error) {
func load(ctx context.Context, source, cacheDir string, seen map[string]struct{}, depth int, client *http.Client) ([]string, error) {
if source == "" {
return nil, nil
}
if depth > MaxNestedSources {
return nil, fmt.Errorf("source nesting exceeds %d levels", MaxNestedSources)
}
entries, err := readSource(ctx, source, cacheDir, seen)
entries, err := readSource(ctx, source, cacheDir, seen, client)
if err != nil {
return nil, err
}
result := make([]string, 0, len(entries))
for _, entry := range entries {
if isURL(entry) {
nested, err := load(ctx, entry, cacheDir, seen, depth+1)
nested, err := load(ctx, entry, cacheDir, seen, depth+1, client)
if err != nil {
return nil, err
}
@ -133,8 +138,8 @@ func isDomain(s string) bool {
}
// Each label must be non-empty and contain only letters, digits, and
// hyphens (RFC 1035).
labels := strings.Split(s, ".")
for _, label := range labels {
labels := strings.SplitSeq(s, ".")
for label := range labels {
if label == "" {
return false
}
@ -147,7 +152,7 @@ func isDomain(s string) bool {
return true
}
func readSource(ctx context.Context, source, cacheDir string, seen map[string]struct{}) ([]string, error) {
func readSource(ctx context.Context, source, cacheDir string, seen map[string]struct{}, client *http.Client) ([]string, error) {
parsed, err := url.Parse(source)
if err == nil && (parsed.Scheme == "http" || parsed.Scheme == "https") {
if docid, gid := gsheet2csv.ParseIDs(source); docid != "" {
@ -192,7 +197,7 @@ func readSource(ctx context.Context, source, cacheDir string, seen map[string]st
// the new content can't be parsed, we fall back to (or restore) it.
previous, previousErr := os.ReadFile(cachePath)
cache := httpcache.NewWith(parsed.String(), cachePath, https.NewInternalClient())
cache := httpcache.NewWith(parsed.String(), cachePath, client)
cache.Header = header
cache.MaxAge = defaultMaxAge
cache.MaxBytes = defaultMaxBytes
@ -243,7 +248,7 @@ func Parse(r io.Reader) ([]string, error) {
// isTSV heuristically detects TSV format by checking whether the first
// non-empty, non-comment line contains a tab character.
func isTSV(data []byte) bool {
for _, line := range bytes.Split(data, []byte("\n")) {
for line := range bytes.SplitSeq(data, []byte("\n")) {
trimmed := bytes.TrimSpace(line)
if len(trimmed) == 0 || bytes.HasPrefix(trimmed, []byte("#")) {
continue

View File

@ -55,12 +55,38 @@ func TestLoadRejectsInvalidEntry(t *testing.T) {
if err := os.WriteFile(path, []byte("192.0.2.1\nnot valid entry\n"), 0o600); err != nil {
t.Fatal(err)
}
_, err := Load(t.Context(), path, dir)
_, err := Load(t.Context(), path, dir, nil)
if err == nil {
t.Fatal("expected error for invalid entry, got nil")
}
}
func TestLoadCustomClient(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
_, _ = w.Write([]byte("192.0.2.7\n"))
}))
defer srv.Close()
called := false
client := &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
called = true
return http.DefaultTransport.RoundTrip(r)
})}
entries, err := Load(t.Context(), srv.URL, t.TempDir(), client)
if err != nil {
t.Fatal(err)
}
if !called || len(entries) != 1 || entries[0] != "192.0.2.7" {
t.Fatalf("called=%v entries=%#v", called, entries)
}
}
type roundTripFunc func(*http.Request) (*http.Response, error)
func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) {
return f(r)
}
func TestLoadURLAuthCacheAndNested(t *testing.T) {
var authOK bool
inner := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
@ -74,7 +100,7 @@ func TestLoadURLAuthCacheAndNested(t *testing.T) {
defer outer.Close()
cacheDir := t.TempDir()
got, err := Load(t.Context(), strings.Replace(outer.URL, "http://", "http://u:p@", 1), cacheDir)
got, err := Load(t.Context(), strings.Replace(outer.URL, "http://", "http://u:p@", 1), cacheDir, nil)
if err != nil || !authOK || len(got) != 1 || got[0] != "192.0.2.2" {
t.Fatalf("auth=%v entries=%#v err=%v", authOK, got, err)
}
@ -96,7 +122,7 @@ func TestLoadAcceptHeader(t *testing.T) {
defer srv.Close()
cacheDir := t.TempDir()
if _, err := Load(t.Context(), srv.URL, cacheDir); err != nil {
if _, err := Load(t.Context(), srv.URL, cacheDir, nil); err != nil {
t.Fatal(err)
}
if !acceptOK {
@ -112,7 +138,7 @@ func TestLoadCSVDetection(t *testing.T) {
defer srv.Close()
cacheDir := t.TempDir()
got, err := Load(t.Context(), srv.URL, cacheDir)
got, err := Load(t.Context(), srv.URL, cacheDir, nil)
if err != nil {
t.Fatalf("err=%v", err)
}