service/http_client_transport_test.go

Rama main · réplica de lectura del repositorio original.

package service

import (
	"bytes"
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/relaykit/dto"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/http2"
	"golang.org/x/net/http2/hpack"
)

func withRelayHTTPTransportSettings(t *testing.T) {
	t.Helper()
	prevMaxIdle := common.RelayMaxIdleConns
	prevPerHost := common.RelayMaxIdleConnsPerHost
	prevTimeout := common.RelayIdleConnTimeout
	prevPingIdle, prevPingTimeout := common.RelayHTTP2PingIdleTimeout, common.RelayHTTP2PingTimeout
	common.RelayMaxIdleConns = 500
	common.RelayMaxIdleConnsPerHost = 100
	common.RelayIdleConnTimeout = 90
	common.RelayHTTP2PingIdleTimeout, common.RelayHTTP2PingTimeout = 1, 1
	t.Cleanup(func() {
		common.RelayMaxIdleConns = prevMaxIdle
		common.RelayMaxIdleConnsPerHost = prevPerHost
		common.RelayIdleConnTimeout = prevTimeout
		common.RelayHTTP2PingIdleTimeout, common.RelayHTTP2PingTimeout = prevPingIdle, prevPingTimeout
	})
}

func initDefaultHTTPClientFixture(t *testing.T) *http.Client {
	t.Helper()
	withRelayHTTPTransportSettings(t)
	if httpClient == nil {
		InitHttpClient()
	} else {
		ResetProxyClientCache()
	}
	require.NotNil(t, httpClient)
	t.Cleanup(ResetProxyClientCache)
	return httpClient
}

func TestShardedRoundTripperPerOriginRotation(t *testing.T) {
	s := &shardedRoundTripper{n: 4}
	originA := "https://a.example:443"
	originB := "https://b.example:443"

	gotA := make([]uint32, 0, 8)
	for range 8 {
		gotA = append(gotA, s.pickShard(originA))
	}
	assert.Equal(t, []uint32{0, 1, 2, 3, 0, 1, 2, 3}, gotA)

	gotB := make([]uint32, 0, 4)
	for range 4 {
		gotB = append(gotB, s.pickShard(originB))
	}
	assert.Equal(t, []uint32{0, 1, 2, 3}, gotB, "independent origins must have independent counters")

	var wg sync.WaitGroup
	const workers = 32
	const perWorker = 50
	var badShardCount atomic.Uint32
	wg.Add(workers)
	for range workers {
		go func() {
			defer wg.Done()
			for range perWorker {
				idx := s.pickShard(originA)
				if idx >= 4 {
					badShardCount.Add(1)
				}
			}
		}()
	}
	wg.Wait()
	assert.Equal(t, uint32(0), badShardCount.Load())
}

func TestOriginKeyUsesSchemeAndHost(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "HTTPS://Example.COM:8443/path", nil)
	assert.Equal(t, "https://Example.COM:8443", originKey(req))
}

func testTLSClientConfig(t *testing.T, server *httptest.Server) *tls.Config {
	t.Helper()
	pool := x509.NewCertPool()
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: server.Certificate().Raw})
	require.True(t, pool.AppendCertsFromPEM(certPEM))
	return &tls.Config{RootCAs: pool}
}

func startHTTP2TLSServer(t *testing.T, handler http.Handler) *httptest.Server {
	t.Helper()
	server := httptest.NewUnstartedServer(handler)
	server.EnableHTTP2 = true
	server.StartTLS()
	t.Cleanup(server.Close)
	return server
}

func drainClose(t *testing.T, resp *http.Response) {
	t.Helper()
	require.NotNil(t, resp)
	_, err := io.Copy(io.Discard, resp.Body)
	require.NoError(t, err)
	require.NoError(t, resp.Body.Close())
}

func TestAutoOneShardNegotiatesHTTP2SingleConnection(t *testing.T) {
	withRelayHTTPTransportSettings(t)

	var mu sync.Mutex
	addrs := make(map[string]struct{})
	var sawHTTP2 atomic.Bool

	server := startHTTP2TLSServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		addrs[r.RemoteAddr] = struct{}{}
		mu.Unlock()
		if r.ProtoMajor == 2 {
			sawHTTP2.Store(true)
		}
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok"))
	}))

	client := newHTTPClientWithPolicyAndTLS(defaultHTTPTransportPolicy(), testTLSClientConfig(t, server))
	for range 4 {
		resp, err := client.Get(server.URL)
		require.NoError(t, err)
		assert.Equal(t, 2, resp.ProtoMajor)
		drainClose(t, resp)
	}

	mu.Lock()
	defer mu.Unlock()
	assert.True(t, sawHTTP2.Load())
	assert.Len(t, addrs, 1, "auto+1 must reuse a single HTTP/2 connection")
}

