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:
Ethan
2024-11-01 14:53:53 +11:00
committed by GitHub
parent f941e78079
commit b1298a3c1e
25 changed files with 2247 additions and 298 deletions
+190
View File
@@ -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)
}
}