fix: ensure agent token is from latest build in middleware (#12443)

This commit is contained in:
Garrett Delfosse
2024-03-14 12:27:32 -04:00
committed by GitHub
parent 63696d762f
commit 0723dd3abf
15 changed files with 242 additions and 260 deletions

View File

@ -61,11 +61,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
api.WebsocketWaitMutex.Unlock()
defer api.WebsocketWaitGroup.Done()
workspaceAgent := httpmw.WorkspaceAgent(r)
build, ok := ensureLatestBuild(ctx, api.Database, logger, rw, workspaceAgent)
if !ok {
return
}
build := httpmw.LatestBuild(r)
workspace, err := api.Database.GetWorkspaceByID(ctx, build.WorkspaceID)
if err != nil {
@ -167,54 +163,6 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
}
}
func ensureLatestBuild(ctx context.Context, db database.Store, logger slog.Logger, rw http.ResponseWriter, workspaceAgent database.WorkspaceAgent) (database.WorkspaceBuild, bool) {
resource, err := db.GetWorkspaceResourceByID(ctx, workspaceAgent.ResourceID)
if err != nil {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Internal error fetching workspace agent resource.",
Detail: err.Error(),
})
return database.WorkspaceBuild{}, false
}
build, err := db.GetWorkspaceBuildByJobID(ctx, resource.JobID)
if err != nil {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Internal error fetching workspace build job.",
Detail: err.Error(),
})
return database.WorkspaceBuild{}, false
}
// Ensure the resource is still valid!
// We only accept agents for resources on the latest build.
err = checkBuildIsLatest(ctx, db, build)
if err != nil {
logger.Debug(ctx, "agent tried to connect from non-latest build",
slog.F("resource", resource),
slog.F("agent", workspaceAgent),
)
httpapi.Write(ctx, rw, http.StatusForbidden, codersdk.Response{
Message: "Agent trying to connect from non-latest build.",
Detail: err.Error(),
})
return database.WorkspaceBuild{}, false
}
return build, true
}
func checkBuildIsLatest(ctx context.Context, db database.Store, build database.WorkspaceBuild) error {
latestBuild, err := db.GetLatestWorkspaceBuildByWorkspaceID(ctx, build.WorkspaceID)
if err != nil {
return err
}
if build.ID != latestBuild.ID {
return xerrors.New("build is outdated")
}
return nil
}
func (api *API) startAgentWebsocketMonitor(ctx context.Context,
workspaceAgent database.WorkspaceAgent, workspaceBuild database.WorkspaceBuild,
conn *websocket.Conn,
@ -494,3 +442,14 @@ func (m *agentConnectionMonitor) close() {
m.cancel()
m.wg.Wait()
}
func checkBuildIsLatest(ctx context.Context, db database.Store, build database.WorkspaceBuild) error {
latestBuild, err := db.GetLatestWorkspaceBuildByWorkspaceID(ctx, build.WorkspaceID)
if err != nil {
return err
}
if build.ID != latestBuild.ID {
return xerrors.New("build is outdated")
}
return nil
}