func TestFourShardHTTP2ReusesExactlyFourConnections(t *testing.T) {
	withRelayHTTPTransportSettings(t)

	var mu sync.Mutex
	addrs := make(map[string]struct{})
	var nonHTTP2Count atomic.Uint32

	server := startHTTP2TLSServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		addrs[r.RemoteAddr] = struct{}{}
		mu.Unlock()
		if r.ProtoMajor != 2 {
			nonHTTP2Count.Add(1)
		}
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok"))
	}))

	policy := HTTPTransportPolicy{Protocol: dto.HTTPProtocolAuto, Shards: 4}
	client := newHTTPClientWithPolicyAndTLS(policy, testTLSClientConfig(t, server))
	for range 8 {
		resp, err := client.Get(server.URL)
		require.NoError(t, err)
		assert.Equal(t, 2, resp.ProtoMajor)
		drainClose(t, resp)
	}

	mu.Lock()
	defer mu.Unlock()
	assert.Equal(t, uint32(0), nonHTTP2Count.Load())
	assert.Len(t, addrs, 4, "four shards must establish and reuse exactly four connections")
}

func TestForcedHTTP1AgainstHTTP2Server(t *testing.T) {
	withRelayHTTPTransportSettings(t)

	var nonHTTP1Count atomic.Uint32
	server := startHTTP2TLSServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.ProtoMajor != 1 {
			nonHTTP1Count.Add(1)
		}
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok"))
	}))

	policy := HTTPTransportPolicy{Protocol: dto.HTTPProtocolHTTP1, Shards: 1}
	client := newHTTPClientWithPolicyAndTLS(policy, testTLSClientConfig(t, server))
	resp, err := client.Get(server.URL)
	require.NoError(t, err)
	assert.Equal(t, 1, resp.ProtoMajor)
	drainClose(t, resp)
	assert.Equal(t, uint32(0), nonHTTP1Count.Load())

	transport, ok := client.Transport.(*http.Transport)
	require.True(t, ok)
	assert.False(t, transport.DisableKeepAlives)
	assert.False(t, transport.ForceAttemptHTTP2)
	assert.NotNil(t, transport.TLSNextProto)
	assert.Len(t, transport.TLSNextProto, 0)
}

func TestForcedHTTP1ConcurrentDistinctConnections(t *testing.T) {
	withRelayHTTPTransportSettings(t)

	const k = 8
	var mu sync.Mutex
	addrs := make(map[string]struct{})
	arrived := make(chan struct{}, k)
	release := make(chan struct{})
	var nonHTTP1Count atomic.Uint32

	server := startHTTP2TLSServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		addrs[r.RemoteAddr] = struct{}{}
		mu.Unlock()
		if r.ProtoMajor != 1 {
			nonHTTP1Count.Add(1)
		}
		arrived <- struct{}{}
		<-release
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok"))
	}))

	policy := HTTPTransportPolicy{Protocol: dto.HTTPProtocolHTTP1, Shards: 1}
	client := newHTTPClientWithPolicyAndTLS(policy, testTLSClientConfig(t, server))

	errCh := make(chan error, k)
	for range k {
		go func() {
			resp, err := client.Get(server.URL)
			if err != nil {
				errCh <- err
				return
			}
			if resp.ProtoMajor != 1 {
				errCh <- fmt.Errorf("expected HTTP/1.x, got %s", resp.Proto)
				_ = resp.Body.Close()
				return
			}
			_, _ = io.Copy(io.Discard, resp.Body)
			_ = resp.Body.Close()
			errCh <- nil
		}()
	}

	for range k {
		<-arrived
	}
	mu.Lock()
	activeAddrs := len(addrs)
	mu.Unlock()
	close(release)

	for range k {
		require.NoError(t, <-errCh)
	}
	assert.Equal(t, uint32(0), nonHTTP1Count.Load())
	assert.Equal(t, k, activeAddrs, "all HTTP/1.1 handlers active together must use K distinct connections")
}

