diff --git a/net/iplist/iplist.go b/net/iplist/iplist.go index ad8cb56..be9bed4 100644 --- a/net/iplist/iplist.go +++ b/net/iplist/iplist.go @@ -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 diff --git a/net/iplist/iplist_test.go b/net/iplist/iplist_test.go index 272ec3f..b65513f 100644 --- a/net/iplist/iplist_test.go +++ b/net/iplist/iplist_test.go @@ -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) }