mirror of
https://github.com/gravitational/teleport.git
synced 2026-09-21 14:12:20 +08:00
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 <tiago.silva@goteleport.com> --------- Co-authored-by: Tiago Silva <tiago.silva@goteleport.com>
This commit is contained in:
Maja
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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 <http://www.gnu.org/licenses/>.
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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 <http://www.gnu.org/licenses/>.
|
||||
|
||||
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)
|
||||
}
|
||||
Reference in New Issue
Block a user