fix: allow startup scripts larger than 32k

spikecurtis · spikecurtis · commit 369cb187cddb · 2024-02-07T22:06:38.000+04:00
diff --git a/coderd/workspaceagentsrpc_test.go b/coderd/workspaceagentsrpc_test.go
@@ -0,0 +1,48 @@
+package coderd_test
+
+import (
+	"context"
+	"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"
+)
+
+func TestAgentAPI_LargeManifest(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	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 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
@@ -202,6 +202,7 @@ func (c *Client) ConnectRPC(ctx context.Context) (drpc.Conn, error) {
 		}
 		return nil, codersdk.ReadBodyAsError(res)
 	}
+	conn.SetReadLimit(1 << 22) // 4 MiB -- about the limit for protobufs
 
 	_, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageBinary)
 

Original file line number	Diff line number	Diff line change
`@@ -202,6 +202,7 @@ func (c *Client) ConnectRPC(ctx context.Context) (drpc.Conn, error) {`
`202`	`202`	`}`
`203`	`203`	`return nil, codersdk.ReadBodyAsError(res)`
`204`	`204`	`}`
	`205`	`+ conn.SetReadLimit(1 << 22) // 4 MiB -- about the limit for protobufs`
`205`	`206`
`206`	`207`	`_, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageBinary)`
`207`	`208`