Skip to content

Commit 7f1d7c4

Browse files
ammarioammario
committed
fix: correctly reject quota-violating builds
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. Side note, I changed workspaceproxy.go to return a nil value in a case to pass staticcheck.
1 parent 4a9c773 commit 7f1d7c4

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)