From 3fd8ee76baf4757f448e6a43e3a8d1aab6e271b2 Mon Sep 17 00:00:00 2001 From: Maja <27113+ravicious@users.noreply.github.com> Date: Mon, 11 May 2026 12:38:30 +0200 Subject: [PATCH] Extract gRPC bidi proxy to separate package and fix edge cases (#66143) * Extract gRPC proxy func to separate package * Use debug logs for err on `stream.Recv()` * Add regression test for treating server response as authoritative The current buggy code path: `forwardServerToClient` returns nil on server trailer, main waits on `<-errCh`, `forwardClientToServer` is parked in `client.Recv()` until either the connection is torn down (hang) or the client happens to send another message whose `server.Send` fails on the half-closed upstream (reshape: server's nil becomes a proxy-side error). * Treat server response as authoritative & propagate client-side errors * Migrate tests to use a custom gRPC service * Test if errs from client.Recv and server.Send are surfaced * Fix flaky `TestProxyBidiStream_ReturnsEOFWhenServerReturnsEarly` * Don't set `clientErrCh` to `nil` * Use existing gRPC service instead of special one for tests * Forward header and trailer to client; use bufconn in tests Co-authored-by: Tiago Silva --------- Co-authored-by: Tiago Silva --- lib/secretsscanner/proxy/proxy.go | 76 +---- lib/secretsscanner/proxy/proxy_test.go | 115 +------ lib/utils/grpc/proxy.go | 165 ++++++++++ lib/utils/grpc/proxy_test.go | 438 +++++++++++++++++++++++++ 4 files changed, 612 insertions(+), 182 deletions(-) create mode 100644 lib/utils/grpc/proxy.go create mode 100644 lib/utils/grpc/proxy_test.go 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) +}