func TestHTTPClientCachePolicyAndCompatibility(t *testing.T) {
	defaultClient := initDefaultHTTPClientFixture(t)

	compat, err := GetHttpClientWithProxy("")
	require.NoError(t, err)
	aware, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{})
	require.NoError(t, err)
	assert.Same(t, defaultClient, compat)
	assert.Same(t, compat, aware)
	assert.Same(t, GetHttpClient(), aware)

	http1, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{HTTPProtocol: dto.HTTPProtocolHTTP1})
	require.NoError(t, err)
	assert.NotSame(t, aware, http1)

	sharded, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{HTTP2ConnectionShards: 4})
	require.NoError(t, err)
	assert.NotSame(t, aware, sharded)
	assert.NotSame(t, http1, sharded)

	proxyA := "http://proxy.example:8080"
	proxyAlias := "http://proxy.example:8080/"
	clientA, err := GetHttpClientWithProxy(proxyA)
	require.NoError(t, err)
	clientAlias, err := GetHttpClientWithProxy(proxyAlias)
	require.NoError(t, err)
	assert.Same(t, clientA, clientAlias, "canonical proxy aliases must share the default policy client")

	proxyHTTP1, err := GetHttpClientWithProxySettings(proxyA, dto.ChannelSettings{HTTPProtocol: dto.HTTPProtocolHTTP1})
	require.NoError(t, err)
	assert.NotSame(t, clientA, proxyHTTP1)

	insecure, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{TLSInsecureSkipVerify: true})
	require.NoError(t, err)
	assert.NotSame(t, aware, insecure, "skip-verify channels must not share the verifying client")
	insecureHTTP1, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{HTTPProtocol: dto.HTTPProtocolHTTP1, TLSInsecureSkipVerify: true})
	require.NoError(t, err)
	assert.NotSame(t, http1, insecureHTTP1)
	proxyInsecure, err := GetHttpClientWithProxySettings(proxyA, dto.ChannelSettings{TLSInsecureSkipVerify: true})
	require.NoError(t, err)
	assert.NotSame(t, clientA, proxyInsecure)
}

func TestChannelTLSInsecureSkipVerifyAcceptsUntrustedUpstream(t *testing.T) {
	initDefaultHTTPClientFixture(t)
	upstream := startHTTP2TLSServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	connectProxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodConnect {
			w.WriteHeader(http.StatusMethodNotAllowed)
			return
		}
		target, err := net.Dial("tcp", r.Host)
		if err != nil {
			w.WriteHeader(http.StatusBadGateway)
			return
		}
		defer target.Close()
		conn, _, err := http.NewResponseController(w).Hijack()
		if err != nil {
			return
		}
		defer conn.Close()
		_, _ = conn.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n"))
		go func() { _, _ = io.Copy(target, conn) }()
		_, _ = io.Copy(conn, target)
	}))
	t.Cleanup(connectProxy.Close)

	skip := dto.ChannelSettings{TLSInsecureSkipVerify: true}
	testCases := []struct {
		name      string
		proxy     string
		settings  dto.ChannelSettings
		wantProto int // zero means certificate verification must fail
	}{
		{name: "default direct verifies", settings: dto.ChannelSettings{}},
		{name: "default proxy verifies", proxy: connectProxy.URL, settings: dto.ChannelSettings{}},
		{name: "http1 verifies", settings: dto.ChannelSettings{HTTPProtocol: dto.HTTPProtocolHTTP1}},
		{name: "skip direct", settings: skip, wantProto: 2},
		{name: "skip proxy", proxy: connectProxy.URL, settings: skip, wantProto: 2},
		{name: "skip http1", settings: dto.ChannelSettings{HTTPProtocol: dto.HTTPProtocolHTTP1, TLSInsecureSkipVerify: true}, wantProto: 1},
		{name: "skip sharded", settings: dto.ChannelSettings{HTTP2ConnectionShards: 3, TLSInsecureSkipVerify: true}, wantProto: 2},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			client, err := GetHttpClientWithProxySettings(tc.proxy, tc.settings)
			require.NoError(t, err)
			resp, err := client.Get(upstream.URL)
			if tc.wantProto == 0 {
				var verifyErr *tls.CertificateVerificationError
				require.ErrorAs(t, err, &verifyErr)
				return
			}
			require.NoError(t, err)
			assert.Equal(t, tc.wantProto, resp.ProtoMajor)
			drainClose(t, resp)
		})
	}
}

