root/golib
1
0
Fork 0
mirror of https://github.com/therootcompany/golib.git synced 2026-11-04 06:55:00 +00:00
Code Issues Projects Releases Wiki Activity

feat(iplist): allow custom HTTP clients

Browse Source
This commit is contained in:
AJ ONeal 2026-09-09 01:07:30 -06:00
parent 5b92532901
commit 9cbf82fc16
No known key found for this signature in database
2 changed files with 41 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
net/iplist/iplist.go
View File

@ -33,17 +33,26 @@ const (
// MaxNestedSources levels. URL responses are cached below cacheDir and the
// cache remains available across process restarts.
func Load(ctx context.Context, source, cacheDir string) ([]string, error) {
return load(ctx, strings.TrimSpace(source), cacheDir, make(map[string]struct{}), 0)
return LoadWithClient(ctx, source, cacheDir, nil)
}
func load(ctx context.Context, source, cacheDir string, seen map[string]struct{}, depth int) ([]string, error) {
// LoadWithClient is Load with an optional HTTP client. A nil client uses the
// internal client with the package's standard timeout and transport policy.
func LoadWithClient(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, 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
}
@ -51,7 +60,7 @@ func load(ctx context.Context, source, cacheDir string, seen map[string]struct{}
for _, entry := range entries {
parsed, parseErr := url.Parse(entry)
if parseErr == nil && (parsed.Scheme == "http" || parsed.Scheme == "https") {
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
}
@ -63,7 +72,7 @@ func load(ctx context.Context, source, cacheDir string, seen map[string]struct{}
return result, nil
}
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 != "" {
@ -95,7 +104,7 @@ func readSource(ctx context.Context, source, cacheDir string, seen map[string]st
digest := sha256.Sum256([]byte(key))
path := filepath.Join(cacheDir, "ip-sources", hex.EncodeToString(digest[:8])+".tsv")
previous, previousErr := os.ReadFile(path)
cache := httpcache.NewWith(parsed.String(), path, https.NewInternalClient())
cache := httpcache.NewWith(parsed.String(), path, client)
cache.Header = header
cache.MaxAge = defaultMaxAge
cache.MaxBytes = defaultMaxBytes

26
net/iplist/iplist_test.go
View File

@ -17,6 +17,32 @@ func TestParseTSV(t *testing.T) {
}
}
func TestLoadWithClient(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
_, _ = w.Write([]byte("192.0.2.7\n"))
}))
defer server.Close()
called := false
client := &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
called = true
return http.DefaultTransport.RoundTrip(r)
})}
entries, err := LoadWithClient(context.Background(), server.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) { _, _ = w.Write([]byte("192.0.2.2\n")) }))