diff --git a/lib/secretsscanner/proxy/proxy.go b/lib/secretsscanner/proxy/proxy.go
index 49160335643..cb6bce0fb13 100644
--- a/lib/secretsscanner/proxy/proxy.go
+++ b/lib/secretsscanner/proxy/proxy.go
@@ -22,13 +22,12 @@ package proxy
import (
"context"
- "errors"
- "io"
"log/slog"
"github.com/gravitational/trace"
accessgraphsecretsv1pb "github.com/gravitational/teleport/api/gen/proto/go/teleport/accessgraph/v1"
+ grpcutils "github.com/gravitational/teleport/lib/utils/grpc"
)
// AuthClient is a subset of the full Auth API that must be connected
@@ -71,72 +70,9 @@ type Service struct {
// ReportSecrets proxies the ReportSecrets method from the proxy to the Auth's secret service.
func (s *Service) ReportSecrets(client accessgraphsecretsv1pb.SecretsScannerService_ReportSecretsServer) error {
- ctx, cancel := context.WithCancel(client.Context())
- defer cancel()
-
- upstream, err := s.authClient.AccessGraphSecretsScannerClient().ReportSecrets(ctx)
- if err != nil {
- return trace.Wrap(err)
- }
-
- errCh := make(chan error, 1)
- go func() {
- errCh <- trace.Wrap(s.forwardClientToServer(ctx, cancel, client, upstream))
- }()
-
- err = s.forwardServerToClient(ctx, client, upstream)
- if err != nil {
- // Return immediately so gRPC closes the stream, which unblocks client.Recv()
- // in the forwardClientToServer goroutine. The buffered errCh prevents a leak.
- return trace.Wrap(err)
- }
- return trace.Wrap(<-errCh)
-}
-
-func (s *Service) forwardClientToServer(ctx context.Context, cancel context.CancelFunc,
- client accessgraphsecretsv1pb.SecretsScannerService_ReportSecretsServer,
- server accessgraphsecretsv1pb.SecretsScannerService_ReportSecretsClient) (err error) {
- defer func() {
- // CloseSend always returns nil error.
- _ = server.CloseSend()
- }()
- for {
- req, err := client.Recv()
- if errors.Is(err, io.EOF) {
- // The client closed the send direction and won't send more messages.
- // Close the send direction of the server stream by returning and _do not_
- // cancel the context so that the client can receive any messages that the
- // server sends after getting io.EOF from the client.
- return nil
- }
- if err != nil {
- s.log.WarnContext(ctx, "Failed to receive from client stream", "error", err)
- cancel()
- return trace.Wrap(err)
- }
- if err := server.Send(req); err != nil {
- s.log.WarnContext(ctx, "Failed to send to upstream stream", "error", err)
- cancel()
- return trace.Wrap(err)
- }
- }
-}
-
-func (s *Service) forwardServerToClient(ctx context.Context,
- client accessgraphsecretsv1pb.SecretsScannerService_ReportSecretsServer,
- server accessgraphsecretsv1pb.SecretsScannerService_ReportSecretsClient) (err error) {
- for {
- out, err := server.Recv()
- if errors.Is(err, io.EOF) {
- return nil
- }
- if err != nil {
- s.log.WarnContext(ctx, "Failed to receive from upstream stream", "error", err)
- return trace.Wrap(err)
- }
- if err := client.Send(out); err != nil {
- s.log.WarnContext(ctx, "Failed to send to client stream", "error", err)
- return trace.Wrap(err)
- }
- }
+ err := grpcutils.ProxyBidiStream(s.log, client, func(ctx context.Context) (accessgraphsecretsv1pb.SecretsScannerService_ReportSecretsClient, error) {
+ server, err := s.authClient.AccessGraphSecretsScannerClient().ReportSecrets(ctx)
+ return server, trace.Wrap(err)
+ })
+ return trace.Wrap(err)
}
diff --git a/lib/secretsscanner/proxy/proxy_test.go b/lib/secretsscanner/proxy/proxy_test.go
index b814edb71ae..9fa105ed02e 100644
--- a/lib/secretsscanner/proxy/proxy_test.go
+++ b/lib/secretsscanner/proxy/proxy_test.go
@@ -19,13 +19,11 @@
package proxy
import (
- "context"
"crypto/tls"
"errors"
"io"
"net"
"testing"
- "time"
"github.com/gravitational/trace"
"github.com/stretchr/testify/assert"
@@ -45,7 +43,7 @@ func TestProxy(t *testing.T) {
// Disable the TLS routing connection upgrade
t.Setenv(defaults.TLSRoutingConnUpgradeEnvVar, "false")
- _, authClient := newFakefakeSecretsScannerSvc(t)
+ authClient := newFakefakeSecretsScannerSvc(t)
lis, err := net.Listen("tcp", "localhost:0")
require.NoError(t, err)
@@ -105,106 +103,7 @@ func TestProxy(t *testing.T) {
require.ErrorIs(t, err, io.EOF)
}
-func TestProxy_HandlesServerReturningErr(t *testing.T) {
- // Disable the TLS routing connection upgrade
- t.Setenv(defaults.TLSRoutingConnUpgradeEnvVar, "false")
-
- _, authClient := newFakefakeSecretsScannerSvc(t)
-
- lis, err := net.Listen("tcp", "localhost:0")
- require.NoError(t, err)
-
- newProxyService(t, lis, authClient)
- // Add a short timeout so if the proxy hangs (as it did before introducing this regression test),
- // the test doesn't wait for a whole minute to fail.
- ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
- defer cancel()
-
- client, err := secretscannerclient.NewSecretsScannerServiceClient(ctx, secretscannerclient.ClientConfig{
- ProxyServer: lis.Addr().String(),
- Insecure: true,
- })
- require.NoError(t, err)
-
- stream, err := client.ReportSecrets(ctx)
- require.NoError(t, err)
-
- // Send incomplete message which should cause the server to return an error.
- err = stream.Send(&accessgraphsecretsv1pb.ReportSecretsRequest{})
- require.NoError(t, err)
- _, err = stream.Recv()
- require.ErrorContains(t, err, "missing device init")
-}
-
-// TestProxy_PropagatesUpstreamErrorAfterClientEOF asserts that a terminal
-// error produced by the upstream SecretsScannerService *after* the client has
-// half-closed (CloseSend) is still propagated through the proxy to the client.
-//
-// This exercises the handler path where forwardClientToServer returns first
-// (normal CloseSend) and forwardServerToClient is the one that ends up carrying
-// Auth's terminal status. A handler that treats forwardClientToServer as
-// authoritative will finish and the client will see io.EOF instead of the real
-// error, masking real upstream failures.
-func TestProxy_PropagatesUpstreamErrorAfterClientEOF(t *testing.T) {
- t.Setenv(defaults.TLSRoutingConnUpgradeEnvVar, "false")
-
- service, authClient := newFakefakeSecretsScannerSvc(t)
- service.postClientEOFErr = trace.AccessDenied("post-EOF validation failed")
-
- lis, err := net.Listen("tcp", "localhost:0")
- require.NoError(t, err)
-
- newProxyService(t, lis, authClient)
- ctx := t.Context()
-
- client, err := secretscannerclient.NewSecretsScannerServiceClient(ctx, secretscannerclient.ClientConfig{
- ProxyServer: lis.Addr().String(),
- Insecure: true,
- })
- require.NoError(t, err)
-
- stream, err := client.ReportSecrets(ctx)
- require.NoError(t, err)
-
- // Full handshake so Auth reaches the final in.Recv() that ends in EOF.
- err = stream.Send(&accessgraphsecretsv1pb.ReportSecretsRequest{
- Payload: &accessgraphsecretsv1pb.ReportSecretsRequest_DeviceAssertion{
- DeviceAssertion: &devicepb.AssertDeviceRequest{
- Payload: &devicepb.AssertDeviceRequest_Init{
- Init: &devicepb.AssertDeviceInit{},
- },
- },
- },
- })
- require.NoError(t, err)
-
- _, err = stream.Recv()
- require.NoError(t, err)
-
- err = stream.Send(&accessgraphsecretsv1pb.ReportSecretsRequest{
- Payload: &accessgraphsecretsv1pb.ReportSecretsRequest_DeviceAssertion{
- DeviceAssertion: &devicepb.AssertDeviceRequest{
- Payload: &devicepb.AssertDeviceRequest_ChallengeResponse{
- ChallengeResponse: &devicepb.AuthenticateDeviceChallengeResponse{Signature: []byte("response")},
- },
- },
- },
- })
- require.NoError(t, err)
-
- _, err = stream.Recv()
- require.NoError(t, err)
-
- err = stream.CloseSend()
- require.NoError(t, err)
-
- // The client must see the upstream error, not a clean io.EOF.
- _, recvErr := stream.Recv()
- require.NotErrorIs(t, recvErr, io.EOF, "client saw clean EOF; upstream error was swallowed")
- require.ErrorContains(t, recvErr, "post-EOF validation failed")
-}
-
-func newFakefakeSecretsScannerSvc(t *testing.T) (*fakeSecretsScannerSvc, *fakeSecretsClient) {
+func newFakefakeSecretsScannerSvc(t *testing.T) *fakeSecretsClient {
lis, err := net.Listen("tcp", "localhost:0")
require.NoError(t, err)
@@ -220,7 +119,7 @@ func newFakefakeSecretsScannerSvc(t *testing.T) (*fakeSecretsScannerSvc, *fakeSe
client, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
require.NoError(t, err)
- return service, &fakeSecretsClient{
+ return &fakeSecretsClient{
SecretsScannerServiceClient: accessgraphsecretsv1pb.NewSecretsScannerServiceClient(client),
}
@@ -236,11 +135,6 @@ func (s *fakeSecretsClient) AccessGraphSecretsScannerClient() accessgraphsecrets
type fakeSecretsScannerSvc struct {
accessgraphsecretsv1pb.UnimplementedSecretsScannerServiceServer
-
- // postClientEOFErr, if non-nil, is returned by ReportSecrets after it
- // receives EOF from the client, modeling Auth producing a terminal error
- // during post-upload processing (after the client has already half-closed).
- postClientEOFErr error
}
func (f *fakeSecretsScannerSvc) ReportSecrets(in accessgraphsecretsv1pb.SecretsScannerService_ReportSecretsServer) error {
@@ -289,9 +183,6 @@ func (f *fakeSecretsScannerSvc) ReportSecrets(in accessgraphsecretsv1pb.SecretsS
_, err = in.Recv()
if errors.Is(err, io.EOF) {
- if f.postClientEOFErr != nil {
- return f.postClientEOFErr
- }
return nil
}
return trace.BadParameter("unexpected message")
diff --git a/lib/utils/grpc/proxy.go b/lib/utils/grpc/proxy.go
new file mode 100644
index 00000000000..240d930d46c
--- /dev/null
+++ b/lib/utils/grpc/proxy.go
@@ -0,0 +1,165 @@
+// Teleport
+// Copyright (C) 2026 Gravitational, Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package grpc
+
+import (
+ "context"
+ "errors"
+ "io"
+ "log/slog"
+
+ "github.com/gravitational/trace"
+ "google.golang.org/grpc"
+ "google.golang.org/grpc/metadata"
+)
+
+// ProxyBidiStream proxies a bidi-streaming RPC. It forwards messages from
+// client to server and responses back to client until the server stream
+// finishes (cleanly or with error) or the client stream errors.
+//
+// getServer is called with a context derived from client.Context() and must
+// return the server client stream. Canceling that context tears down both
+// directions, so callers should pass it directly to the server dial call.
+//
+// If the server returns early, any still-in-flight messages the client sent are
+// dropped by the proxy. Also, the client can half-close and still receive
+// messages from the server. Those two behaviors match what the client would see
+// talking to the server directly.
+//
+// During the brief window between the server ending the stream and the proxy's
+// handler returning, client Send calls return nil rather than io.EOF and are
+// dropped. A client that interleaves Send with Recv is unaffected because the
+// next Recv carries the terminal status.
+func ProxyBidiStream[Req, Resp any](log *slog.Logger, client grpc.BidiStreamingServer[Req, Resp],
+ getServer func(context.Context) (grpc.BidiStreamingClient[Req, Resp], error),
+) error {
+ ctx, cancel := context.WithCancel(client.Context())
+ defer cancel()
+
+ if md, ok := metadata.FromIncomingContext(client.Context()); ok {
+ ctx = metadata.NewOutgoingContext(ctx, md)
+ }
+
+ server, err := getServer(ctx)
+ if err != nil {
+ return trace.Wrap(err, "establishing server stream")
+ }
+
+ clientErrCh := make(chan error, 1)
+ serverErrCh := make(chan error, 1)
+
+ go func() { clientErrCh <- forwardClientToServer(ctx, log, client, server) }()
+ go func() { serverErrCh <- forwardServerToClient(ctx, log, client, server) }()
+
+ for {
+ select {
+ case err := <-serverErrCh:
+ // The server stream is authoritative for the RPC's terminal status.
+ // Whatever it returns is what the client should see.
+ return trace.Wrap(err)
+ case err := <-clientErrCh:
+ if err != nil {
+ // Something went wrong on the client side (client.Recv failure, or a
+ // locally-generated server.Send failure). Cancel the server stream and
+ // surface the client error — it's more specific than whatever Canceled
+ // serverErrCh is about to produce.
+ cancel()
+ return trace.Wrap(err)
+ }
+ // forwardClientToServer finished cleanly: the client half-closed or the
+ // server stream is already terminal (Send returned io.EOF). In either
+ // case, keep waiting on the server stream to deliver its terminal status.
+ }
+ }
+}
+
+func forwardClientToServer[Req, Resp any](ctx context.Context, log *slog.Logger,
+ client grpc.BidiStreamingServer[Req, Resp],
+ server grpc.BidiStreamingClient[Req, Resp],
+) error {
+ defer func() {
+ // CloseSend always returns nil error.
+ _ = server.CloseSend()
+ }()
+
+ for {
+ req, err := client.Recv()
+ if errors.Is(err, io.EOF) {
+ // The client half-closed its send side and won't send more messages.
+ // Returning here triggers the deferred CloseSend on the server stream.
+ // The caller keeps waiting on the server stream for its terminal status.
+ return nil
+ }
+ if err != nil {
+ // Debug log because it's impossible to distinguish between transport and
+ // application errors.
+ //
+ // If both proxying functions were to warn on err from Recv, each
+ // application-level err from the server would result in two log lines.
+ // First with the server error and the second with a context canceled for
+ // the client stream.
+ log.DebugContext(ctx, "Failed to receive from client stream", "error", err)
+ return trace.Wrap(err)
+ }
+
+ err = server.Send(req)
+ if errors.Is(err, io.EOF) {
+ // io.EOF means the server ended the stream and the real status is
+ // discoverable via Recv. forwardServerToClient is running that Recv.
+ // Let it surface the terminal status.
+ // We can't forward this io.EOF to the client because the client already
+ // got nil from its Send when we got its message through client.Recv.
+ return nil
+ }
+ if err != nil {
+ log.WarnContext(ctx, "Failed to send to server stream", "error", err)
+ return trace.Wrap(err)
+ }
+ }
+}
+
+func forwardServerToClient[Req, Resp any](ctx context.Context, log *slog.Logger,
+ client grpc.BidiStreamingServer[Req, Resp],
+ server grpc.BidiStreamingClient[Req, Resp],
+) error {
+ defer func() { client.SetTrailer(server.Trailer()) }()
+
+ if md, err := server.Header(); err != nil {
+ log.DebugContext(ctx, "Failed to receive headers from server stream", "error", err)
+ } else if len(md) > 0 {
+ if sendErr := client.SendHeader(md); sendErr != nil {
+ log.WarnContext(ctx, "Failed to send headers to client", "error", sendErr)
+ }
+ }
+
+ for {
+ out, err := server.Recv()
+ if errors.Is(err, io.EOF) {
+ return nil
+ }
+ if err != nil {
+ // Debug log because it's impossible to distinguish between transport and
+ // application errors.
+ log.DebugContext(ctx, "Failed to receive from server stream", "error", err)
+ return trace.Wrap(err)
+ }
+ if err := client.Send(out); err != nil {
+ log.WarnContext(ctx, "Failed to send to client stream", "error", err)
+ return trace.Wrap(err)
+ }
+ }
+}
diff --git a/lib/utils/grpc/proxy_test.go b/lib/utils/grpc/proxy_test.go
new file mode 100644
index 00000000000..971a4fd0a5c
--- /dev/null
+++ b/lib/utils/grpc/proxy_test.go
@@ -0,0 +1,438 @@
+// Teleport
+// Copyright (C) 2026 Gravitational, Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package grpc_test
+
+import (
+ "context"
+ "errors"
+ "io"
+ "net"
+ "strings"
+ "testing"
+ "time"
+
+ "github.com/gravitational/trace"
+ "github.com/stretchr/testify/require"
+ "google.golang.org/grpc"
+ "google.golang.org/grpc/credentials/insecure"
+ "google.golang.org/grpc/metadata"
+ "google.golang.org/grpc/test/bufconn"
+
+ teletermv1 "github.com/gravitational/teleport/gen/proto/go/teleport/lib/teleterm/v1"
+ grpcutils "github.com/gravitational/teleport/lib/utils/grpc"
+ "github.com/gravitational/teleport/lib/utils/log/logtest"
+)
+
+func TestMain(m *testing.M) {
+ logtest.InitLogger(testing.Verbose)
+ m.Run()
+}
+
+// TestProxyBidiStream creates two gRPC services: one acting as a server and one
+// as a proxy. The proxy uses [grpcutils.ProxyBidiStream] to proxy messages from
+// the client and the server.
+//
+// Both services implement [teletermv1.TerminalServiceServer]. The server uses
+// [fakeServerSvc] as its implementation, whereas the proxy uses [proxyService].
+//
+// The other tests in this file use the same setup. TestProxyBidiStream tests
+// the happy path.
+func TestProxyBidiStream(t *testing.T) {
+ t.Parallel()
+ _, fakeServerSvcClient := newFakeServerSvc(t)
+
+ lis := bufconn.Listen(1024)
+ newProxyService(t, lis, fakeServerSvcClient)
+ ctx := t.Context()
+
+ client := newProxyServiceClient(t, lis)
+ stream, err := client.ConnectToDesktop(ctx)
+ require.NoError(t, err)
+
+ // Send a message.
+ err = stream.Send(&teletermv1.ConnectToDesktopRequest{Data: []byte("hello")})
+ require.NoError(t, err)
+
+ // Receive the server's response.
+ msg, err := stream.Recv()
+ require.NoError(t, err)
+ require.Equal(t, []byte("ack"), msg.GetData())
+
+ // Half-close and wait for the server to terminate the stream cleanly.
+ err = stream.CloseSend()
+ require.NoError(t, err)
+
+ _, err = stream.Recv()
+ require.ErrorIs(t, err, io.EOF)
+}
+
+// TestProxyBidiStream_HandlesServerReturningErr covers the case where the
+// server errors on its first Recv. Before this regression test the proxy
+// handler could deadlock instead of propagating the error.
+func TestProxyBidiStream_HandlesServerReturningErr(t *testing.T) {
+ t.Parallel()
+ _, fakeServerSvcClient := newFakeServerSvc(t)
+
+ lis := bufconn.Listen(1024)
+ newProxyService(t, lis, fakeServerSvcClient)
+ // Add a short timeout so if the proxy hangs (as it did before introducing
+ // this regression test), the test doesn't wait for a whole minute to fail.
+ ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
+ defer cancel()
+
+ client := newProxyServiceClient(t, lis)
+ stream, err := client.ConnectToDesktop(ctx)
+ require.NoError(t, err)
+
+ // Empty input triggers the fake server to return an error on its first
+ // Recv.
+ err = stream.Send(&teletermv1.ConnectToDesktopRequest{})
+ require.NoError(t, err)
+ _, err = stream.Recv()
+ require.ErrorContains(t, err, "empty data")
+}
+
+// TestProxyBidiStream_PropagatesServerErrorAfterClientEOF asserts that a
+// terminal error produced by the server *after* the client has half-closed
+// (CloseSend) is still propagated through the proxy to the client.
+//
+// This exercises the handler path where forwardClientToServer returns first
+// (normal CloseSend) and forwardServerToClient is the one that ends up carrying
+// server's terminal status. A handler that treats forwardClientToServer as
+// authoritative will finish and the client will see io.EOF instead of the real
+// error, masking real server failures.
+func TestProxyBidiStream_PropagatesServerErrorAfterClientEOF(t *testing.T) {
+ t.Parallel()
+ service, fakeServerSvcClient := newFakeServerSvc(t)
+ service.postClientEOFErr = trace.AccessDenied("post-EOF validation failed")
+
+ lis := bufconn.Listen(1024)
+ newProxyService(t, lis, fakeServerSvcClient)
+ ctx := t.Context()
+
+ client := newProxyServiceClient(t, lis)
+ stream, err := client.ConnectToDesktop(ctx)
+ require.NoError(t, err)
+
+ err = stream.Send(&teletermv1.ConnectToDesktopRequest{Data: []byte("hello")})
+ require.NoError(t, err)
+ _, err = stream.Recv()
+ require.NoError(t, err)
+
+ err = stream.CloseSend()
+ require.NoError(t, err)
+
+ // The client must see the server error, not a clean io.EOF.
+ _, recvErr := stream.Recv()
+ require.NotErrorIs(t, recvErr, io.EOF, "client saw clean EOF; server error was swallowed")
+ require.ErrorContains(t, recvErr, "post-EOF validation failed")
+}
+
+// TestProxyBidiStream_ReturnsEOFWhenServerReturnsEarly asserts that when the
+// server ends its handler cleanly (nil) *before* the client has half-closed,
+// the proxy propagates that as io.EOF to the client rather than hanging or
+// reshaping the server's nil into an error.
+func TestProxyBidiStream_ReturnsEOFWhenServerReturnsEarly(t *testing.T) {
+ t.Parallel()
+ service, fakeServerSvcClient := newFakeServerSvc(t)
+ service.returnAfterFirstResponse = true
+
+ lis := bufconn.Listen(1024)
+ newProxyService(t, lis, fakeServerSvcClient)
+ // Short timeout so a hang surfaces as a test failure rather than waiting
+ // out the default go-test timeout.
+ ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
+ defer cancel()
+
+ client := newProxyServiceClient(t, lis)
+ stream, err := client.ConnectToDesktop(ctx)
+ require.NoError(t, err)
+
+ err = stream.Send(&teletermv1.ConnectToDesktopRequest{Data: []byte("hello")})
+ require.NoError(t, err)
+
+ // Drain the first response the server sent before returning.
+ _, err = stream.Recv()
+ require.NoError(t, err)
+
+ // Client sends the next message it would naturally send, not knowing the
+ // server has already returned. Under the bug this reshapes the server's
+ // clean completion into an error via a failed upstream Send; under the
+ // fix the handler has already returned and the Send is irrelevant.
+ //
+ // At this point, the Send returns either nil if the trailer wasn't propagated
+ // to the client yet or io.EOF if it was, so we skip asserting on err here.
+ _ = stream.Send(&teletermv1.ConnectToDesktopRequest{Data: []byte("more")})
+
+ // Server has returned nil. The client must see clean io.EOF, not a
+ // proxy-reshaped error and not a hang (which would surface as a
+ // DeadlineExceeded from ctx).
+ _, err = stream.Recv()
+ require.ErrorIs(t, err, io.EOF)
+}
+
+// TestProxyBidiStream_SurfacesClientRecvError asserts that when client.Recv
+// on the proxy side fails with a non-EOF error, the proxy returns that
+// specific error instead of the Canceled artifact produced by a naive design
+// that cancels the server stream and then returns whatever server.Recv yields.
+//
+// To trigger this, we set a tiny MaxRecvMsgSize on the proxy's gRPC server and
+// have the client send a message exceeding it. The proxy's client.Recv returns
+// a ResourceExhausted status error; the handler must propagate it.
+func TestProxyBidiStream_SurfacesClientRecvError(t *testing.T) {
+ t.Parallel()
+ _, fakeServerSvcClient := newFakeServerSvc(t)
+
+ lis := bufconn.Listen(1024)
+ newProxyService(t, lis, fakeServerSvcClient, grpc.MaxRecvMsgSize(64))
+
+ client := newProxyServiceClient(t, lis)
+ stream, err := client.ConnectToDesktop(t.Context())
+ require.NoError(t, err)
+
+ // Send a message larger than the proxy's MaxRecvMsgSize.
+ err = stream.Send(&teletermv1.ConnectToDesktopRequest{Data: []byte(strings.Repeat("x", 256))})
+ require.NoError(t, err)
+
+ _, err = stream.Recv()
+ require.ErrorContains(t, err, "larger than max")
+}
+
+// TestProxyBidiStream_SurfacesServerSendError asserts that when server.Send on
+// the proxy side fails with a non-EOF error (locally generated, e.g. the
+// outbound message exceeds MaxCallSendMsgSize on the proxy's upstream
+// connection), the proxy returns that specific error rather than masking it as
+// Canceled.
+//
+// To trigger this, we dial the fake server with a tiny MaxCallSendMsgSize so
+// that the proxy's server.Send fails whenever the client-forwarded message is
+// larger than that limit.
+func TestProxyBidiStream_SurfacesServerSendError(t *testing.T) {
+ t.Parallel()
+ _, fakeServerSvcClient := newFakeServerSvc(t,
+ grpc.WithDefaultCallOptions(grpc.MaxCallSendMsgSize(64)),
+ )
+
+ lis := bufconn.Listen(1024)
+ newProxyService(t, lis, fakeServerSvcClient)
+
+ client := newProxyServiceClient(t, lis)
+ stream, err := client.ConnectToDesktop(t.Context())
+ require.NoError(t, err)
+
+ // The proxy accepts this message (its server-side MaxRecvMsgSize is the
+ // default 4MB), then tries to forward it to the fake server whose upstream
+ // connection caps sends at 64 bytes, triggering a local ResourceExhausted on
+ // the proxy's server.Send.
+ err = stream.Send(&teletermv1.ConnectToDesktopRequest{Data: []byte(strings.Repeat("x", 256))})
+ require.NoError(t, err)
+
+ _, err = stream.Recv()
+ require.ErrorContains(t, err, "larger than max")
+}
+
+// TestProxyBidiStream_ForwardsMetadata asserts that the proxy passes the
+// client's incoming metadata upstream to the server and forwards the server's
+// response headers and trailers back to the client.
+func TestProxyBidiStream_ForwardsMetadata(t *testing.T) {
+ t.Parallel()
+ service, fakeServerSvcClient := newFakeServerSvc(t)
+ service.echoMetadata = true
+
+ lis := bufconn.Listen(1024)
+ newProxyService(t, lis, fakeServerSvcClient)
+
+ // Short timeout so a hang (e.g. Header() never unblocks due to a regression)
+ // surfaces as a test failure rather than waiting out the default go-test timeout.
+ ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
+ defer cancel()
+ // Attach metadata to the outgoing call so the proxy can forward it upstream.
+ ctx = metadata.NewOutgoingContext(ctx, metadata.Pairs("x-test-key", "test-value"))
+
+ client := newProxyServiceClient(t, lis)
+ stream, err := client.ConnectToDesktop(ctx)
+ require.NoError(t, err)
+
+ err = stream.Send(&teletermv1.ConnectToDesktopRequest{Data: []byte("hello")})
+ require.NoError(t, err)
+
+ // Receive the server's first response; by this point the server has already
+ // called SendHeader so headers are available on the client stream.
+ _, err = stream.Recv()
+ require.NoError(t, err)
+
+ headers, err := stream.Header()
+ require.NoError(t, err)
+ require.Equal(t, []string{"test-value"}, headers.Get("x-test-key"),
+ "response headers not forwarded through proxy")
+
+ err = stream.CloseSend()
+ require.NoError(t, err)
+
+ _, err = stream.Recv()
+ require.ErrorIs(t, err, io.EOF)
+
+ trailers := stream.Trailer()
+ require.Equal(t, []string{"test-value"}, trailers.Get("x-test-key"),
+ "response trailers not forwarded through proxy")
+}
+
+func newFakeServerSvc(t *testing.T, clientOpts ...grpc.DialOption) (*fakeServerSvc, teletermv1.TerminalServiceClient) {
+ lis := bufconn.Listen(1024)
+ server := grpc.NewServer()
+ service := &fakeServerSvc{}
+ teletermv1.RegisterTerminalServiceServer(server, service)
+ go func() {
+ err := server.Serve(lis)
+ require.NoError(t, err)
+ }()
+ t.Cleanup(server.GracefulStop)
+
+ opts := append([]grpc.DialOption{
+ grpc.WithTransportCredentials(insecure.NewCredentials()),
+ grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) {
+ return lis.DialContext(ctx)
+ }),
+ }, clientOpts...)
+ client, err := grpc.NewClient("passthrough:///bufconn", opts...)
+ require.NoError(t, err)
+
+ return service, teletermv1.NewTerminalServiceClient(client)
+}
+
+type fakeServerSvc struct {
+ teletermv1.UnimplementedTerminalServiceServer
+
+ // postClientEOFErr, if non-nil, is returned by ConnectToDesktop after it gets
+ // EOF from the client, modeling the server producing a terminal error during
+ // post-upload processing (after the client has already half-closed).
+ postClientEOFErr error
+
+ // returnAfterFirstResponse, if true, makes ConnectToDesktop return nil right
+ // after sending its first response, without waiting for any further client
+ // input or for the client to half-close. It models a server that ends the
+ // stream early while the client is still mid-conversation.
+ returnAfterFirstResponse bool
+
+ // echoMetadata, if true, makes ConnectToDesktop read the incoming metadata
+ // from the stream context and echo it back as both response headers and
+ // trailers. Used to verify the proxy forwards metadata in both directions.
+ echoMetadata bool
+}
+
+// ConnectToDesktop does NOT implement the semantics of the real
+// ConnectToDesktop RPC. The RPC is borrowed purely for its bidi-stream shape so
+// the tests in this file can exercise ProxyBidiStream without introducing a
+// custom test-only proto.
+//
+// Contract used by the tests:
+// - Every request must populate data with a non-empty payload. An empty data
+// triggers a trace.BadParameter return.
+// - Every response carries data = "ack".
+func (f *fakeServerSvc) ConnectToDesktop(stream teletermv1.TerminalService_ConnectToDesktopServer) error {
+ if f.echoMetadata {
+ if md, ok := metadata.FromIncomingContext(stream.Context()); ok {
+ stream.SetTrailer(md)
+ if err := stream.SendHeader(md); err != nil {
+ return trace.Wrap(err)
+ }
+ }
+ }
+ for {
+ req, err := stream.Recv()
+ if errors.Is(err, io.EOF) {
+ if f.postClientEOFErr != nil {
+ return f.postClientEOFErr
+ }
+ return nil
+ }
+ if err != nil {
+ return trace.Wrap(err)
+ }
+ if len(req.GetData()) == 0 {
+ return trace.BadParameter("empty data")
+ }
+ if err := stream.Send(&teletermv1.ConnectToDesktopResponse{Data: []byte("ack")}); err != nil {
+ return trace.Wrap(err)
+ }
+ if f.returnAfterFirstResponse {
+ return nil
+ }
+ }
+}
+
+// newProxyService creates a gRPC server under lis and registers in it a gRPC
+// service that proxies ConnectToDesktop calls to the server using
+// [grpcutils.ProxyBidiStream]. Callers may supply extra grpc.ServerOptions to
+// drive specific fault scenarios.
+func newProxyService(t *testing.T, lis net.Listener, client teletermv1.TerminalServiceClient, opts ...grpc.ServerOption) {
+ t.Helper()
+
+ s := grpc.NewServer(opts...)
+ t.Cleanup(s.GracefulStop)
+
+ proxySvc := &proxyService{
+ serverSvcClient: client,
+ }
+
+ teletermv1.RegisterTerminalServiceServer(s, proxySvc)
+
+ go func() {
+ err := s.Serve(lis)
+ require.NoError(t, err)
+ }()
+}
+
+type proxyService struct {
+ teletermv1.UnimplementedTerminalServiceServer
+
+ serverSvcClient teletermv1.TerminalServiceClient
+}
+
+// ConnectToDesktop forwards every client request (whose data is non-empty by
+// contract) to the upstream server and every response (carrying data = "ack" by
+// contract) back to the client, using ProxyBidiStream.
+//
+// ConnectToDesktop does NOT implement the semantics of the real
+// ConnectToDesktop RPC. See the godoc for [fakeServerSvc.ConnectToDesktop].
+//
+// client goes from a client to the proxy. From that point of view, the proxy
+// is a server for the client.
+// server from getServer goes from the proxy to the server. From that point of
+// view, the proxy is a client of the server.
+func (p *proxyService) ConnectToDesktop(client teletermv1.TerminalService_ConnectToDesktopServer) error {
+ getServer := func(ctx context.Context) (teletermv1.TerminalService_ConnectToDesktopClient, error) {
+ return p.serverSvcClient.ConnectToDesktop(ctx)
+ }
+ err := grpcutils.ProxyBidiStream(logtest.NewLogger(), client, getServer)
+ return trace.Wrap(err)
+}
+
+func newProxyServiceClient(t *testing.T, lis *bufconn.Listener) teletermv1.TerminalServiceClient {
+ t.Helper()
+ clientConn, err := grpc.NewClient(
+ "passthrough:///bufconn",
+ grpc.WithContextDialer(
+ func(ctx context.Context, _ string) (net.Conn, error) {
+ return lis.DialContext(ctx)
+ },
+ ),
+ grpc.WithTransportCredentials(insecure.NewCredentials()),
+ )
+ require.NoError(t, err)
+ return teletermv1.NewTerminalServiceClient(clientConn)
+}