mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: add WorkspaceUpdates tailnet RPC (#14847)
Closes #14716 Closes #14717 Adds a new user-scoped tailnet API endpoint (`api/v2/tailnet`) with a new RPC stream for receiving updates on workspaces owned by a specific user, as defined in #14716. When a stream is started, the `WorkspaceUpdatesProvider` will begin listening on the user-scoped pubsub events implemented in #14964. When a relevant event type is seen (such as a workspace state transition), the provider will query the DB for all the workspaces (and agents) owned by the user. This gets compared against the result of the previous query to produce a set of workspace updates. Workspace updates can be requested for any user ID, however only workspaces the authorised user is permitted to `ActionRead` will have their updates streamed. Opening a tunnel to an agent requires that the user can perform `ActionSSH` against the workspace containing it.
This commit is contained in:
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"maps"
|
||||
"net"
|
||||
"net/http"
|
||||
"runtime"
|
||||
@@ -38,6 +39,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/database/pubsub"
|
||||
"github.com/coder/coder/v2/coderd/externalauth"
|
||||
"github.com/coder/coder/v2/coderd/jwtutils"
|
||||
"github.com/coder/coder/v2/coderd/rbac"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/codersdk/agentsdk"
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk"
|
||||
@@ -1930,6 +1932,106 @@ func TestWorkspaceAgentExternalAuthListen(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestOwnedWorkspacesCoordinate(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
|
||||
firstClient, _, api := coderdtest.NewWithAPI(t, &coderdtest.Options{
|
||||
Coordinator: tailnet.NewCoordinator(logger),
|
||||
})
|
||||
firstUser := coderdtest.CreateFirstUser(t, firstClient)
|
||||
member, memberUser := coderdtest.CreateAnotherUser(t, firstClient, firstUser.OrganizationID, rbac.RoleTemplateAdmin())
|
||||
|
||||
// Create a workspace with an agent
|
||||
firstWorkspace := buildWorkspaceWithAgent(t, member, firstUser.OrganizationID, memberUser.ID, api.Database, api.Pubsub)
|
||||
|
||||
u, err := member.URL.Parse("/api/v2/tailnet")
|
||||
require.NoError(t, err)
|
||||
q := u.Query()
|
||||
q.Set("version", "2.0")
|
||||
u.RawQuery = q.Encode()
|
||||
|
||||
//nolint:bodyclose // websocket package closes this for you
|
||||
wsConn, resp, err := websocket.Dial(ctx, u.String(), &websocket.DialOptions{
|
||||
HTTPHeader: http.Header{
|
||||
"Coder-Session-Token": []string{member.SessionToken()},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
if resp.StatusCode != http.StatusSwitchingProtocols {
|
||||
err = codersdk.ReadBodyAsError(resp)
|
||||
}
|
||||
require.NoError(t, err)
|
||||
}
|
||||
defer wsConn.Close(websocket.StatusNormalClosure, "done")
|
||||
|
||||
rpcClient, err := tailnet.NewDRPCClient(
|
||||
websocket.NetConn(ctx, wsConn, websocket.MessageBinary),
|
||||
logger,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
stream, err := rpcClient.WorkspaceUpdates(ctx, &tailnetproto.WorkspaceUpdatesRequest{
|
||||
WorkspaceOwnerId: tailnet.UUIDToByteSlice(memberUser.ID),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// First update will contain the existing workspace and agent
|
||||
update, err := stream.Recv()
|
||||
require.NoError(t, err)
|
||||
require.Len(t, update.UpsertedWorkspaces, 1)
|
||||
require.EqualValues(t, update.UpsertedWorkspaces[0].Id, firstWorkspace.ID)
|
||||
require.Len(t, update.UpsertedAgents, 1)
|
||||
require.EqualValues(t, update.UpsertedAgents[0].WorkspaceId, firstWorkspace.ID)
|
||||
require.Len(t, update.DeletedWorkspaces, 0)
|
||||
require.Len(t, update.DeletedAgents, 0)
|
||||
|
||||
// Build a second workspace
|
||||
secondWorkspace := buildWorkspaceWithAgent(t, member, firstUser.OrganizationID, memberUser.ID, api.Database, api.Pubsub)
|
||||
|
||||
// Wait for the second workspace to be running with an agent
|
||||
expectedState := map[uuid.UUID]workspace{
|
||||
secondWorkspace.ID: {
|
||||
Status: tailnetproto.Workspace_RUNNING,
|
||||
NumAgents: 1,
|
||||
},
|
||||
}
|
||||
waitForUpdates(t, ctx, stream, map[uuid.UUID]workspace{}, expectedState)
|
||||
|
||||
// Wait for the workspace and agent to be deleted
|
||||
secondWorkspace.Deleted = true
|
||||
dbfake.WorkspaceBuild(t, api.Database, secondWorkspace).
|
||||
Seed(database.WorkspaceBuild{
|
||||
Transition: database.WorkspaceTransitionDelete,
|
||||
BuildNumber: 2,
|
||||
}).Do()
|
||||
|
||||
waitForUpdates(t, ctx, stream, expectedState, map[uuid.UUID]workspace{
|
||||
secondWorkspace.ID: {
|
||||
Status: tailnetproto.Workspace_DELETED,
|
||||
NumAgents: 0,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func buildWorkspaceWithAgent(
|
||||
t *testing.T,
|
||||
client *codersdk.Client,
|
||||
orgID uuid.UUID,
|
||||
ownerID uuid.UUID,
|
||||
db database.Store,
|
||||
ps pubsub.Pubsub,
|
||||
) database.WorkspaceTable {
|
||||
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
|
||||
OrganizationID: orgID,
|
||||
OwnerID: ownerID,
|
||||
}).WithAgent().Pubsub(ps).Do()
|
||||
_ = agenttest.New(t, client.URL, r.AgentToken)
|
||||
coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).Wait()
|
||||
return r.Workspace
|
||||
}
|
||||
|
||||
func requireGetManifest(ctx context.Context, t testing.TB, aAPI agentproto.DRPCAgentClient) agentsdk.Manifest {
|
||||
mp, err := aAPI.GetManifest(ctx, &agentproto.GetManifestRequest{})
|
||||
require.NoError(t, err)
|
||||
@@ -1949,3 +2051,91 @@ func postStartup(ctx context.Context, t testing.TB, client agent.Client, startup
|
||||
_, err = aAPI.UpdateStartup(ctx, &agentproto.UpdateStartupRequest{Startup: startup})
|
||||
return err
|
||||
}
|
||||
|
||||
type workspace struct {
|
||||
Status tailnetproto.Workspace_Status
|
||||
NumAgents int
|
||||
}
|
||||
|
||||
func waitForUpdates(
|
||||
t *testing.T,
|
||||
//nolint:revive // t takes precedence
|
||||
ctx context.Context,
|
||||
stream tailnetproto.DRPCTailnet_WorkspaceUpdatesClient,
|
||||
currentState map[uuid.UUID]workspace,
|
||||
expectedState map[uuid.UUID]workspace,
|
||||
) {
|
||||
t.Helper()
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
errCh <- ctx.Err()
|
||||
return
|
||||
default:
|
||||
}
|
||||
update, err := stream.Recv()
|
||||
if err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
for _, ws := range update.UpsertedWorkspaces {
|
||||
id, err := uuid.FromBytes(ws.Id)
|
||||
if err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
currentState[id] = workspace{
|
||||
Status: ws.Status,
|
||||
NumAgents: currentState[id].NumAgents,
|
||||
}
|
||||
}
|
||||
for _, ws := range update.DeletedWorkspaces {
|
||||
id, err := uuid.FromBytes(ws.Id)
|
||||
if err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
currentState[id] = workspace{
|
||||
Status: tailnetproto.Workspace_DELETED,
|
||||
NumAgents: currentState[id].NumAgents,
|
||||
}
|
||||
}
|
||||
for _, a := range update.UpsertedAgents {
|
||||
id, err := uuid.FromBytes(a.WorkspaceId)
|
||||
if err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
currentState[id] = workspace{
|
||||
Status: currentState[id].Status,
|
||||
NumAgents: currentState[id].NumAgents + 1,
|
||||
}
|
||||
}
|
||||
for _, a := range update.DeletedAgents {
|
||||
id, err := uuid.FromBytes(a.WorkspaceId)
|
||||
if err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
currentState[id] = workspace{
|
||||
Status: currentState[id].Status,
|
||||
NumAgents: currentState[id].NumAgents - 1,
|
||||
}
|
||||
}
|
||||
if maps.Equal(currentState, expectedState) {
|
||||
errCh <- nil
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
select {
|
||||
case err := <-errCh:
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
case <-ctx.Done():
|
||||
t.Fatal("Timeout waiting for desired state", currentState)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user