mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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:
@@ -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{
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user