Stop reusing hosts after idle leftovers

pkg/protocols/common/protocolstate/dialers.go

@@ -16,6 +16,7 @@ type Dialers struct {
 	HTTPClientPool             *HTTPPool
 	PerHostRateLimitPool       any // *httpclientpool.PerHostRateLimitPool
 	HTTPToHTTPSPortTracker     any // *httpclientpool.HTTPToHTTPSPortTracker
+	HTTPDesyncHosts            *ExpiringSet
 	NetworkPolicy              *networkpolicy.NetworkPolicy
 	LocalFileAccessAllowed     bool
 	RestrictLocalNetworkAccess bool

pkg/protocols/common/protocolstate/expiring_set.go

@@ -0,0 +1,89 @@
+package protocolstate
+
+import (
+	"slices"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// ExpiringSet is a lock-free-read set of string keys with per-entry expiry.
+// Cleanup is amortized across writes and also performed by Keys.
+type ExpiringSet struct {
+	values          sync.Map // string -> expiry UnixNano
+	cleanupInterval time.Duration
+	lastCleanup     atomic.Int64
+}
+
+func NewExpiringSet(cleanupInterval time.Duration) *ExpiringSet {
+	set := &ExpiringSet{cleanupInterval: cleanupInterval}
+	set.lastCleanup.Store(time.Now().UnixNano())
+	return set
+}
+
+func (s *ExpiringSet) Store(key string, ttl time.Duration) {
+	s.StoreUntil(key, time.Now().Add(ttl))
+}
+
+func (s *ExpiringSet) StoreUntil(key string, expiry time.Time) {
+	if s == nil || key == "" {
+		return
+	}
+	s.values.Store(key, expiry.UnixNano())
+	s.maybeCleanup(time.Now())
+}
+
+func (s *ExpiringSet) Contains(key string) bool {
+	return s.containsAt(key, time.Now())
+}
+
+func (s *ExpiringSet) containsAt(key string, now time.Time) bool {
+	if s == nil || key == "" {
+		return false
+	}
+	value, ok := s.values.Load(key)
+	if !ok {
+		return false
+	}
+	expiresAt, ok := value.(int64)
+	if !ok || now.UnixNano() >= expiresAt {
+		s.values.CompareAndDelete(key, value)
+		return false
+	}
+	return true
+}
+
+func (s *ExpiringSet) Keys() []string {
+	return s.keysAt(time.Now())
+}
+
+func (s *ExpiringSet) keysAt(now time.Time) []string {
+	if s == nil {
+		return nil
+	}
+	var keys []string
+	s.values.Range(func(key, value any) bool {
+		name, keyOK := key.(string)
+		expiresAt, expiryOK := value.(int64)
+		if keyOK && expiryOK && now.UnixNano() < expiresAt {
+			keys = append(keys, name)
+		} else {
+			s.values.CompareAndDelete(key, value)
+		}
+		return true
+	})
+	slices.Sort(keys)
+	return keys
+}
+
+func (s *ExpiringSet) maybeCleanup(now time.Time) {
+	if s.cleanupInterval <= 0 {
+		return
+	}
+	last := s.lastCleanup.Load()
+	if now.UnixNano()-last < s.cleanupInterval.Nanoseconds() ||
+		!s.lastCleanup.CompareAndSwap(last, now.UnixNano()) {
+		return
+	}
+	s.keysAt(now)
+}

pkg/protocols/common/protocolstate/expiring_set_test.go

@@ -0,0 +1,102 @@
+package protocolstate
+
+import (
+	"fmt"
+	"strconv"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestExpiringSetContainsAndExpires(t *testing.T) {
+	set := NewExpiringSet(time.Minute)
+	now := time.Now()
+	set.StoreUntil("active", now.Add(time.Minute))
+	set.StoreUntil("expired", now.Add(-time.Second))
+
+	require.True(t, set.containsAt("active", now))
+	require.False(t, set.containsAt("expired", now))
+	_, exists := set.values.Load("expired")
+	require.False(t, exists)
+}
+
+func TestExpiringSetKeysSortsAndCleans(t *testing.T) {
+	set := NewExpiringSet(time.Minute)
+	now := time.Now()
+	set.StoreUntil("b", now.Add(time.Minute))
+	set.StoreUntil("a", now.Add(time.Minute))
+	set.StoreUntil("expired", now.Add(-time.Second))
+
+	require.Equal(t, []string{"a", "b"}, set.keysAt(now))
+	_, exists := set.values.Load("expired")
+	require.False(t, exists)
+}
+
+func TestExpiringSetRefreshesTTL(t *testing.T) {
+	set := NewExpiringSet(time.Minute)
+	now := time.Now()
+	set.StoreUntil("host", now.Add(time.Minute))
+	set.StoreUntil("host", now.Add(2*time.Minute))
+
+	require.True(t, set.containsAt("host", now.Add(90*time.Second)))
+	require.False(t, set.containsAt("host", now.Add(3*time.Minute)))
+}
+
+func TestExpiringSetEmptyAndNil(t *testing.T) {
+	var nilSet *ExpiringSet
+	require.False(t, nilSet.Contains("host"))
+	require.Empty(t, nilSet.Keys())
+
+	set := NewExpiringSet(time.Minute)
+	set.Store("", time.Minute)
+	require.False(t, set.Contains(""))
+	require.Empty(t, set.Keys())
+}
+
+func TestExpiringSetConcurrentAccess(t *testing.T) {
+	set := NewExpiringSet(time.Minute)
+	var done atomic.Int64
+
+	t.Run("parallel", func(t *testing.T) {
+		for i := range 100 {
+			i := i
+			t.Run(strconv.Itoa(i), func(t *testing.T) {
+				t.Parallel()
+				key := fmt.Sprintf("host-%d", i%10)
+				set.Store(key, time.Minute)
+				require.True(t, set.Contains(key))
+				done.Add(1)
+			})
+		}
+	})
+	require.Equal(t, int64(100), done.Load())
+}
+
+func TestExpiringSetConcurrentRefreshSurvivesExpiredCleanup(t *testing.T) {
+	for range 1000 {
+		set := NewExpiringSet(time.Minute)
+		now := time.Now()
+		set.StoreUntil("host", now.Add(-time.Second))
+
+		start := make(chan struct{})
+		done := make(chan struct{}, 2)
+		go func() {
+			<-start
+			_ = set.containsAt("host", now)
+			done <- struct{}{}
+		}()
+		go func() {
+			<-start
+			set.StoreUntil("host", now.Add(time.Minute))
+			done <- struct{}{}
+		}()
+		close(start)
+		<-done
+		<-done
+
+		require.True(t, set.containsAt("host", now),
+			"cleanup of the expired value must not delete a concurrent refresh")
+	}
+}

pkg/protocols/common/protocolstate/state.go

@@ -216,6 +216,7 @@ func initDialers(options *types.Options) error {
 		Fastdialer:                 dialer,
 		NetworkPolicy:              networkPolicy,
 		HTTPClientPool:             httpClientPool,
+		HTTPDesyncHosts:            NewExpiringSet(time.Minute),
 		LocalFileAccessAllowed:     options.AllowLocalFileAccess,
 		RestrictLocalNetworkAccess: options.RestrictLocalNetworkAccess,
 		ExcludeTargets:             options.ExcludeTargets,

pkg/protocols/http/httpclientpool/clientpool.go

@@ -124,6 +124,7 @@ func (t *connTrackingTransport) RoundTrip(req *http.Request) (*http.Response, er
 	// Compute the host key once (URL is already parsed) so the GotConn hook can
 	// update both the global counters and the per-host bucket from one trace.
 	host := normalizeHost(req.URL)
+	var used atomic.Pointer[desyncConn]
 	trace := &httptrace.ClientTrace{
 		GotConn: func(info httptrace.GotConnInfo) {
 			if info.Reused {
@@ -132,6 +133,18 @@ func (t *connTrackingTransport) RoundTrip(req *http.Request) (*http.Response, er
 				connStats.New.Add(1)
 			}
 			recordHostConn(host, info.Reused)
+			if tracked, ok := info.Conn.(*desyncConn); ok {
+				tracked.markBusy()
+				used.Store(tracked)
+			}
+		},
+		PutIdleConn: func(err error) {
+			if err != nil {
+				return
+			}
+			if tracked := used.Load(); tracked != nil {
+				tracked.markIdle()
+			}
 		},
 	}
 	req = req.WithContext(httptrace.WithClientTrace(req.Context(), trace))
@@ -303,6 +316,13 @@ func wrappedGet(options *types.Options, configuration *Configuration, host strin
 		clientKey += ":" + host
 	}
 
+	// A host caught sending unsolicited idle responses gets its own cache
+	// entry, so the keep-alive client already cached for it is not reused.
+	noReuse := IsHostDesynced(options, host)
+	if noReuse {
+		clientKey += ":noreuse"
+	}
+
 	// Fast path: lock-free cache hit.
 	if !hasExplicitJar {
 		if client, ok := pool.GetClient(clientKey); ok {
@@ -330,7 +350,8 @@ func wrappedGet(options *types.Options, configuration *Configuration, host strin
 		maxIdleConns = configuration.Threads
 	}
 
-	disableKeepAlives := configuration.Connection != nil && configuration.Connection.DisableKeepAlive
+	disableKeepAlives := noReuse ||
+		(configuration.Connection != nil && configuration.Connection.DisableKeepAlive)
 
 	responseHeaderTimeout := options.GetTimeouts().HttpResponseHeaderTimeout
 	if configuration.ResponseHeaderTimeout != 0 {
@@ -367,18 +388,40 @@ func wrappedGet(options *types.Options, configuration *Configuration, host strin
 			return nil, errors.Wrap(err, "could not create client certificate")
 		}
 
+		// Wrap keep-alive dials so an idle leftover can close the connection
+		// and quarantine the host. Proxies are left alone: CONNECT is performed
+		// by net/http outside a round trip we can mark busy/idle.
+		proxied := options.AliveHttpProxy != "" || options.AliveSocksProxy != ""
+		trackDesync := func(dial func(context.Context, string, string) (net.Conn, error)) func(context.Context, string, string) (net.Conn, error) {
+			if disableKeepAlives || proxied {
+				return dial
+			}
+			return func(ctx context.Context, network, addr string) (net.Conn, error) {
+				conn, err := dial(ctx, network, addr)
+				if err != nil {
+					return conn, err
+				}
+				return newDesyncConn(conn, addr, func(host string) {
+					if dialers.HTTPDesyncHosts != nil {
+						dialers.HTTPDesyncHosts.Store(host, desyncedHostTTL)
+					}
+					countDesyncedResponse()
+				}), nil
+			}
+		}
+
 		transport := &http.Transport{
 			ForceAttemptHTTP2: options.ForceAttemptHTTP2,
-			DialContext:       dialers.Fastdialer.Dial,
-			DialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
+			DialContext:       trackDesync(dialers.Fastdialer.Dial),
+			DialTLSContext: trackDesync(func(ctx context.Context, network, addr string) (net.Conn, error) {
 				if options.TlsImpersonate {
 					return dialers.Fastdialer.DialTLSWithConfigImpersonate(ctx, network, addr, tlsConfig, impersonate.Random, nil)
 				}
 				if options.HasClientCertificates() || options.ForceAttemptHTTP2 {
 					return dialers.Fastdialer.DialTLSWithConfig(ctx, network, addr, tlsConfig)
 				}
 				return dialers.Fastdialer.DialTLS(ctx, network, addr)
-			},
+			}),
 			MaxIdleConns:          maxIdleConns,
 			MaxIdleConnsPerHost:   maxIdleConnsPerHost,
 			MaxConnsPerHost:       maxConnsPerHost,

pkg/protocols/http/httpclientpool/desync_conn.go

@@ -0,0 +1,83 @@
+package httpclientpool
+
+import (
+	"crypto/tls"
+	"net"
+	"sync"
+	"sync/atomic"
+)
+
+// desyncConn watches for bytes that arrive from a server while the connection is
+// not serving any request. A server has nothing to answer at that point, so
+// whatever it sends belongs to an earlier exchange.
+//
+// Idleness comes from net/http: GotConn fires when a request takes the
+// connection and PutIdleConn fires when it returns to the pool. Bytes in that
+// window are unsolicited by the same criterion net/http uses internally. The
+// check is one atomic load in Read, identical over plaintext and TLS because
+// the wrapper sits above the handshake.
+//
+// net/http already closes a connection when it notices this, but only logs it.
+// Observing it here is what lets the pool stop reusing connections to that host.
+// A leftover that arrives while a request is outstanding is indistinguishable
+// from that request's real reply, so this cannot save the first wrong response.
+// Closing and quarantining only stops the host from doing it again.
+type desyncConn struct {
+	net.Conn
+	host      string
+	onPoison  func(host string)
+	idle      atomic.Bool
+	poisoned  atomic.Bool
+	closeOnce sync.Once
+}
+
+func newDesyncConn(conn net.Conn, addr string, onPoison func(host string)) net.Conn {
+	if conn == nil || negotiatedHTTP2(conn) {
+		return conn
+	}
+	return &desyncConn{
+		Conn:     conn,
+		host:     desyncedHostKey(addr),
+		onPoison: onPoison,
+	}
+}
+
+func (c *desyncConn) Read(b []byte) (int, error) {
+	n, err := c.Conn.Read(b)
+	if n > 0 && c.idle.Load() {
+		c.poison()
+	}
+	return n, err
+}
+
+func (c *desyncConn) markBusy() {
+	c.idle.Store(false)
+}
+
+func (c *desyncConn) markIdle() {
+	c.idle.Store(true)
+}
+
+func (c *desyncConn) Close() error {
+	c.idle.Store(false)
+	return c.Conn.Close()
+}
+
+func (c *desyncConn) poison() {
+	if !c.poisoned.CompareAndSwap(false, true) {
+		return
+	}
+	c.idle.Store(false)
+	if c.onPoison != nil {
+		c.onPoison(c.host)
+	}
+	c.closeOnce.Do(func() { _ = c.Conn.Close() })
+}
+
+func negotiatedHTTP2(conn net.Conn) bool {
+	handshaked, ok := conn.(interface{ ConnectionState() tls.ConnectionState })
+	if !ok {
+		return false
+	}
+	return handshaked.ConnectionState().NegotiatedProtocol == "h2"
+}