Skip to content

Commit 665b84d

Browse files
deansheatherdeansheather
authored
feat: use app tickets for web terminal (coder#6628)
1 parent a07209e commit 665b84d

18 files changed

+1009
-401
lines changed

‎coderd/coderd.go

+4-1
Original file line numberDiff line numberDiff line change
@@ -287,6 +287,7 @@ func New(options *Options) *API {
287287
options.Database,
288288
options.DeploymentValues,
289289
oauthConfigs,
290+
options.AgentInactiveDisconnectTimeout,
290291
options.AppSigningKey,
291292
),
292293
metricsCache: metricsCache,
@@ -618,14 +619,16 @@ func New(options *Options) *API {
618619
r.Post("/report-stats", api.workspaceAgentReportStats)
619620
r.Post("/report-lifecycle", api.workspaceAgentReportLifecycle)
620621
})
622+
// No middleware on the PTY endpoint since it uses workspace
623+
// application auth and tickets.
624+
r.Get("/{workspaceagent}/pty", api.workspaceAgentPTY)
621625
r.Route("/{workspaceagent}", func(r chi.Router) {
622626
r.Use(
623627
apiKeyMiddleware,
624628
httpmw.ExtractWorkspaceAgentParam(options.Database),
625629
httpmw.ExtractWorkspaceParam(options.Database),
626630
)
627631
r.Get("/", api.workspaceAgent)
628-
r.Get("/pty", api.workspaceAgentPTY)
629632
r.Get("/startup-logs", api.workspaceAgentStartupLogs)
630633
r.Get("/listening-ports", api.workspaceAgentListeningPorts)
631634
r.Get("/connection", api.workspaceAgentConnection)

‎coderd/coderdtest/coderdtest.go

+15-2
Original file line numberDiff line numberDiff line change
@@ -638,10 +638,17 @@ func AwaitWorkspaceBuildJob(t *testing.T, client *codersdk.Client, build uuid.UU
638638
return workspaceBuild
639639
}
640640

641-
// AwaitWorkspaceAgents waits for all resources with agents to be connected.
642-
func AwaitWorkspaceAgents(t *testing.T, client *codersdk.Client, workspaceID uuid.UUID) []codersdk.WorkspaceResource {
641+
// AwaitWorkspaceAgents waits for all resources with agents to be connected. If
642+
// specific agents are provided, it will wait for those agents to be connected
643+
// but will not fail if other agents are not connected.
644+
func AwaitWorkspaceAgents(t *testing.T, client *codersdk.Client, workspaceID uuid.UUID, agentNames ...string) []codersdk.WorkspaceResource {
643645
t.Helper()
644646

647+
agentNamesMap := make(map[string]struct{}, len(agentNames))
648+
for _, name := range agentNames {
649+
agentNamesMap[name] = struct{}{}
650+
}
651+
645652
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
646653
defer cancel()
647654

@@ -659,6 +666,12 @@ func AwaitWorkspaceAgents(t *testing.T, client *codersdk.Client, workspaceID uui
659666

660667
for _, resource := range workspace.LatestBuild.Resources {
661668
for _, agent := range resource.Agents {
669+
if len(agentNames) > 0 {
670+
if _, ok := agentNamesMap[agent.Name]; !ok {
671+
continue
672+
}
673+
}
674+
662675
if agent.Status != codersdk.WorkspaceAgentConnected {
663676
t.Logf("agent %s not connected yet", agent.Name)
664677
return false

‎coderd/database/dbauthz/querier.go

+9
Original file line numberDiff line numberDiff line change
@@ -1292,6 +1292,15 @@ func (q *querier) GetWorkspaceAgentByInstanceID(ctx context.Context, authInstanc
12921292
return agent, nil
12931293
}
12941294

1295+
func (q *querier) GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx context.Context, workspaceID uuid.UUID) ([]database.WorkspaceAgent, error) {
1296+
workspace, err := q.GetWorkspaceByID(ctx, workspaceID)
1297+
if err != nil {
1298+
return nil, err
1299+
}
1300+
1301+
return q.db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, workspace.ID)
1302+
}
1303+
12951304
func (q *querier) UpdateWorkspaceAgentLifecycleStateByID(ctx context.Context, arg database.UpdateWorkspaceAgentLifecycleStateByIDParams) error {
12961305
agent, err := q.db.GetWorkspaceAgentByID(ctx, arg.ID)
12971306
if err != nil {

‎coderd/database/dbfake/databasefake.go

+32
Original file line numberDiff line numberDiff line change
@@ -314,6 +314,38 @@ func (*fakeQuerier) DeleteOldWorkspaceAgentStats(_ context.Context) error {
314314
return nil
315315
}
316316

317+
func (q *fakeQuerier) GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx context.Context, workspaceID uuid.UUID) ([]database.WorkspaceAgent, error) {
318+
q.mutex.RLock()
319+
defer q.mutex.RUnlock()
320+
321+
// Get latest build for workspace.
322+
workspaceBuild, err := q.getLatestWorkspaceBuildByWorkspaceIDNoLock(ctx, workspaceID)
323+
if err != nil {
324+
return nil, xerrors.Errorf("get latest workspace build: %w", err)
325+
}
326+
327+
// Get resources for build.
328+
resources, err := q.GetWorkspaceResourcesByJobID(ctx, workspaceBuild.JobID)
329+
if err != nil {
330+
return nil, xerrors.Errorf("get workspace resources: %w", err)
331+
}
332+
if len(resources) == 0 {
333+
return []database.WorkspaceAgent{}, nil
334+
}
335+
336+
resourceIDs := make([]uuid.UUID, len(resources))
337+
for i, resource := range resources {
338+
resourceIDs[i] = resource.ID
339+
}
340+
341+
agents, err := q.GetWorkspaceAgentsByResourceIDs(ctx, resourceIDs)
342+
if err != nil {
343+
return nil, xerrors.Errorf("get workspace agents: %w", err)
344+
}
345+
346+
return agents, nil
347+
}
348+
317349
func (q *fakeQuerier) GetDeploymentWorkspaceAgentStats(_ context.Context, createdAfter time.Time) (database.GetDeploymentWorkspaceAgentStatsRow, error) {
318350
q.mutex.RLock()
319351
defer q.mutex.RUnlock()

‎coderd/database/modelmethods.go

+76
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package database
33
import (
44
"sort"
55
"strconv"
6+
"time"
67

78
"golang.org/x/exp/maps"
89

@@ -36,6 +37,26 @@ func (s WorkspaceStatus) Valid() bool {
3637
}
3738
}
3839

40+
type WorkspaceAgentStatus string
41+
42+
// This is also in codersdk/workspaceagents.go and should be kept in sync.
43+
const (
44+
WorkspaceAgentStatusConnecting WorkspaceAgentStatus = "connecting"
45+
WorkspaceAgentStatusConnected WorkspaceAgentStatus = "connected"
46+
WorkspaceAgentStatusDisconnected WorkspaceAgentStatus = "disconnected"
47+
WorkspaceAgentStatusTimeout WorkspaceAgentStatus = "timeout"
48+
)
49+
50+
func (s WorkspaceAgentStatus) Valid() bool {
51+
switch s {
52+
case WorkspaceAgentStatusConnecting, WorkspaceAgentStatusConnected,
53+
WorkspaceAgentStatusDisconnected, WorkspaceAgentStatusTimeout:
54+
return true
55+
default:
56+
return false
57+
}
58+
}
59+
3960
type AuditableGroup struct {
4061
Group
4162
Members []GroupMember `json:"members"`
@@ -199,6 +220,61 @@ func (l License) RBACObject() rbac.Object {
199220
return rbac.ResourceLicense.WithIDString(strconv.FormatInt(int64(l.ID), 10))
200221
}
201222

223+
type WorkspaceAgentConnectionStatus struct {
224+
Status WorkspaceAgentStatus `json:"status"`
225+
FirstConnectedAt *time.Time `json:"first_connected_at"`
226+
LastConnectedAt *time.Time `json:"last_connected_at"`
227+
DisconnectedAt *time.Time `json:"disconnected_at"`
228+
}
229+
230+
func (a WorkspaceAgent) Status(inactiveTimeout time.Duration) WorkspaceAgentConnectionStatus {
231+
connectionTimeout := time.Duration(a.ConnectionTimeoutSeconds) * time.Second
232+
233+
status := WorkspaceAgentConnectionStatus{
234+
Status: WorkspaceAgentStatusDisconnected,
235+
}
236+
if a.FirstConnectedAt.Valid {
237+
status.FirstConnectedAt = &a.FirstConnectedAt.Time
238+
}
239+
if a.LastConnectedAt.Valid {
240+
status.LastConnectedAt = &a.LastConnectedAt.Time
241+
}
242+
if a.DisconnectedAt.Valid {
243+
status.DisconnectedAt = &a.DisconnectedAt.Time
244+
}
245+
246+
switch {
247+
case !a.FirstConnectedAt.Valid:
248+
switch {
249+
case connectionTimeout > 0 && Now().Sub(a.CreatedAt) > connectionTimeout:
250+
// If the agent took too long to connect the first time,
251+
// mark it as timed out.
252+
status.Status = WorkspaceAgentStatusTimeout
253+
default:
254+
// If the agent never connected, it's waiting for the compute
255+
// to start up.
256+
status.Status = WorkspaceAgentStatusConnecting
257+
}
258+
// We check before instead of after because last connected at and
259+
// disconnected at can be equal timestamps in tight-timed tests.
260+
case !a.DisconnectedAt.Time.Before(a.LastConnectedAt.Time):
261+
// If we've disconnected after our last connection, we know the
262+
// agent is no longer connected.
263+
status.Status = WorkspaceAgentStatusDisconnected
264+
case Now().Sub(a.LastConnectedAt.Time) > inactiveTimeout:
265+
// The connection died without updating the last connected.
266+
status.Status = WorkspaceAgentStatusDisconnected
267+
// Client code needs an accurate disconnected at if the agent has been inactive.
268+
status.DisconnectedAt = &a.LastConnectedAt.Time
269+
case a.LastConnectedAt.Valid:
270+
// The agent should be assumed connected if it's under inactivity timeouts
271+
// and last connected at has been properly set.
272+
status.Status = WorkspaceAgentStatusConnected
273+
}
274+
275+
return status
276+
}
277+
202278
func ConvertUserRows(rows []GetUsersRow) []User {
203279
users := make([]User, len(rows))
204280
for i, r := range rows {

‎coderd/database/querier.go

+1
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries.sql.go

+75
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/workspaceagents.sql

+20
Original file line numberDiff line numberDiff line change
@@ -132,3 +132,23 @@ INSERT INTO
132132
DELETE FROM workspace_agent_startup_logs WHERE agent_id IN
133133
(SELECT id FROM workspace_agents WHERE last_connected_at IS NOT NULL
134134
AND last_connected_at < NOW() - INTERVAL '7 day');
135+
136+
-- name: GetWorkspaceAgentsInLatestBuildByWorkspaceID :many
137+
SELECT
138+
workspace_agents.*
139+
FROM
140+
workspace_agents
141+
JOIN
142+
workspace_resources ON workspace_agents.resource_id = workspace_resources.id
143+
JOIN
144+
workspace_builds ON workspace_resources.job_id = workspace_builds.job_id
145+
WHERE
146+
workspace_builds.workspace_id = @workspace_id :: uuid AND
147+
workspace_builds.build_number = (
148+
SELECT
149+
MAX(build_number)
150+
FROM
151+
workspace_builds AS wb
152+
WHERE
153+
wb.workspace_id = @workspace_id :: uuid
154+
);

0 commit comments

Comments
 (0)