Skip to content

Commit 545a256

Browse files
ammarioammario
authored
fix: correctly reject quota-violating builds (#9233)
Due to a logical error in CommitQuota, all workspace Stop->Start operations were being accepted, regardless of the Quota limit. This issue only appeared after #9201, so this was a minor regression in main for about 3 days. This PR adds a test to make sure this kind of bug doesn't recur. To make the new test possible, we give the echo provisioner the ability to simulate responses to specific transitions.
1 parent 69ec8d7 commit 545a256

File tree

7 files changed

+292
-54
lines changed

7 files changed

+292
-54
lines changed

‎enterprise/coderd/coderd.go

+4-1
Original file line numberDiff line numberDiff line change
@@ -498,7 +498,10 @@ func (api *API) updateEntitlements(ctx context.Context) error {
498498

499499
if initial, changed, enabled := featureChanged(codersdk.FeatureTemplateRBAC); shouldUpdate(initial, changed, enabled) {
500500
if enabled {
501-
committer := committer{Database: api.Database}
501+
committer := committer{
502+
Log: api.Logger.Named("quota_committer"),
503+
Database: api.Database,
504+
}
502505
ptr := proto.QuotaCommitter(&committer)
503506
api.AGPL.QuotaCommitter.Store(&ptr)
504507
} else {

‎enterprise/coderd/workspacequota.go

+23-10
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,12 @@ package coderd
33
import (
44
"context"
55
"database/sql"
6+
"errors"
67
"net/http"
78

89
"github.com/google/uuid"
9-
"golang.org/x/xerrors"
10+
11+
"cdr.dev/slog"
1012

1113
"github.com/coder/coder/v2/coderd/database"
1214
"github.com/coder/coder/v2/coderd/httpapi"
@@ -17,6 +19,7 @@ import (
1719
)
1820

1921
type committer struct {
22+
Log slog.Logger
2023
Database database.Store
2124
}
2225

@@ -28,12 +31,12 @@ func (c *committer) CommitQuota(
2831
return nil, err
2932
}
3033

31-
build, err := c.Database.GetWorkspaceBuildByJobID(ctx, jobID)
34+
nextBuild, err := c.Database.GetWorkspaceBuildByJobID(ctx, jobID)
3235
if err != nil {
3336
return nil, err
3437
}
3538

36-
workspace, err := c.Database.GetWorkspaceByID(ctx, build.WorkspaceID)
39+
workspace, err := c.Database.GetWorkspaceByID(ctx, nextBuild.WorkspaceID)
3740
if err != nil {
3841
return nil, err
3942
}
@@ -58,25 +61,35 @@ func (c *committer) CommitQuota(
5861
// If the new build will reduce overall quota consumption, then we
5962
// allow it even if the user is over quota.
6063
netIncrease := true
61-
previousBuild, err := s.GetWorkspaceBuildByWorkspaceIDAndBuildNumber(ctx, database.GetWorkspaceBuildByWorkspaceIDAndBuildNumberParams{
64+
prevBuild, err := s.GetWorkspaceBuildByWorkspaceIDAndBuildNumber(ctx, database.GetWorkspaceBuildByWorkspaceIDAndBuildNumberParams{
6265
WorkspaceID: workspace.ID,
63-
BuildNumber: build.BuildNumber - 1,
66+
BuildNumber: nextBuild.BuildNumber - 1,
6467
})
6568
if err == nil {
66-
if build.DailyCost < previousBuild.DailyCost {
67-
netIncrease = false
68-
}
69-
} else if !xerrors.Is(err, sql.ErrNoRows) {
69+
netIncrease = request.DailyCost >= prevBuild.DailyCost
70+
c.Log.Debug(
71+
ctx, "previous build cost",
72+
slog.F("prev_cost", prevBuild.DailyCost),
73+
slog.F("next_cost", request.DailyCost),
74+
slog.F("net_increase", netIncrease),
75+
)
76+
} else if !errors.Is(err, sql.ErrNoRows) {
7077
return err
7178
}
7279

7380
newConsumed := int64(request.DailyCost) + consumed
7481
if newConsumed > budget && netIncrease {
82+
c.Log.Debug(
83+
ctx, "over quota, rejecting",
84+
slog.F("prev_consumed", consumed),
85+
slog.F("next_consumed", newConsumed),
86+
slog.F("budget", budget),
87+
)
7588
return nil
7689
}
7790

7891
err = s.UpdateWorkspaceBuildCostByID(ctx, database.UpdateWorkspaceBuildCostByIDParams{
79-
ID: build.ID,
92+
ID: nextBuild.ID,
8093
DailyCost: request.DailyCost,
8194
})
8295
if err != nil {

‎enterprise/coderd/workspacequota_test.go

+105-5
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"github.com/stretchr/testify/require"
1111

1212
"github.com/coder/coder/v2/coderd/coderdtest"
13+
"github.com/coder/coder/v2/coderd/database"
1314
"github.com/coder/coder/v2/coderd/util/ptr"
1415
"github.com/coder/coder/v2/codersdk"
1516
"github.com/coder/coder/v2/enterprise/coderd/coderdenttest"
@@ -31,12 +32,13 @@ func verifyQuota(ctx context.Context, t *testing.T, client *codersdk.Client, con
3132
}
3233

3334
func TestWorkspaceQuota(t *testing.T) {
34-
// TODO: refactor for new impl
35-
3635
t.Parallel()
3736

38-
t.Run("BlocksBuild", func(t *testing.T) {
37+
// This first test verifies the behavior of creating and deleting workspaces.
38+
// It also tests multi-group quota stacking and the everyone group.
39+
t.Run("CreateDelete", func(t *testing.T) {
3940
t.Parallel()
41+
4042
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
4143
defer cancel()
4244
max := 1
@@ -49,8 +51,6 @@ func TestWorkspaceQuota(t *testing.T) {
4951
},
5052
})
5153
coderdtest.NewProvisionerDaemon(t, api.AGPL)
52-
coderdtest.NewProvisionerDaemon(t, api.AGPL)
53-
coderdtest.NewProvisionerDaemon(t, api.AGPL)
5454

5555
verifyQuota(ctx, t, client, 0, 0)
5656

@@ -157,4 +157,104 @@ func TestWorkspaceQuota(t *testing.T) {
157157
verifyQuota(ctx, t, client, 4, 4)
158158
require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
159159
})
160+
161+
t.Run("StartStop", func(t *testing.T) {
162+
t.Parallel()
163+
164+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
165+
defer cancel()
166+
max := 1
167+
client, _, api, user := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
168+
UserWorkspaceQuota: max,
169+
LicenseOptions: &coderdenttest.LicenseOptions{
170+
Features: license.Features{
171+
codersdk.FeatureTemplateRBAC: 1,
172+
},
173+
},
174+
})
175+
coderdtest.NewProvisionerDaemon(t, api.AGPL)
176+
177+
verifyQuota(ctx, t, client, 0, 0)
178+
179+
// Patch the 'Everyone' group to verify its quota allowance is being accounted for.
180+
_, err := client.PatchGroup(ctx, user.OrganizationID, codersdk.PatchGroupRequest{
181+
QuotaAllowance: ptr.Ref(4),
182+
})
183+
require.NoError(t, err)
184+
verifyQuota(ctx, t, client, 0, 4)
185+
186+
stopResp := []*proto.Provision_Response{{
187+
Type: &proto.Provision_Response_Complete{
188+
Complete: &proto.Provision_Complete{
189+
Resources: []*proto.Resource{{
190+
Name: "example",
191+
Type: "aws_instance",
192+
DailyCost: 1,
193+
}},
194+
},
195+
},
196+
}}
197+
198+
startResp := []*proto.Provision_Response{{
199+
Type: &proto.Provision_Response_Complete{
200+
Complete: &proto.Provision_Complete{
201+
Resources: []*proto.Resource{{
202+
Name: "example",
203+
Type: "aws_instance",
204+
DailyCost: 2,
205+
}},
206+
},
207+
},
208+
}}
209+
210+
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
211+
Parse: echo.ParseComplete,
212+
ProvisionPlanMap: map[proto.WorkspaceTransition][]*proto.Provision_Response{
213+
proto.WorkspaceTransition_START: startResp,
214+
proto.WorkspaceTransition_STOP: stopResp,
215+
},
216+
ProvisionApplyMap: map[proto.WorkspaceTransition][]*proto.Provision_Response{
217+
proto.WorkspaceTransition_START: startResp,
218+
proto.WorkspaceTransition_STOP: stopResp,
219+
},
220+
})
221+
222+
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
223+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
224+
225+
// Spin up two workspaces.
226+
var wg sync.WaitGroup
227+
var workspaces []codersdk.Workspace
228+
for i := 0; i < 2; i++ {
229+
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
230+
workspaces = append(workspaces, workspace)
231+
build := coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
232+
assert.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
233+
}
234+
wg.Wait()
235+
verifyQuota(ctx, t, client, 4, 4)
236+
237+
// Next one must fail
238+
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
239+
build := coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
240+
require.Contains(t, build.Job.Error, "quota")
241+
242+
// Consumed shouldn't bump
243+
verifyQuota(ctx, t, client, 4, 4)
244+
require.Equal(t, codersdk.WorkspaceStatusFailed, build.Status)
245+
246+
build = coderdtest.CreateWorkspaceBuild(t, client, workspaces[0], database.WorkspaceTransitionStop)
247+
build = coderdtest.AwaitWorkspaceBuildJob(t, client, build.ID)
248+
249+
// Quota goes down one
250+
verifyQuota(ctx, t, client, 3, 4)
251+
require.Equal(t, codersdk.WorkspaceStatusStopped, build.Status)
252+
253+
build = coderdtest.CreateWorkspaceBuild(t, client, workspaces[0], database.WorkspaceTransitionStart)
254+
build = coderdtest.AwaitWorkspaceBuildJob(t, client, build.ID)
255+
256+
// Quota goes back up
257+
verifyQuota(ctx, t, client, 4, 4)
258+
require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
259+
})
160260
}

‎enterprise/tailnet/workspaceproxy.go

+4
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"bytes"
55
"context"
66
"encoding/json"
7+
"errors"
78
"net"
89
"time"
910

@@ -26,6 +27,9 @@ func ServeWorkspaceProxy(ctx context.Context, conn net.Conn, ma agpl.MultiAgentC
2627
var msg wsproxysdk.CoordinateMessage
2728
err := decoder.Decode(&msg)
2829
if err != nil {
30+
if errors.Is(err, net.ErrClosed) {
31+
return nil
32+
}
2933
return xerrors.Errorf("read json: %w", err)
3034
}
3135

0 commit comments

Comments
 (0)