From 34be38dc23e4f642eb8f0aac96472d793d6e0891 Mon Sep 17 00:00:00 2001 From: Dylan Huff Date: Wed, 9 Sep 2026 12:13:34 -0700 Subject: [PATCH 1/2] fix(coderd/agentapi): reject missing RPC messages --- coderd/agentapi/lifecycle.go | 10 ++ coderd/agentapi/metadata.go | 12 +++ coderd/agentapi/validation_test.go | 151 +++++++++++++++++++++++++++++ 3 files changed, 173 insertions(+) create mode 100644 coderd/agentapi/validation_test.go diff --git a/coderd/agentapi/lifecycle.go b/coderd/agentapi/lifecycle.go index 5003a16f04dae..d9b1aa296a162 100644 --- a/coderd/agentapi/lifecycle.go +++ b/coderd/agentapi/lifecycle.go @@ -10,7 +10,9 @@ import ( "github.com/google/uuid" "golang.org/x/mod/semver" "golang.org/x/xerrors" + "google.golang.org/grpc/codes" "google.golang.org/protobuf/types/known/timestamppb" + "storj.io/drpc/drpcerr" "cdr.dev/slog/v3" agentproto "github.com/coder/coder/v2/agent/proto" @@ -45,6 +47,10 @@ func (a *LifecycleAPI) now() time.Time { } func (a *LifecycleAPI) UpdateLifecycle(ctx context.Context, req *agentproto.UpdateLifecycleRequest) (*agentproto.Lifecycle, error) { + if req.GetLifecycle() == nil { + return nil, drpcerr.WithCode(xerrors.New("lifecycle is required"), uint64(codes.InvalidArgument)) + } + workspaceAgent, err := a.AgentFn(ctx) if err != nil { return nil, err @@ -146,6 +152,10 @@ func (a *LifecycleAPI) UpdateLifecycle(ctx context.Context, req *agentproto.Upda } func (a *LifecycleAPI) UpdateStartup(ctx context.Context, req *agentproto.UpdateStartupRequest) (*agentproto.Startup, error) { + if req.GetStartup() == nil { + return nil, drpcerr.WithCode(xerrors.New("startup is required"), uint64(codes.InvalidArgument)) + } + apiVersion, ok := ctx.Value(contextKeyAPIVersion{}).(string) if !ok { return nil, xerrors.Errorf("internal error; api version unspecified") diff --git a/coderd/agentapi/metadata.go b/coderd/agentapi/metadata.go index 12efe362abb02..8baee8eb3f236 100644 --- a/coderd/agentapi/metadata.go +++ b/coderd/agentapi/metadata.go @@ -8,6 +8,8 @@ import ( "github.com/google/uuid" "golang.org/x/xerrors" + "google.golang.org/grpc/codes" + "storj.io/drpc/drpcerr" "cdr.dev/slog/v3" agentproto "github.com/coder/coder/v2/agent/proto" @@ -34,6 +36,16 @@ func (a *MetadataAPI) now() time.Time { } func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.BatchUpdateMetadataRequest) (*agentproto.BatchUpdateMetadataResponse, error) { + if req == nil { + return nil, drpcerr.WithCode(xerrors.New("metadata request is required"), uint64(codes.InvalidArgument)) + } + // Validate the entire batch before trimming values or discarding excess keys. + for i, md := range req.Metadata { + if md.GetResult() == nil { + return nil, drpcerr.WithCode(xerrors.Errorf("metadata result at index %d is required", i), uint64(codes.InvalidArgument)) + } + } + const ( // maxAllKeysLen is the maximum length of all metadata keys. This is // 6144 to stay below the Postgres NOTIFY limit of 8000 bytes, with some diff --git a/coderd/agentapi/validation_test.go b/coderd/agentapi/validation_test.go new file mode 100644 index 0000000000000..3dfc110178b12 --- /dev/null +++ b/coderd/agentapi/validation_test.go @@ -0,0 +1,151 @@ +package agentapi_test + +import ( + "context" + "strings" + "testing" + + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/protobuf/proto" + "storj.io/drpc/drpcerr" + "storj.io/drpc/drpcmux" + "storj.io/drpc/drpcserver" + + agentproto "github.com/coder/coder/v2/agent/proto" + "github.com/coder/coder/v2/coderd/agentapi" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/codersdk/drpcsdk" + "github.com/coder/coder/v2/testutil" +) + +func TestUpdateLifecycleMissingMessage(t *testing.T) { + t.Parallel() + + for name, req := range map[string]*agentproto.UpdateLifecycleRequest{"NilRequest": nil, "MissingLifecycle": {}} { + t.Run(name, func(t *testing.T) { + t.Parallel() + api := &agentapi.LifecycleAPI{ + AgentFn: func(context.Context) (database.WorkspaceAgent, error) { + return database.WorkspaceAgent{}, nil + }, + Log: testutil.Logger(t), + } + require.NotPanics(t, func() { + resp, err := api.UpdateLifecycle(context.Background(), req) + require.ErrorContains(t, err, "lifecycle is required") + require.EqualValues(t, codes.InvalidArgument, drpcerr.Code(err)) + require.Nil(t, resp) + }) + }) + } +} + +func TestUpdateStartupMissingMessage(t *testing.T) { + t.Parallel() + + for name, req := range map[string]*agentproto.UpdateStartupRequest{"NilRequest": nil, "MissingStartup": {}} { + t.Run(name, func(t *testing.T) { + t.Parallel() + api := &agentapi.LifecycleAPI{ + AgentFn: func(context.Context) (database.WorkspaceAgent, error) { + return database.WorkspaceAgent{}, nil + }, + Log: testutil.Logger(t), + } + ctx := agentapi.WithAPIVersion(context.Background(), "2.0") + require.NotPanics(t, func() { + resp, err := api.UpdateStartup(ctx, req) + require.ErrorContains(t, err, "startup is required") + require.EqualValues(t, codes.InvalidArgument, drpcerr.Code(err)) + require.Nil(t, resp) + }) + }) + } +} + +func TestBatchUpdateMetadataMissingMessage(t *testing.T) { + t.Parallel() + + for _, tt := range []struct { + name string + req *agentproto.BatchUpdateMetadataRequest + }{ + {name: "NilRequest"}, + {name: "NilMetadata", req: &agentproto.BatchUpdateMetadataRequest{Metadata: []*agentproto.Metadata{nil}}}, + {name: "MissingResult", req: &agentproto.BatchUpdateMetadataRequest{Metadata: []*agentproto.Metadata{{Key: "key"}}}}, + {name: "MixedBatch", req: &agentproto.BatchUpdateMetadataRequest{Metadata: []*agentproto.Metadata{ + {Key: "valid", Result: &agentproto.WorkspaceAgentMetadata_Result{Value: " value "}}, + {Key: "invalid"}, + }}}, + {name: "MissingResultAfterKeyLimit", req: &agentproto.BatchUpdateMetadataRequest{Metadata: []*agentproto.Metadata{ + {Key: strings.Repeat("k", 6145), Result: &agentproto.WorkspaceAgentMetadata_Result{Value: " value "}}, + {Key: "invalid"}, + }}}, + } { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + original := proto.Clone(tt.req) + // A nil batcher also catches attempts to enqueue a partial batch. + api := &agentapi.MetadataAPI{Log: testutil.Logger(t)} + require.NotPanics(t, func() { + resp, err := api.BatchUpdateMetadata(context.Background(), tt.req) + require.Error(t, err) + require.EqualValues(t, codes.InvalidArgument, drpcerr.Code(err)) + require.Nil(t, resp) + }) + require.True(t, proto.Equal(original, tt.req), "invalid batches must not be processed") + }) + } +} + +func TestAgentRPCMissingMessages(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + serverCtx, cancel := context.WithCancel(agentapi.WithAPIVersion(ctx, "2.0")) + defer cancel() + conn, listener := drpcsdk.MemTransportPipe() + defer conn.Close() + defer listener.Close() + + api := &agentapi.API{ + LifecycleAPI: &agentapi.LifecycleAPI{ + AgentFn: func(context.Context) (database.WorkspaceAgent, error) { + return database.WorkspaceAgent{}, nil + }, + Log: testutil.Logger(t), + }, + MetadataAPI: &agentapi.MetadataAPI{}, + StatsAPI: &agentapi.StatsAPI{}, + } + mux := drpcmux.New() + require.NoError(t, agentproto.DRPCRegisterAgent(mux, api)) + server := drpcsdk.NewServer(testutil.NewFakeSink(t).Logger(), mux, drpcserver.Options{ + Manager: drpcsdk.DefaultDRPCOptions(nil), + }) + done := make(chan error, 1) + go func() { + done <- server.Serve(serverCtx, listener) + }() + client := agentproto.NewDRPCAgentClient(conn) + + _, err := client.UpdateLifecycle(ctx, &agentproto.UpdateLifecycleRequest{}) + require.ErrorContains(t, err, "lifecycle is required") + require.EqualValues(t, codes.InvalidArgument, drpcerr.Code(err)) + _, err = client.UpdateStartup(ctx, &agentproto.UpdateStartupRequest{}) + require.ErrorContains(t, err, "startup is required") + require.EqualValues(t, codes.InvalidArgument, drpcerr.Code(err)) + _, err = client.BatchUpdateMetadata(ctx, &agentproto.BatchUpdateMetadataRequest{ + Metadata: []*agentproto.Metadata{{Key: "key"}}, + }) + require.ErrorContains(t, err, "metadata result at index 0 is required") + require.EqualValues(t, codes.InvalidArgument, drpcerr.Code(err)) + + _, err = client.UpdateStats(ctx, &agentproto.UpdateStatsRequest{}) + require.NoError(t, err) + _, err = client.BatchUpdateMetadata(ctx, &agentproto.BatchUpdateMetadataRequest{}) + require.NoError(t, err) + cancel() + require.NoError(t, testutil.RequireReceive(ctx, t, done)) +} From 8a29a72860cc86b9eacc752dd2e5acb2df479559 Mon Sep 17 00:00:00 2001 From: Dylan Huff Date: Wed, 9 Sep 2026 22:02:21 +0000 Subject: [PATCH 2/2] test(coderd/agentapi): trim missing message validation tests --- coderd/agentapi/validation_test.go | 107 +++++++---------------------- 1 file changed, 23 insertions(+), 84 deletions(-) diff --git a/coderd/agentapi/validation_test.go b/coderd/agentapi/validation_test.go index 3dfc110178b12..2a32556b2195d 100644 --- a/coderd/agentapi/validation_test.go +++ b/coderd/agentapi/validation_test.go @@ -2,19 +2,14 @@ package agentapi_test import ( "context" - "strings" "testing" "github.com/stretchr/testify/require" "google.golang.org/grpc/codes" - "google.golang.org/protobuf/proto" "storj.io/drpc/drpcerr" - "storj.io/drpc/drpcmux" - "storj.io/drpc/drpcserver" agentproto "github.com/coder/coder/v2/agent/proto" "github.com/coder/coder/v2/coderd/agentapi" - "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/codersdk/drpcsdk" "github.com/coder/coder/v2/testutil" ) @@ -22,46 +17,19 @@ import ( func TestUpdateLifecycleMissingMessage(t *testing.T) { t.Parallel() - for name, req := range map[string]*agentproto.UpdateLifecycleRequest{"NilRequest": nil, "MissingLifecycle": {}} { - t.Run(name, func(t *testing.T) { - t.Parallel() - api := &agentapi.LifecycleAPI{ - AgentFn: func(context.Context) (database.WorkspaceAgent, error) { - return database.WorkspaceAgent{}, nil - }, - Log: testutil.Logger(t), - } - require.NotPanics(t, func() { - resp, err := api.UpdateLifecycle(context.Background(), req) - require.ErrorContains(t, err, "lifecycle is required") - require.EqualValues(t, codes.InvalidArgument, drpcerr.Code(err)) - require.Nil(t, resp) - }) - }) - } + resp, err := (&agentapi.LifecycleAPI{}).UpdateLifecycle(context.Background(), &agentproto.UpdateLifecycleRequest{}) + require.ErrorContains(t, err, "lifecycle is required") + require.EqualValues(t, codes.InvalidArgument, drpcerr.Code(err)) + require.Nil(t, resp) } func TestUpdateStartupMissingMessage(t *testing.T) { t.Parallel() - for name, req := range map[string]*agentproto.UpdateStartupRequest{"NilRequest": nil, "MissingStartup": {}} { - t.Run(name, func(t *testing.T) { - t.Parallel() - api := &agentapi.LifecycleAPI{ - AgentFn: func(context.Context) (database.WorkspaceAgent, error) { - return database.WorkspaceAgent{}, nil - }, - Log: testutil.Logger(t), - } - ctx := agentapi.WithAPIVersion(context.Background(), "2.0") - require.NotPanics(t, func() { - resp, err := api.UpdateStartup(ctx, req) - require.ErrorContains(t, err, "startup is required") - require.EqualValues(t, codes.InvalidArgument, drpcerr.Code(err)) - require.Nil(t, resp) - }) - }) - } + resp, err := (&agentapi.LifecycleAPI{}).UpdateStartup(context.Background(), &agentproto.UpdateStartupRequest{}) + require.ErrorContains(t, err, "startup is required") + require.EqualValues(t, codes.InvalidArgument, drpcerr.Code(err)) + require.Nil(t, resp) } func TestBatchUpdateMetadataMissingMessage(t *testing.T) { @@ -72,34 +40,23 @@ func TestBatchUpdateMetadataMissingMessage(t *testing.T) { req *agentproto.BatchUpdateMetadataRequest }{ {name: "NilRequest"}, - {name: "NilMetadata", req: &agentproto.BatchUpdateMetadataRequest{Metadata: []*agentproto.Metadata{nil}}}, - {name: "MissingResult", req: &agentproto.BatchUpdateMetadataRequest{Metadata: []*agentproto.Metadata{{Key: "key"}}}}, {name: "MixedBatch", req: &agentproto.BatchUpdateMetadataRequest{Metadata: []*agentproto.Metadata{ - {Key: "valid", Result: &agentproto.WorkspaceAgentMetadata_Result{Value: " value "}}, - {Key: "invalid"}, - }}}, - {name: "MissingResultAfterKeyLimit", req: &agentproto.BatchUpdateMetadataRequest{Metadata: []*agentproto.Metadata{ - {Key: strings.Repeat("k", 6145), Result: &agentproto.WorkspaceAgentMetadata_Result{Value: " value "}}, + {Key: "valid", Result: &agentproto.WorkspaceAgentMetadata_Result{Value: "value"}}, {Key: "invalid"}, }}}, } { t.Run(tt.name, func(t *testing.T) { t.Parallel() - original := proto.Clone(tt.req) - // A nil batcher also catches attempts to enqueue a partial batch. - api := &agentapi.MetadataAPI{Log: testutil.Logger(t)} - require.NotPanics(t, func() { - resp, err := api.BatchUpdateMetadata(context.Background(), tt.req) - require.Error(t, err) - require.EqualValues(t, codes.InvalidArgument, drpcerr.Code(err)) - require.Nil(t, resp) - }) - require.True(t, proto.Equal(original, tt.req), "invalid batches must not be processed") + + resp, err := (&agentapi.MetadataAPI{}).BatchUpdateMetadata(context.Background(), tt.req) + require.Error(t, err) + require.EqualValues(t, codes.InvalidArgument, drpcerr.Code(err)) + require.Nil(t, resp) }) } } -func TestAgentRPCMissingMessages(t *testing.T) { +func TestAgentRPCMissingMessage(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitShort) @@ -110,42 +67,24 @@ func TestAgentRPCMissingMessages(t *testing.T) { defer listener.Close() api := &agentapi.API{ - LifecycleAPI: &agentapi.LifecycleAPI{ - AgentFn: func(context.Context) (database.WorkspaceAgent, error) { - return database.WorkspaceAgent{}, nil - }, - Log: testutil.Logger(t), - }, - MetadataAPI: &agentapi.MetadataAPI{}, - StatsAPI: &agentapi.StatsAPI{}, + LifecycleAPI: &agentapi.LifecycleAPI{}, + StatsAPI: &agentapi.StatsAPI{}, } - mux := drpcmux.New() - require.NoError(t, agentproto.DRPCRegisterAgent(mux, api)) - server := drpcsdk.NewServer(testutil.NewFakeSink(t).Logger(), mux, drpcserver.Options{ - Manager: drpcsdk.DefaultDRPCOptions(nil), - }) + server, err := api.Server(serverCtx) + require.NoError(t, err) done := make(chan error, 1) go func() { done <- server.Serve(serverCtx, listener) }() client := agentproto.NewDRPCAgentClient(conn) - _, err := client.UpdateLifecycle(ctx, &agentproto.UpdateLifecycleRequest{}) + _, err = client.UpdateLifecycle(ctx, &agentproto.UpdateLifecycleRequest{}) require.ErrorContains(t, err, "lifecycle is required") require.EqualValues(t, codes.InvalidArgument, drpcerr.Code(err)) - _, err = client.UpdateStartup(ctx, &agentproto.UpdateStartupRequest{}) - require.ErrorContains(t, err, "startup is required") - require.EqualValues(t, codes.InvalidArgument, drpcerr.Code(err)) - _, err = client.BatchUpdateMetadata(ctx, &agentproto.BatchUpdateMetadataRequest{ - Metadata: []*agentproto.Metadata{{Key: "key"}}, - }) - require.ErrorContains(t, err, "metadata result at index 0 is required") - require.EqualValues(t, codes.InvalidArgument, drpcerr.Code(err)) - - _, err = client.UpdateStats(ctx, &agentproto.UpdateStatsRequest{}) - require.NoError(t, err) - _, err = client.BatchUpdateMetadata(ctx, &agentproto.BatchUpdateMetadataRequest{}) + stats, err := client.UpdateStats(ctx, &agentproto.UpdateStatsRequest{}) require.NoError(t, err) + require.NotNil(t, stats.ReportInterval) + cancel() require.NoError(t, testutil.RequireReceive(ctx, t, done)) }