func TestHTTPClientCacheConcurrentGetOrCreate(t *testing.T) {
	initDefaultHTTPClientFixture(t)

	proxyURL := "http://concurrent-proxy.example:9090"
	const workers = 32
	results := make([]*http.Client, workers)
	errs := make([]error, workers)
	var wg sync.WaitGroup
	wg.Add(workers)
	for i := range workers {
		i := i
		go func() {
			defer wg.Done()
			client, err := GetHttpClientWithProxySettings(proxyURL, dto.ChannelSettings{HTTP2ConnectionShards: 3})
			errs[i] = err
			results[i] = client
		}()
	}
	wg.Wait()
	for i := range workers {
		require.NoError(t, errs[i])
	}
	for i := 1; i < workers; i++ {
		assert.Same(t, results[0], results[i])
	}
}

type closeCountingRoundTripper struct {
	closes atomic.Int32
}

func (c *closeCountingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	return &http.Response{
		StatusCode: http.StatusOK,
		Body:       http.NoBody,
		Header:     make(http.Header),
		Request:    req,
	}, nil
}

func (c *closeCountingRoundTripper) CloseIdleConnections() {
	c.closes.Add(1)
}

func TestShardedRoundTripperCloseIdleConnectionsFansOut(t *testing.T) {
	trackers := []*closeCountingRoundTripper{{}, {}, {}}
	shards := make([]http.RoundTripper, len(trackers))
	for i, tracker := range trackers {
		shards[i] = tracker
	}
	s := &shardedRoundTripper{shards: shards, n: uint32(len(shards))}
	s.CloseIdleConnections()
	for _, tracker := range trackers {
		assert.Equal(t, int32(1), tracker.closes.Load())
	}
}

func TestInvalidateProxyClientClosesAllPolicyVariants(t *testing.T) {
	initDefaultHTTPClientFixture(t)

	proxyURL := "http://invalidate-proxy.example:8080"
	defaultClient, err := GetHttpClientWithProxy(proxyURL)
	require.NoError(t, err)
	http1Client, err := GetHttpClientWithProxySettings(proxyURL, dto.ChannelSettings{HTTPProtocol: dto.HTTPProtocolHTTP1})
	require.NoError(t, err)
	shardedClient, err := GetHttpClientWithProxySettings(proxyURL, dto.ChannelSettings{HTTP2ConnectionShards: 2})
	require.NoError(t, err)
	insecureClient, err := GetHttpClientWithProxySettings(proxyURL, dto.ChannelSettings{TLSInsecureSkipVerify: true})
	require.NoError(t, err)

	InvalidateProxyClient(proxyURL)

	afterDefault, err := GetHttpClientWithProxy(proxyURL)
	require.NoError(t, err)
	afterHTTP1, err := GetHttpClientWithProxySettings(proxyURL, dto.ChannelSettings{HTTPProtocol: dto.HTTPProtocolHTTP1})
	require.NoError(t, err)
	afterSharded, err := GetHttpClientWithProxySettings(proxyURL, dto.ChannelSettings{HTTP2ConnectionShards: 2})
	require.NoError(t, err)
	afterInsecure, err := GetHttpClientWithProxySettings(proxyURL, dto.ChannelSettings{TLSInsecureSkipVerify: true})
	require.NoError(t, err)

	assert.NotSame(t, defaultClient, afterDefault)
	assert.NotSame(t, http1Client, afterHTTP1)
	assert.NotSame(t, shardedClient, afterSharded)
	assert.NotSame(t, insecureClient, afterInsecure)
}

func TestResetProxyClientCacheKeepsDefaultPointerAndRecreatesVariants(t *testing.T) {
	defaultClient := initDefaultHTTPClientFixture(t)

	http1Client, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{HTTPProtocol: dto.HTTPProtocolHTTP1})
	require.NoError(t, err)
	shardedClient, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{HTTP2ConnectionShards: 3})
	require.NoError(t, err)
	proxyClient, err := GetHttpClientWithProxy("http://reset-proxy.example:8080")
	require.NoError(t, err)
	insecureClient, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{TLSInsecureSkipVerify: true})
	require.NoError(t, err)

	ResetProxyClientCache()

	assert.Same(t, defaultClient, GetHttpClient(), "default httpClient pointer must stay stable across reset")
	aware, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{})
	require.NoError(t, err)
	assert.Same(t, defaultClient, aware)

	afterHTTP1, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{HTTPProtocol: dto.HTTPProtocolHTTP1})
	require.NoError(t, err)
	afterSharded, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{HTTP2ConnectionShards: 3})
	require.NoError(t, err)
	afterProxy, err := GetHttpClientWithProxy("http://reset-proxy.example:8080")
	require.NoError(t, err)
	afterInsecure, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{TLSInsecureSkipVerify: true})
	require.NoError(t, err)
	assert.NotSame(t, http1Client, afterHTTP1)
	assert.NotSame(t, shardedClient, afterSharded)
	assert.NotSame(t, proxyClient, afterProxy)
	assert.NotSame(t, insecureClient, afterInsecure)
}

