fix: stop updating agent stats from deleted workspaces (#11026)

f0ssel · Emyrk · web-flow · commit 228cbec99bc8 · 2023-12-07T13:55:29.000-05:00
Co-authored-by: Steven Masley &lt;stevenmasley@gmail.com&gt;
diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -619,11 +619,15 @@ func CreateAnotherUserMutators(t testing.TB, client *codersdk.Client, organizati
 }
 
 // AuthzUserSubject does not include the user's groups.
-func AuthzUserSubject(user codersdk.User) rbac.Subject {
+func AuthzUserSubject(user codersdk.User, orgID uuid.UUID) rbac.Subject {
 	roles := make(rbac.RoleNames, 0, len(user.Roles))
+	// Member role is always implied
+	roles = append(roles, rbac.RoleMember())
 	for _, r := range user.Roles {
 		roles = append(roles, r.Name)
 	}
+	// We assume only 1 org exists
+	roles = append(roles, rbac.RoleOrgMember(orgID))
 
 	return rbac.Subject{
 		ID:     user.ID.String(),
diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -3762,6 +3762,9 @@ func (q *FakeQuerier) GetWorkspaceAgentAndOwnerByAuthToken(_ context.Context, au
 					if build.WorkspaceID != ws.ID {
 						continue
 					}
+					if ws.Deleted {
+						continue
+					}
 					var row database.GetWorkspaceAgentAndOwnerByAuthTokenRow
 					row.WorkspaceID = ws.ID
 					usr, err := q.getUserByIDNoLock(ws.OwnerID)
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaceagents.sql b/coderd/database/queries/workspaceagents.sql
@@ -252,9 +252,10 @@ FROM users
 WHERE
 	-- TODO: we can add more conditions here, such as:
 	-- 1) The user must be active
-	-- 2) The user must not be deleted
-	-- 3) The workspace must be running
+	-- 2) The workspace must be running
 	workspace_agents.auth_token = @auth_token
+AND
+	workspaces.deleted = FALSE
 GROUP BY
 	workspace_agents.id,
 	workspaces.id,
diff --git a/coderd/templates_test.go b/coderd/templates_test.go
@@ -560,7 +560,7 @@ func TestPatchTemplateMeta(t *testing.T) {
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 		defer cancel()
 		// nolint:gocritic // Setting up unit test data
-		err := db.UpdateTemplateAccessControlByID(dbauthz.As(ctx, coderdtest.AuthzUserSubject(tplAdmin)), database.UpdateTemplateAccessControlByIDParams{
+		err := db.UpdateTemplateAccessControlByID(dbauthz.As(ctx, coderdtest.AuthzUserSubject(tplAdmin, user.OrganizationID)), database.UpdateTemplateAccessControlByIDParams{
 			ID:                   template.ID,
 			RequireActiveVersion: false,
 			Deprecated:           "Some deprecated message",
diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go
@@ -26,10 +26,12 @@ import (
 	"github.com/coder/coder/v2/coderd"
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbauthz"
 	"github.com/coder/coder/v2/coderd/database/dbfake"
 	"github.com/coder/coder/v2/coderd/database/dbmem"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/database/pubsub"
+	"github.com/coder/coder/v2/coderd/rbac"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/coder/v2/provisioner/echo"
@@ -876,6 +878,63 @@ func TestWorkspaceAgentReportStats(t *testing.T) {
 			"%s is not after %s", newWorkspace.LastUsedAt, r.Workspace.LastUsedAt,
 		)
 	})
+
+	t.Run("FailDeleted", func(t *testing.T) {
+		t.Parallel()
+
+		owner, db := coderdtest.NewWithDatabase(t, nil)
+		ownerUser := coderdtest.CreateFirstUser(t, owner)
+		client, admin := coderdtest.CreateAnotherUser(t, owner, ownerUser.OrganizationID, rbac.RoleTemplateAdmin(), rbac.RoleUserAdmin())
+		r := dbfake.WorkspaceBuild(t, db, database.Workspace{
+			OrganizationID: admin.OrganizationIDs[0],
+			OwnerID:        admin.ID,
+		}).WithAgent().Do()
+
+		agentClient := agentsdk.New(client.URL)
+		agentClient.SetSessionToken(r.AgentToken)
+
+		_, err := agentClient.PostStats(context.Background(), &agentsdk.Stats{
+			ConnectionsByProto:          map[string]int64{"TCP": 1},
+			ConnectionCount:             1,
+			RxPackets:                   1,
+			RxBytes:                     1,
+			TxPackets:                   1,
+			TxBytes:                     1,
+			SessionCountVSCode:          0,
+			SessionCountJetBrains:       0,
+			SessionCountReconnectingPTY: 0,
+			SessionCountSSH:             0,
+			ConnectionMedianLatencyMS:   10,
+		})
+		require.NoError(t, err)
+
+		newWorkspace, err := client.Workspace(context.Background(), r.Workspace.ID)
+		require.NoError(t, err)
+
+		// nolint:gocritic // using db directly over creating a delete job
+		err = db.UpdateWorkspaceDeletedByID(dbauthz.As(context.Background(),
+			coderdtest.AuthzUserSubject(admin, ownerUser.OrganizationID)),
+			database.UpdateWorkspaceDeletedByIDParams{
+				ID:      newWorkspace.ID,
+				Deleted: true,
+			})
+		require.NoError(t, err)
+
+		_, err = agentClient.PostStats(context.Background(), &agentsdk.Stats{
+			ConnectionsByProto:          map[string]int64{"TCP": 1},
+			ConnectionCount:             1,
+			RxPackets:                   1,
+			RxBytes:                     1,
+			TxPackets:                   1,
+			TxBytes:                     1,
+			SessionCountVSCode:          1,
+			SessionCountJetBrains:       0,
+			SessionCountReconnectingPTY: 0,
+			SessionCountSSH:             0,
+			ConnectionMedianLatencyMS:   10,
+		})
+		require.ErrorContains(t, err, "agent is invalid")
+	})
 }
 
 func TestWorkspaceAgent_LifecycleState(t *testing.T) {
