fix(coderd): allow agent auth during workspace shutdown (#21538)

Agents were losing authentication during workspace shutdown, causing
shutdown scripts to fail. The auth query required agents to belong to
the latest build, but during shutdown a `stop` build becomes latest while
the `start` build's agents are still running.

Modified the auth query to allow `start` build agents to authenticate
temporarily during `stop` execution. The query allows auth when:

- Agent's `start` build job succeeded
- Latest build is `stop` with `pending`/`running` job status
- Builds are adjacent (`stop` is `build_number + 1`)
- Template versions match

Auth closes once `stop` completes.

Renamed `GetWorkspaceAgentAndLatestBuildByAuthToken` to
`GetAuthenticatedWorkspaceAgentAndBuildByAuthToken` since it returns the
agent's build (not always latest) during shutdown.

Closes coder/internal#1249
Fixes #19467
This commit is contained in:
Mathias Fredriksson
2026-01-21 13:18:43 +00:00
committed by GitHub
parent a14a22eb54
commit 97e8a5b093
15 changed files with 726 additions and 72 deletions
+1 -1
View File
@@ -92,7 +92,7 @@ func ExtractWorkspaceAgentAndLatestBuild(opts ExtractWorkspaceAgentAndLatestBuil
}
//nolint:gocritic // System needs to be able to get workspace agents.
row, err := opts.DB.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), token)
row, err := opts.DB.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), token)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
optionalWrite(http.StatusUnauthorized, codersdk.Response{
+188
View File
@@ -1,9 +1,11 @@
package httpmw_test
import (
"database/sql"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/go-chi/chi/v5"
"github.com/google/uuid"
@@ -12,6 +14,7 @@ import (
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbgen"
"github.com/coder/coder/v2/coderd/database/dbtestutil"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/httpmw"
"github.com/coder/coder/v2/codersdk"
)
@@ -95,6 +98,180 @@ func TestWorkspaceAgent(t *testing.T) {
t.Cleanup(func() { _ = res.Body.Close() })
require.Equal(t, http.StatusUnauthorized, res.StatusCode)
})
t.Run("DuringShutdown", func(t *testing.T) {
t.Parallel()
db, ps := dbtestutil.NewDB(t)
authToken := uuid.New()
req, rtr, ws, tpv := setup(t, db, authToken, httpmw.ExtractWorkspaceAgentAndLatestBuild(
httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{
DB: db,
Optional: false,
}),
)
// Create a STOP build with running job (becomes latest).
stopJob := dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
OrganizationID: ws.OrganizationID,
JobStatus: database.ProvisionerJobStatusRunning,
})
_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
WorkspaceID: ws.ID,
JobID: stopJob.ID,
TemplateVersionID: tpv.ID,
BuildNumber: 2,
Transition: database.WorkspaceTransitionStop,
})
// Agent should still authenticate during shutdown.
rw := httptest.NewRecorder()
req.Header.Set(codersdk.SessionTokenHeader, authToken.String())
rtr.ServeHTTP(rw, req)
//nolint:bodyclose // Closed in `t.Cleanup`
res := rw.Result()
t.Cleanup(func() { _ = res.Body.Close() })
require.Equal(t, http.StatusOK, res.StatusCode, "agent should authenticate during stop build execution")
})
t.Run("AfterShutdownCompletes", func(t *testing.T) {
t.Parallel()
db, ps := dbtestutil.NewDB(t)
authToken := uuid.New()
req, rtr, ws, tpv := setup(t, db, authToken, httpmw.ExtractWorkspaceAgentAndLatestBuild(
httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{
DB: db,
Optional: false,
}),
)
// Create a STOP build with completed job (becomes latest).
stopJob := dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
OrganizationID: ws.OrganizationID,
JobStatus: database.ProvisionerJobStatusSucceeded,
CompletedAt: sql.NullTime{
Time: dbtime.Now(),
Valid: true,
},
})
_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
WorkspaceID: ws.ID,
JobID: stopJob.ID,
TemplateVersionID: tpv.ID,
BuildNumber: 2,
Transition: database.WorkspaceTransitionStop,
})
// Agent should NOT authenticate after stop job completes.
rw := httptest.NewRecorder()
req.Header.Set(codersdk.SessionTokenHeader, authToken.String())
rtr.ServeHTTP(rw, req)
//nolint:bodyclose // Closed in `t.Cleanup`
res := rw.Result()
t.Cleanup(func() { _ = res.Body.Close() })
require.Equal(t, http.StatusUnauthorized, res.StatusCode, "agent should not authenticate after stop job completes")
})
t.Run("FailedStartBuild", func(t *testing.T) {
t.Parallel()
db, ps := dbtestutil.NewDB(t)
authToken := uuid.New()
org := dbgen.Organization(t, db, database.Organization{})
user := dbgen.User(t, db, database.User{
Status: database.UserStatusActive,
})
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
UserID: user.ID,
OrganizationID: org.ID,
})
templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
OrganizationID: org.ID,
CreatedBy: user.ID,
})
template := dbgen.Template(t, db, database.Template{
OrganizationID: org.ID,
ActiveVersionID: templateVersion.ID,
CreatedBy: user.ID,
})
workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
OwnerID: user.ID,
OrganizationID: org.ID,
TemplateID: template.ID,
})
// Create START build with FAILED job status.
startJob := dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
OrganizationID: org.ID,
JobStatus: database.ProvisionerJobStatusFailed,
StartedAt: sql.NullTime{
Time: dbtime.Now().Add(-time.Minute),
Valid: true,
},
CompletedAt: sql.NullTime{
Time: dbtime.Now(),
Valid: true,
},
Error: sql.NullString{
String: "build failed",
Valid: true,
},
ErrorCode: sql.NullString{
String: "FAILED",
Valid: true,
},
})
resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
JobID: startJob.ID,
})
_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
WorkspaceID: workspace.ID,
JobID: startJob.ID,
TemplateVersionID: templateVersion.ID,
BuildNumber: 1,
Transition: database.WorkspaceTransitionStart,
})
_ = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
ResourceID: resource.ID,
AuthToken: authToken,
})
// Create a STOP build with running job.
stopJob := dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
OrganizationID: org.ID,
JobStatus: database.ProvisionerJobStatusRunning,
})
_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
WorkspaceID: workspace.ID,
JobID: stopJob.ID,
TemplateVersionID: templateVersion.ID,
BuildNumber: 2,
Transition: database.WorkspaceTransitionStop,
})
req := httptest.NewRequest("GET", "/", nil)
rtr := chi.NewRouter()
rtr.Use(httpmw.ExtractWorkspaceAgentAndLatestBuild(
httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{
DB: db,
Optional: false,
}))
rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
_ = httpmw.WorkspaceAgent(r)
rw.WriteHeader(http.StatusOK)
})
// Agent should NOT authenticate (start build failed).
rw := httptest.NewRecorder()
req.Header.Set(codersdk.SessionTokenHeader, authToken.String())
rtr.ServeHTTP(rw, req)
//nolint:bodyclose // Closed in `t.Cleanup`
res := rw.Result()
t.Cleanup(func() { _ = res.Body.Close() })
require.Equal(t, http.StatusUnauthorized, res.StatusCode, "agent should not authenticate when start build failed")
})
}
func setup(t testing.TB, db database.Store, authToken uuid.UUID, mw func(http.Handler) http.Handler) (*http.Request, http.Handler, database.WorkspaceTable, database.TemplateVersion) {
@@ -123,6 +300,15 @@ func setup(t testing.TB, db database.Store, authToken uuid.UUID, mw func(http.Ha
})
job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
OrganizationID: org.ID,
JobStatus: database.ProvisionerJobStatusSucceeded,
StartedAt: sql.NullTime{
Time: dbtime.Now().Add(-30 * time.Second),
Valid: true,
},
CompletedAt: sql.NullTime{
Time: dbtime.Now(),
Valid: true,
},
})
resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
JobID: job.ID,
@@ -131,6 +317,8 @@ func setup(t testing.TB, db database.Store, authToken uuid.UUID, mw func(http.Ha
WorkspaceID: workspace.ID,
JobID: job.ID,
TemplateVersionID: templateVersion.ID,
BuildNumber: 1,
Transition: database.WorkspaceTransitionStart,
})
_ = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
ResourceID: resource.ID,