func TestResetProxyClientCacheClosesDefaultIdlePool(t *testing.T) {
	defaultClient := initDefaultHTTPClientFixture(t)
	tracker := &closeCountingRoundTripper{}
	previousTransport := defaultClient.Transport
	defaultClient.Transport = tracker
	t.Cleanup(func() {
		defaultClient.Transport = previousTransport
	})

	ResetProxyClientCache()

	assert.Same(t, defaultClient, GetHttpClient())
	assert.GreaterOrEqual(t, tracker.closes.Load(), int32(1), "reset must close idle connections on the stable default client")
	aware, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{})
	require.NoError(t, err)
	assert.Same(t, defaultClient, aware)
}

func TestResetProxyClientCacheConcurrentWithGetHttpClient(t *testing.T) {
	initDefaultHTTPClientFixture(t)

	const workers = 64
	var wg sync.WaitGroup
	wg.Add(workers * 2)
	for range workers {
		go func() {
			defer wg.Done()
			_ = GetHttpClient()
		}()
		go func() {
			defer wg.Done()
			ResetProxyClientCache()
		}()
	}
	wg.Wait()
	assert.NotNil(t, GetHttpClient())
	aware, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{})
	require.NoError(t, err)
	assert.Same(t, GetHttpClient(), aware)
}

func TestCloseIdleConnectionsRedialsHTTP2(t *testing.T) {
	withRelayHTTPTransportSettings(t)

	var mu sync.Mutex
	addrs := make([]string, 0, 2)

	server := startHTTP2TLSServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		addrs = append(addrs, r.RemoteAddr)
		mu.Unlock()
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok"))
	}))

	client := newHTTPClientWithPolicyAndTLS(defaultHTTPTransportPolicy(), testTLSClientConfig(t, server))
	resp, err := client.Get(server.URL)
	require.NoError(t, err)
	drainClose(t, resp)

	client.CloseIdleConnections()

	resp, err = client.Get(server.URL)
	require.NoError(t, err)
	drainClose(t, resp)

	mu.Lock()
	defer mu.Unlock()
	require.Len(t, addrs, 2)
	assert.NotEqual(t, addrs[0], addrs[1], "after CloseIdleConnections the next request must redial")
}

func TestNormalizeHTTPTransportPolicyClampsWithoutPanic(t *testing.T) {
	assert.Equal(t, defaultHTTPTransportPolicy(), NormalizeHTTPTransportPolicy(dto.ChannelSettings{}))
	assert.Equal(t, HTTPTransportPolicy{Protocol: dto.HTTPProtocolAuto, Shards: 1}, NormalizeHTTPTransportPolicy(dto.ChannelSettings{HTTPProtocol: "AUTO"}))
	assert.Equal(t, HTTPTransportPolicy{Protocol: dto.HTTPProtocolHTTP1, Shards: 1}, NormalizeHTTPTransportPolicy(dto.ChannelSettings{HTTPProtocol: "HTTP1", HTTP2ConnectionShards: 8}))
	assert.Equal(t, HTTPTransportPolicy{Protocol: dto.HTTPProtocolAuto, Shards: 1}, NormalizeHTTPTransportPolicy(dto.ChannelSettings{HTTPProtocol: "http3"}))
	assert.Equal(t, HTTPTransportPolicy{Protocol: dto.HTTPProtocolAuto, Shards: 1}, NormalizeHTTPTransportPolicy(dto.ChannelSettings{HTTP2ConnectionShards: -3}))
	assert.Equal(t, HTTPTransportPolicy{Protocol: dto.HTTPProtocolAuto, Shards: 8}, NormalizeHTTPTransportPolicy(dto.ChannelSettings{HTTP2ConnectionShards: 99}))
}

