diff --git a/coderd/workspaceagentsrpc_test.go b/coderd/workspaceagentsrpc_test.go new file mode 100644 index 0000000000..a92fbdcd1c --- /dev/null +++ b/coderd/workspaceagentsrpc_test.go @@ -0,0 +1,50 @@ +package coderd_test + +import ( + "testing" + + "github.com/stretchr/testify/require" + + agentproto "github.com/coder/coder/v2/agent/proto" + "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbfake" + "github.com/coder/coder/v2/codersdk/agentsdk" + "github.com/coder/coder/v2/provisionersdk/proto" + "github.com/coder/coder/v2/testutil" +) + +func TestAgentAPI_LargeManifest(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + client, store := coderdtest.NewWithDatabase(t, nil) + adminUser := coderdtest.CreateFirstUser(t, client) + n := 512000 + longScript := make([]byte, n) + for i := range longScript { + longScript[i] = 'q' + } + r := dbfake.WorkspaceBuild(t, store, database.Workspace{ + OrganizationID: adminUser.OrganizationID, + OwnerID: adminUser.UserID, + }).WithAgent(func(agents []*proto.Agent) []*proto.Agent { + agents[0].Scripts = []*proto.Script{ + { + Script: string(longScript), + }, + } + return agents + }).Do() + ac := agentsdk.New(client.URL) + ac.SetSessionToken(r.AgentToken) + conn, err := ac.ConnectRPC(ctx) + defer func() { + _ = conn.Close() + }() + require.NoError(t, err) + agentAPI := agentproto.NewDRPCAgentClient(conn) + manifest, err := agentAPI.GetManifest(ctx, &agentproto.GetManifestRequest{}) + require.NoError(t, err) + require.Len(t, manifest.Scripts, 1) + require.Len(t, manifest.Scripts[0].Script, n) +} diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go index 82a8e697ee..22cc0faab6 100644 --- a/codersdk/agentsdk/agentsdk.go +++ b/codersdk/agentsdk/agentsdk.go @@ -629,6 +629,9 @@ func (c *wsNetConn) Close() error { // during read or write will cancel the context, but not close the // conn. Close should be called to release context resources. func websocketNetConn(ctx context.Context, conn *websocket.Conn, msgType websocket.MessageType) (context.Context, net.Conn) { + // Set the read limit to 4 MiB -- about the limit for protobufs. This needs to be larger than + // the default because some of our protocols can include large messages like startup scripts. + conn.SetReadLimit(1 << 22) ctx, cancel := context.WithCancel(ctx) nc := websocket.NetConn(ctx, conn, msgType) return ctx, &wsNetConn{