mirror of
https://github.com/therootcompany/golib.git
synced 2026-11-04 06:55:00 +00:00
feat(iplist): allow custom HTTP clients
This commit is contained in:
parent
c41ac3732d
commit
03722a2934
@ -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
|
||||
|
||||
@ -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)
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user