// serveHTTP2ProbePeer can withhold PING acknowledgements without sending FIN or
// RST, reproducing an upstream connection silently lost by an intermediary.
func serveHTTP2ProbePeer(conn net.Conn, acknowledge bool, pings *atomic.Int32) error {
	defer conn.Close()
	if err := conn.SetDeadline(time.Now().Add(10 * time.Second)); err != nil {
		return err
	}
	preface := make([]byte, len(http2.ClientPreface))
	if _, err := io.ReadFull(conn, preface); err != nil {
		return err
	}
	if string(preface) != http2.ClientPreface {
		return fmt.Errorf("unexpected HTTP/2 preface: %q", preface)
	}
	framer := http2.NewFramer(conn, conn)
	framer.ReadMetaHeaders = hpack.NewDecoder(4096, nil)
	if err := framer.WriteSettings(); err != nil {
		return err
	}
	var streamID uint32
	for {
		frame, err := framer.ReadFrame()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
		switch frame := frame.(type) {
		case *http2.SettingsFrame:
			if !frame.IsAck() {
				if err := framer.WriteSettingsAck(); err != nil {
					return err
				}
			}
		case *http2.MetaHeadersFrame:
			streamID = frame.StreamID
			var block bytes.Buffer
			if err := hpack.NewEncoder(&block).WriteField(hpack.HeaderField{Name: ":status", Value: "200"}); err != nil {
				return err
			}
			if err := framer.WriteHeaders(http2.HeadersFrameParam{StreamID: streamID, BlockFragment: block.Bytes(), EndHeaders: true}); err != nil {
				return err
			}
			if err := framer.WriteData(streamID, false, []byte("start")); err != nil {
				return err
			}
		case *http2.PingFrame:
			if frame.IsAck() {
				continue
			}
			count := pings.Add(1)
			if acknowledge {
				if err := framer.WritePing(true, frame.Data); err != nil {
					return err
				}
				if count >= 2 && streamID != 0 {
					if err := framer.WriteData(streamID, true, []byte("end")); err != nil {
						return err
					}
					streamID = 0
				}
			}
		}
	}
}

func TestRelayHTTP2HealthCheck(t *testing.T) {
	for _, tc := range []struct {
		name        string
		idleTimeout int
		acknowledge bool
	}{
		{name: "disabled", idleTimeout: 0},
		{name: "healthy slow response", idleTimeout: 1, acknowledge: true},
		{name: "unresponsive connection", idleTimeout: 1},
	} {
		t.Run(tc.name, func(t *testing.T) {
			withRelayHTTPTransportSettings(t)
			common.RelayHTTP2PingIdleTimeout = tc.idleTimeout

			var connections, pings atomic.Int32
			peerErrors := make(chan error, 4)
			server := httptest.NewUnstartedServer(nil)
			server.EnableHTTP2 = true
			server.Config.TLSNextProto = map[string]func(*http.Server, *tls.Conn, http.Handler){
				"h2": func(_ *http.Server, conn *tls.Conn, _ http.Handler) {
					connection := connections.Add(1)
					acknowledge := tc.acknowledge || connection > 1
					peerErrors <- serveHTTP2ProbePeer(conn, acknowledge, &pings)
				},
			}
			server.StartTLS()
			t.Cleanup(server.Close)
			client := newHTTPClientWithPolicyAndTLS(defaultHTTPTransportPolicy(), testTLSClientConfig(t, server))
			client.Timeout = 7 * time.Second
			if tc.idleTimeout == 0 {
				client.Timeout = 3 * time.Second
			}
			t.Cleanup(client.CloseIdleConnections)

			resp, err := client.Get(server.URL)
			require.NoError(t, err)
			require.Equal(t, 2, resp.ProtoMajor)
			body, err := io.ReadAll(resp.Body)
			require.NoError(t, resp.Body.Close())
			if tc.idleTimeout == 0 {
				require.ErrorIs(t, err, context.DeadlineExceeded)
				assert.Zero(t, pings.Load(), "disabled probes must preserve the existing transport behavior")
			} else if tc.acknowledge {
				require.NoError(t, err, "a slow response must survive acknowledged health checks")
				assert.Equal(t, "startend", string(body))
				assert.GreaterOrEqual(t, pings.Load(), int32(2))
				assert.Equal(t, int32(1), connections.Load())
			} else {
				require.ErrorContains(t, err, "http2: client connection lost")
				assert.Equal(t, "start", string(body), "already delivered content must not be replayed")
				assert.Equal(t, int32(1), connections.Load())
				require.NoError(t, <-peerErrors, "the peer must observe the failed connection close")

				resp, err = client.Get(server.URL)
				require.NoError(t, err)
				body, err = io.ReadAll(resp.Body)
				require.NoError(t, err)
				require.NoError(t, resp.Body.Close())
				assert.Equal(t, "startend", string(body))
				assert.Equal(t, int32(2), connections.Load(), "the next request must establish a fresh connection")
			}
		})
	}
}