SSRF-safe recipe fetcher #65

2 changed files with 131 additions and 12 deletions
Showing only changes of commit b013b542e3 - Show all commits
+70 -12
View File
@@ -10,10 +10,12 @@ import (
"io" "io"
"net" "net"
"net/http" "net/http"
"net/http/httptrace"
"net/netip" "net/netip"
"net/url" "net/url"
"strings" "strings"
"sync" "sync"
"sync/atomic"
"time" "time"
) )
@@ -99,9 +101,12 @@ type Fetcher struct {
sleep func(ctx context.Context, d time.Duration) error sleep func(ctx context.Context, d time.Duration) error
mu sync.Mutex mu sync.Mutex
// lastSent is when each host's latest request was let through; only // lastSent is when each host's latest request was let through; a
// requests actually sent advance it. // request that is never sent hands its slot back.
lastSent map[string]time.Time lastSent map[string]time.Time
// released is closed, then replaced, whenever a slot is handed back,
// waking requests waiting for their turn.
released chan struct{}
cache map[string]cacheEntry cache map[string]cacheEntry
} }
@@ -125,6 +130,7 @@ func newFetcher(resolver Resolver, dial dialFunc, tlsConfig *tls.Config) *Fetche
now: time.Now, now: time.Now,
sleep: sleep, sleep: sleep,
lastSent: map[string]time.Time{}, lastSent: map[string]time.Time{},
released: make(chan struct{}),
cache: map[string]cacheEntry{}, cache: map[string]cacheEntry{},
} }
transport := &http.Transport{ transport := &http.Transport{
@@ -237,11 +243,15 @@ func (f *Fetcher) store(key string, page Page) {
} }
// waitTurn blocks until a request to host may be sent, keeping requests to // waitTurn blocks until a request to host may be sent, keeping requests to
// one host at least HostInterval apart. A wait holds no reservation: the // one host at least HostInterval apart. A wait holds no reservation, and a
// host's slot advances only when a request is let through, so a wait that // request ending with ctx before its turn takes no slot. The request let
// ends with ctx does not delay later requests. // through takes the host's slot and must call release if it is not sent,
func (f *Fetcher) waitTurn(ctx context.Context, host string) error { // so that it delays no other request.
func (f *Fetcher) waitTurn(ctx context.Context, host string) (release func(), err error) {
for { for {
if err := ctx.Err(); err != nil {
return nil, err
}
f.mu.Lock() f.mu.Lock()
now := f.now() now := f.now()
last, sent := f.lastSent[host] last, sent := f.lastSent[host]
@@ -249,15 +259,55 @@ func (f *Fetcher) waitTurn(ctx context.Context, host string) error {
if !sent || wait <= 0 { if !sent || wait <= 0 {
f.lastSent[host] = now f.lastSent[host] = now
f.mu.Unlock() f.mu.Unlock()
return nil return func() { f.release(host, now, last, sent) }, nil
} }
released := f.released
f.mu.Unlock() f.mu.Unlock()
if err := f.sleep(ctx, wait); err != nil { if err := f.waitOrRelease(ctx, wait, released); err != nil {
return err return nil, err
} }
} }
} }
// release hands back the slot taken at taken for a request that was not
// sent, restoring the host's previous send time, and wakes waiting requests.
func (f *Fetcher) release(host string, taken, prev time.Time, hadPrev bool) {
f.mu.Lock()
defer f.mu.Unlock()
if !f.lastSent[host].Equal(taken) {
// A later request already holds the slot.
return
}
if hadPrev {
f.lastSent[host] = prev
} else {
delete(f.lastSent, host)
}
close(f.released)
f.released = make(chan struct{})
}
// waitOrRelease sleeps for d, returning early without error when a slot is
// handed back so the caller can look again.
func (f *Fetcher) waitOrRelease(ctx context.Context, d time.Duration, released <-chan struct{}) error {
sleepCtx, cancel := context.WithCancel(ctx)
defer cancel()
go func() {
select {
case <-released:
cancel()
case <-sleepCtx.Done():
}
}()
err := f.sleep(sleepCtx, d)
select {
case <-released:
return nil
default:
return err
}
}
// dialPublic resolves address's host once, refuses unless every answer is // dialPublic resolves address's host once, refuses unless every answer is
// public, then dials the checked addresses in turn. // public, then dials the checked addresses in turn.
func (f *Fetcher) dialPublic(ctx context.Context, network, address string) (net.Conn, error) { func (f *Fetcher) dialPublic(ctx context.Context, network, address string) (net.Conn, error) {
@@ -301,17 +351,25 @@ func (f *Fetcher) dialPublic(ctx context.Context, network, address string) (net.
} }
// limitedTransport waits for each request's per-host rate-limit slot, // limitedTransport waits for each request's per-host rate-limit slot,
// redirect hops included, before sending it. // redirect hops included, before sending it. A request that fails before
// its headers are written was not sent and hands its slot back.
type limitedTransport struct { type limitedTransport struct {
f *Fetcher f *Fetcher
next http.RoundTripper next http.RoundTripper
} }
func (t *limitedTransport) RoundTrip(req *http.Request) (*http.Response, error) { func (t *limitedTransport) RoundTrip(req *http.Request) (*http.Response, error) {
if err := t.f.waitTurn(req.Context(), strings.ToLower(req.URL.Hostname())); err != nil { release, err := t.f.waitTurn(req.Context(), strings.ToLower(req.URL.Hostname()))
if err != nil {
return nil, err return nil, err
} }
return t.next.RoundTrip(req) var sent atomic.Bool
trace := &httptrace.ClientTrace{WroteHeaders: func() { sent.Store(true) }}
resp, err := t.next.RoundTrip(req.WithContext(httptrace.WithClientTrace(req.Context(), trace)))
if err != nil && !sent.Load() {
release()
}
return resp, err
} }
// checkURL refuses anything but an absolute https:// URL with a host. // checkURL refuses anything but an absolute https:// URL with a host.
+61
View File
@@ -592,3 +592,64 @@ func TestFetchCanceledWaitHoldsNoSlot(t *testing.T) {
t.Errorf("waits = %v, want none: canceled waits held slots", got) t.Errorf("waits = %v, want none: canceled waits held slots", got)
} }
} }
// TestFetchUnsentRequestHoldsNoSlot covers requests to a host with no
// previous request that are never sent — canceled before Fetch, or failing
// to connect — and checks that the next request is not delayed.
func TestFetchUnsentRequestHoldsNoSlot(t *testing.T) {
t.Parallel()
const unroutableIP = "93.184.215.99"
answers := map[string][][]string{"example.com": {{unroutableIP}, {publicIP}}}
env := newTestEnv(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
fmt.Fprint(w, "ok")
}), answers)
clock := installClock(env.f)
ctx, cancel := context.WithCancel(t.Context())
cancel()
if _, err := env.f.Fetch(ctx, "https://example.com/canceled"); !errors.Is(err, context.Canceled) {
t.Fatalf("canceled Fetch error = %v, want context.Canceled", err)
}
if _, err := env.f.Fetch(t.Context(), "https://example.com/unroutable"); err == nil || !strings.Contains(err.Error(), "no route") {
t.Fatalf("unroutable Fetch error = %v, want a dial failure", err)
}
if _, err := env.f.Fetch(t.Context(), "https://example.com/next"); err != nil {
t.Fatalf("next Fetch: %v", err)
}
if got := clock.waits(); len(got) != 0 {
t.Errorf("waits = %v, want none: unsent requests held slots", got)
}
if got := env.dials(); len(got) != 2 {
t.Errorf("dials = %v, want the unroutable and the next request only", got)
}
}
// TestReleaseWakesWaiter checks that a request waiting behind a slot that
// is handed back goes ahead at once instead of sleeping out the interval.
func TestReleaseWakesWaiter(t *testing.T) {
t.Parallel()
f := newFetcher(&stubResolver{}, nil, nil)
release, err := f.waitTurn(t.Context(), "example.com")
if err != nil {
t.Fatalf("first waitTurn: %v", err)
}
done := make(chan error, 1)
start := time.Now()
go func() {
_, err := f.waitTurn(t.Context(), "example.com")
done <- err
}()
time.Sleep(50 * time.Millisecond)
release()
select {
case err := <-done:
if err != nil {
t.Fatalf("waiting waitTurn: %v", err)
}
if waited := time.Since(start); waited >= HostInterval {
t.Errorf("waiter took %s, want less than %s", waited, HostInterval)
}
case <-time.After(HostInterval / 2):
t.Fatal("waiter still waiting after the slot was handed back")
}
}