diff --git a/internal/runtime/executor/helps/utls_client.go b/internal/runtime/executor/helps/utls_client.go index d0e23875..95083201 100644 --- a/internal/runtime/executor/helps/utls_client.go +++ b/internal/runtime/executor/helps/utls_client.go @@ -2,7 +2,9 @@ package helps import ( "context" + "errors" "fmt" + "io" "net" "net/http" "strings" @@ -21,12 +23,35 @@ import ( ) // utlsRoundTripper implements http.RoundTripper using a Chrome fingerprint for -// providers that require a browser-like TLS and HTTP/2 transport. +// providers that require a browser-like TLS and HTTP/2 transport. Each request +// gets a dedicated connection that is closed with the response body. type utlsRoundTripper struct { - mu sync.Mutex - connections map[string]*http2.ClientConn - pending map[string]*sync.Cond - dialer proxy.Dialer + dialer proxy.Dialer +} + +type closeConnectionBody struct { + io.ReadCloser + closeConnection func() error + once sync.Once + err error +} + +func (b *closeConnectionBody) Close() error { + if b == nil { + return nil + } + b.once.Do(func() { + var errConnection error + if b.closeConnection != nil { + errConnection = b.closeConnection() + } + var errBody error + if b.ReadCloser != nil { + errBody = b.ReadCloser.Close() + } + b.err = errors.Join(errBody, errConnection) + }) + return b.err } func newUtlsRoundTripper(proxyURL string) *utlsRoundTripper { @@ -39,68 +64,39 @@ func newUtlsRoundTripper(proxyURL string) *utlsRoundTripper { dialer = proxyDialer } } - return &utlsRoundTripper{ - connections: make(map[string]*http2.ClientConn), - pending: make(map[string]*sync.Cond), - dialer: dialer, - } + return &utlsRoundTripper{dialer: dialer} } -func (t *utlsRoundTripper) getOrCreateConnection(host, addr string) (*http2.ClientConn, error) { - t.mu.Lock() - - if h2Conn, ok := t.connections[host]; ok && h2Conn.CanTakeNewRequest() { - t.mu.Unlock() - return h2Conn, nil - } - - if cond, ok := t.pending[host]; ok { - cond.Wait() - if h2Conn, ok := t.connections[host]; ok && h2Conn.CanTakeNewRequest() { - t.mu.Unlock() - return h2Conn, nil - } - } - - cond := sync.NewCond(&t.mu) - t.pending[host] = cond - t.mu.Unlock() - - h2Conn, err := t.createConnection(host, addr) - - t.mu.Lock() - defer t.mu.Unlock() - - delete(t.pending, host) - cond.Broadcast() - - if err != nil { - return nil, err +func (t *utlsRoundTripper) createConnection(ctx context.Context, host, addr string) (*http2.ClientConn, error) { + contextDialer, ok := t.dialer.(proxy.ContextDialer) + if !ok { + return nil, fmt.Errorf("utls: dialer does not support context cancellation") } - - t.connections[host] = h2Conn - return h2Conn, nil -} - -func (t *utlsRoundTripper) createConnection(host, addr string) (*http2.ClientConn, error) { - conn, err := t.dialer.Dial("tcp", addr) - if err != nil { - return nil, err + conn, errDial := contextDialer.DialContext(ctx, "tcp", addr) + if errDial != nil { + return nil, fmt.Errorf("utls: dial upstream: %w", errDial) } tlsConfig := &tls.Config{ServerName: host} tlsConn := tls.UClient(conn, tlsConfig, tls.HelloChrome_Auto) - if err := tlsConn.Handshake(); err != nil { - conn.Close() - return nil, err + if errHandshake := tlsConn.HandshakeContext(ctx); errHandshake != nil { + if errors.Is(errHandshake, context.Canceled) || errors.Is(errHandshake, context.DeadlineExceeded) { + return nil, fmt.Errorf("utls: TLS handshake: %w", errHandshake) + } + if errClose := conn.Close(); errClose != nil { + return nil, fmt.Errorf("utls: TLS handshake: %w; close connection: %v", errHandshake, errClose) + } + return nil, fmt.Errorf("utls: TLS handshake: %w", errHandshake) } tr := &http2.Transport{} - h2Conn, err := tr.NewClientConn(tlsConn) - if err != nil { - tlsConn.Close() - return nil, err + h2Conn, errClientConn := tr.NewClientConn(tlsConn) + if errClientConn != nil { + if errClose := tlsConn.Close(); errClose != nil { + return nil, fmt.Errorf("utls: initialize HTTP/2 connection: %w; close TLS connection: %v", errClientConn, errClose) + } + return nil, fmt.Errorf("utls: initialize HTTP/2 connection: %w", errClientConn) } return h2Conn, nil @@ -114,21 +110,31 @@ func (t *utlsRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) } addr := net.JoinHostPort(hostname, port) - h2Conn, err := t.getOrCreateConnection(hostname, addr) + h2Conn, err := t.createConnection(req.Context(), hostname, addr) if err != nil { return nil, err } resp, err := h2Conn.RoundTrip(req) if err != nil { - t.mu.Lock() - if cached, ok := t.connections[hostname]; ok && cached == h2Conn { - delete(t.connections, hostname) + if errClose := h2Conn.Close(); errClose != nil { + log.Debugf("utls: close connection after round trip failure: %v", errClose) } - t.mu.Unlock() return nil, err } - + if resp == nil { + if errClose := h2Conn.Close(); errClose != nil { + log.Debugf("utls: close connection after empty response: %v", errClose) + } + return nil, fmt.Errorf("utls: upstream returned an empty response") + } + if resp.Body == nil { + resp.Body = http.NoBody + } + resp.Body = &closeConnectionBody{ + ReadCloser: resp.Body, + closeConnection: h2Conn.Close, + } return resp, nil } diff --git a/internal/runtime/executor/helps/utls_client_test.go b/internal/runtime/executor/helps/utls_client_test.go index 30cee33f..f4492adc 100644 --- a/internal/runtime/executor/helps/utls_client_test.go +++ b/internal/runtime/executor/helps/utls_client_test.go @@ -15,6 +15,7 @@ import ( "reflect" "strconv" "strings" + "sync/atomic" "testing" "time" @@ -28,6 +29,174 @@ func (f utlsClientRoundTripFunc) RoundTrip(req *http.Request) (*http.Response, e return f(req) } +type trackedReadCloser struct { + io.Reader + closeCount int + closeErr error + onClose func() +} + +func (r *trackedReadCloser) Close() error { + r.closeCount++ + if r.onClose != nil { + r.onClose() + } + return r.closeErr +} + +type contextDialerFunc func(context.Context, string, string) (net.Conn, error) + +func (f contextDialerFunc) Dial(network, addr string) (net.Conn, error) { + return f(context.Background(), network, addr) +} + +func (f contextDialerFunc) DialContext(ctx context.Context, network, addr string) (net.Conn, error) { + return f(ctx, network, addr) +} + +type trackedNetConn struct { + net.Conn + closeCount atomic.Int32 +} + +func (c *trackedNetConn) Close() error { + c.closeCount.Add(1) + return c.Conn.Close() +} + +func TestCloseConnectionBodyClosesConnectionBeforeBodyOnce(t *testing.T) { + bodyErr := errors.New("body close failed") + connectionErr := errors.New("connection close failed") + var closeOrder []string + body := &trackedReadCloser{ + Reader: strings.NewReader("response"), + closeErr: bodyErr, + onClose: func() { + closeOrder = append(closeOrder, "body") + }, + } + connectionCloseCount := 0 + wrapped := &closeConnectionBody{ + ReadCloser: body, + closeConnection: func() error { + connectionCloseCount++ + closeOrder = append(closeOrder, "connection") + return connectionErr + }, + } + + payload, errRead := io.ReadAll(wrapped) + if errRead != nil { + t.Fatal(errRead) + } + if got, want := string(payload), "response"; got != want { + t.Fatalf("response body = %q, want %q", got, want) + } + + errClose := wrapped.Close() + if !errors.Is(errClose, bodyErr) { + t.Fatalf("close error = %v, want body close error", errClose) + } + if !errors.Is(errClose, connectionErr) { + t.Fatalf("close error = %v, want connection close error", errClose) + } + if errCloseAgain := wrapped.Close(); errCloseAgain != errClose { + t.Fatalf("second close error = %v, want %v", errCloseAgain, errClose) + } + if body.closeCount != 1 { + t.Fatalf("body close count = %d, want 1", body.closeCount) + } + if connectionCloseCount != 1 { + t.Fatalf("connection close count = %d, want 1", connectionCloseCount) + } + if want := []string{"connection", "body"}; !reflect.DeepEqual(closeOrder, want) { + t.Fatalf("close order = %v, want %v", closeOrder, want) + } +} + +func TestUtlsRoundTripperDialUsesRequestContext(t *testing.T) { + dialStarted := make(chan struct{}) + roundTripper := &utlsRoundTripper{dialer: contextDialerFunc(func(ctx context.Context, _, _ string) (net.Conn, error) { + close(dialStarted) + <-ctx.Done() + return nil, ctx.Err() + })} + ctx, cancel := context.WithCancel(t.Context()) + req, errRequest := http.NewRequestWithContext(ctx, http.MethodGet, "https://chatgpt.com/backend-api/codex/responses", nil) + if errRequest != nil { + t.Fatal(errRequest) + } + roundTripDone := make(chan error, 1) + go func() { + resp, errRoundTrip := roundTripper.RoundTrip(req) + if resp != nil && resp.Body != nil { + errRoundTrip = errors.Join(errRoundTrip, resp.Body.Close()) + } + roundTripDone <- errRoundTrip + }() + + select { + case <-dialStarted: + case <-time.After(time.Second): + t.Fatal("dial did not start") + } + cancel() + select { + case errRoundTrip := <-roundTripDone: + if !errors.Is(errRoundTrip, context.Canceled) { + t.Fatalf("RoundTrip error = %v, want context canceled", errRoundTrip) + } + case <-time.After(time.Second): + t.Fatal("RoundTrip did not stop after context cancellation") + } +} + +func TestUtlsRoundTripperHandshakeUsesRequestContext(t *testing.T) { + clientConn, serverConn := net.Pipe() + t.Cleanup(func() { + if errClose := clientConn.Close(); errClose != nil && !errors.Is(errClose, net.ErrClosed) && !errors.Is(errClose, io.ErrClosedPipe) { + t.Errorf("close client connection: %v", errClose) + } + if errClose := serverConn.Close(); errClose != nil && !errors.Is(errClose, net.ErrClosed) && !errors.Is(errClose, io.ErrClosedPipe) { + t.Errorf("close server connection: %v", errClose) + } + }) + + trackedConn := &trackedNetConn{Conn: clientConn} + dialDone := make(chan struct{}) + roundTripper := &utlsRoundTripper{dialer: contextDialerFunc(func(context.Context, string, string) (net.Conn, error) { + close(dialDone) + return trackedConn, nil + })} + ctx, cancel := context.WithCancel(t.Context()) + connectionDone := make(chan error, 1) + go func() { + h2Conn, errConnect := roundTripper.createConnection(ctx, "chatgpt.com", "chatgpt.com:443") + if h2Conn != nil { + errConnect = errors.Join(errConnect, h2Conn.Close()) + } + connectionDone <- errConnect + }() + + select { + case <-dialDone: + case <-time.After(time.Second): + t.Fatal("dial did not complete") + } + cancel() + select { + case errConnect := <-connectionDone: + if !errors.Is(errConnect, context.Canceled) { + t.Fatalf("createConnection error = %v, want context canceled", errConnect) + } + case <-time.After(time.Second): + t.Fatal("TLS handshake did not stop after context cancellation") + } + if got := trackedConn.closeCount.Load(); got != 1 { + t.Fatalf("connection close count = %d, want 1", got) + } +} + type claudeCodeTLSFingerprintFixture struct { ClientHelloLength int JA3 string