fix: Multiple builds using the incorrect agent token

kylecarbs · kylecarbs · commit be08e2bedbeb · 2022-04-12T18:15:04.000Z
This was an issue with our in-memory database that caused
newer builds to return an outdated agent, which would then
be rejected.

A test case has been added to ensure this can't happen again!
diff --git a/coderd/database/databasefake/databasefake.go b/coderd/database/databasefake/databasefake.go
@@ -626,7 +626,9 @@ func (q *fakeQuerier) GetWorkspaceAgentByAuthToken(_ context.Context, authToken
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()
 
-	for _, agent := range q.provisionerJobAgent {
+	// The schema sorts this by created at, so we iterate the array backwards.
+	for i := len(q.provisionerJobAgent) - 1; i >= 0; i-- {
+		agent := q.provisionerJobAgent[i]
 		if agent.AuthToken.String() == authToken.String() {
 			return agent, nil
 		}
@@ -990,7 +992,7 @@ func (q *fakeQuerier) InsertProvisionerJob(_ context.Context, arg database.Inser
 	return job, nil
 }
 
-func (q *fakeQuerier) InsertWorkspaceAgent(_ context.Context, arg database.InsertWorkspaceAgentParams) (database.WorkspaceAgent, error) {
+func (q *fakeQuerier) InsertWorkspaceAgent(ctx context.Context, arg database.InsertWorkspaceAgentParams) (database.WorkspaceAgent, error) {
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
 
diff --git a/coderd/workspacebuilds_test.go b/coderd/workspacebuilds_test.go
@@ -6,10 +6,16 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 
+	"cdr.dev/slog/sloggers/slogtest"
+
+	"github.com/coder/coder/agent"
 	"github.com/coder/coder/coderd/coderdtest"
+	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/peer"
 	"github.com/coder/coder/provisioner/echo"
 	"github.com/coder/coder/provisionersdk/proto"
 )
@@ -115,6 +121,56 @@ func TestWorkspaceBuildResources(t *testing.T) {
 		require.Equal(t, "example", resources[0].Type)
 		require.Len(t, resources[0].Agents, 1)
 	})
+	t.Run("MultipleBuilds", func(t *testing.T) {
+		// This ensures that the latest agent can connect through multiple builds.
+		// Agents won't always have new authentication tokens on rebuild, because
+		// infrastructure as code won't always recreate the resource.
+		t.Parallel()
+		client := coderdtest.New(t, nil)
+		user := coderdtest.CreateFirstUser(t, client)
+		coderdtest.NewProvisionerDaemon(t, client)
+		authToken := uuid.NewString()
+		version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+			Parse: echo.ParseComplete,
+			Provision: []*proto.Provision_Response{{
+				Type: &proto.Provision_Response_Complete{
+					Complete: &proto.Provision_Complete{
+						Resources: []*proto.Resource{{
+							Name: "some",
+							Type: "example",
+							Agents: []*proto.Agent{{
+								Id: "something",
+								Auth: &proto.Agent_Token{
+									Token: authToken,
+								},
+							}},
+						}},
+					},
+				},
+			}},
+		})
+		coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+		workspace := coderdtest.CreateWorkspace(t, client, codersdk.Me, template.ID)
+		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+		build, err := client.CreateWorkspaceBuild(context.Background(), workspace.ID, codersdk.CreateWorkspaceBuildRequest{
+			TemplateVersionID: template.ActiveVersionID,
+			Transition:        database.WorkspaceTransitionStart,
+			DryRun:            false,
+		})
+		require.NoError(t, err)
+		coderdtest.AwaitWorkspaceBuildJob(t, client, build.ID)
+
+		agentClient := codersdk.New(client.URL)
+		agentClient.SessionToken = authToken
+		agentCloser := agent.New(agentClient.ListenWorkspaceAgent, &peer.ConnOptions{
+			Logger: slogtest.Make(t, nil),
+		})
+		t.Cleanup(func() {
+			_ = agentCloser.Close()
+		})
+		coderdtest.AwaitWorkspaceAgents(t, client, build.ID)
+	})
 }
 
 func TestWorkspaceBuildLogs(t *testing.T) {
