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
+19
View File
@@ -3770,6 +3770,25 @@ const docTemplate = `{
}
}
},
"/tailnet": {
"get": {
"security": [
{
"CoderSessionToken": []
}
],
"tags": [
"Agents"
],
"summary": "User-scoped tailnet RPC connection",
"operationId": "user-scoped-tailnet-rpc-connection",
"responses": {
"101": {
"description": "Switching Protocols"
}
}
}
},
"/templates": {
"get": {
"security": [
+17
View File
@@ -3316,6 +3316,23 @@
}
}
},
"/tailnet": {
"get": {
"security": [
{
"CoderSessionToken": []
}
],
"tags": ["Agents"],
"summary": "User-scoped tailnet RPC connection",
"operationId": "user-scoped-tailnet-rpc-connection",
"responses": {
"101": {
"description": "Switching Protocols"
}
}
}
},
"/templates": {
"get": {
"security": [
+17 -6
View File
@@ -493,6 +493,8 @@ func New(options *Options) *API {
}
}
updatesProvider := NewUpdatesProvider(options.Logger.Named("workspace_updates"), options.Pubsub, options.Database, options.Authorizer)
// Start a background process that rotates keys. We intentionally start this after the caches
// are created to force initial requests for a key to populate the caches. This helps catch
// bugs that may only occur when a key isn't precached in tests and the latency cost is minimal.
@@ -523,6 +525,7 @@ func New(options *Options) *API {
metricsCache: metricsCache,
Auditor: atomic.Pointer[audit.Auditor]{},
TailnetCoordinator: atomic.Pointer[tailnet.Coordinator]{},
UpdatesProvider: updatesProvider,
TemplateScheduleStore: options.TemplateScheduleStore,
UserQuietHoursScheduleStore: options.UserQuietHoursScheduleStore,
AccessControlStore: options.AccessControlStore,
@@ -652,12 +655,13 @@ func New(options *Options) *API {
panic("CoordinatorResumeTokenProvider is nil")
}
api.TailnetClientService, err = tailnet.NewClientService(tailnet.ClientServiceOptions{
Logger: api.Logger.Named("tailnetclient"),
CoordPtr: &api.TailnetCoordinator,
DERPMapUpdateFrequency: api.Options.DERPMapUpdateFrequency,
DERPMapFn: api.DERPMap,
NetworkTelemetryHandler: api.NetworkTelemetryBatcher.Handler,
ResumeTokenProvider: api.Options.CoordinatorResumeTokenProvider,
Logger: api.Logger.Named("tailnetclient"),
CoordPtr: &api.TailnetCoordinator,
DERPMapUpdateFrequency: api.Options.DERPMapUpdateFrequency,
DERPMapFn: api.DERPMap,
NetworkTelemetryHandler: api.NetworkTelemetryBatcher.Handler,
ResumeTokenProvider: api.Options.CoordinatorResumeTokenProvider,
WorkspaceUpdatesProvider: api.UpdatesProvider,
})
if err != nil {
api.Logger.Fatal(context.Background(), "failed to initialize tailnet client service", slog.Error(err))
@@ -1327,6 +1331,10 @@ func New(options *Options) *API {
})
r.Get("/dispatch-methods", api.notificationDispatchMethods)
})
r.Route("/tailnet", func(r chi.Router) {
r.Use(apiKeyMiddleware)
r.Get("/", api.tailnetRPCConn)
})
})
if options.SwaggerEndpoint {
@@ -1408,6 +1416,8 @@ type API struct {
AccessControlStore *atomic.Pointer[dbauthz.AccessControlStore]
PortSharer atomic.Pointer[portsharing.PortSharer]
UpdatesProvider tailnet.WorkspaceUpdatesProvider
HTTPAuth *HTTPAuthorizer
// APIHandler serves "/api/v2"
@@ -1489,6 +1499,7 @@ func (api *API) Close() error {
_ = api.OIDCConvertKeyCache.Close()
_ = api.AppSigningKeyCache.Close()
_ = api.AppEncryptionKeyCache.Close()
_ = api.UpdatesProvider.Close()
return nil
}
+8
View File
@@ -224,6 +224,14 @@ func (b WorkspaceBuildBuilder) Do() WorkspaceResponse {
}
_ = dbgen.WorkspaceBuildParameters(b.t, b.db, b.params)
if b.ws.Deleted {
err = b.db.UpdateWorkspaceDeletedByID(ownerCtx, database.UpdateWorkspaceDeletedByIDParams{
ID: b.ws.ID,
Deleted: true,
})
require.NoError(b.t, err)
}
if b.ps != nil {
msg, err := json.Marshal(wspubsub.WorkspaceEvent{
Kind: wspubsub.WorkspaceEventKindStateChange,
+114 -26
View File
@@ -33,6 +33,7 @@ import (
"github.com/coder/coder/v2/coderd/httpapi"
"github.com/coder/coder/v2/coderd/httpmw"
"github.com/coder/coder/v2/coderd/jwtutils"
"github.com/coder/coder/v2/coderd/rbac"
"github.com/coder/coder/v2/coderd/rbac/policy"
"github.com/coder/coder/v2/coderd/wspubsub"
"github.com/coder/coder/v2/codersdk"
@@ -844,31 +845,10 @@ func (api *API) workspaceAgentClientCoordinate(rw http.ResponseWriter, r *http.R
return
}
// Accept a resume_token query parameter to use the same peer ID.
var (
peerID = uuid.New()
resumeToken = r.URL.Query().Get("resume_token")
)
if resumeToken != "" {
var err error
peerID, err = api.Options.CoordinatorResumeTokenProvider.VerifyResumeToken(ctx, resumeToken)
// If the token is missing the key ID, it's probably an old token in which
// case we just want to generate a new peer ID.
if xerrors.Is(err, jwtutils.ErrMissingKeyID) {
peerID = uuid.New()
} else if err != nil {
httpapi.Write(ctx, rw, http.StatusUnauthorized, codersdk.Response{
Message: workspacesdk.CoordinateAPIInvalidResumeToken,
Detail: err.Error(),
Validations: []codersdk.ValidationError{
{Field: "resume_token", Detail: workspacesdk.CoordinateAPIInvalidResumeToken},
},
})
return
} else {
api.Logger.Debug(ctx, "accepted coordinate resume token for peer",
slog.F("peer_id", peerID.String()))
}
peerID, err := api.handleResumeToken(ctx, rw, r)
if err != nil {
// handleResumeToken has already written the response.
return
}
api.WebsocketWaitMutex.Lock()
@@ -891,13 +871,47 @@ func (api *API) workspaceAgentClientCoordinate(rw http.ResponseWriter, r *http.R
go httpapi.Heartbeat(ctx, conn)
defer conn.Close(websocket.StatusNormalClosure, "")
err = api.TailnetClientService.ServeClient(ctx, version, wsNetConn, peerID, workspaceAgent.ID)
err = api.TailnetClientService.ServeClient(ctx, version, wsNetConn, tailnet.StreamID{
Name: "client",
ID: peerID,
Auth: tailnet.ClientCoordinateeAuth{
AgentID: workspaceAgent.ID,
},
})
if err != nil && !xerrors.Is(err, io.EOF) && !xerrors.Is(err, context.Canceled) {
_ = conn.Close(websocket.StatusInternalError, err.Error())
return
}
}
// handleResumeToken accepts a resume_token query parameter to use the same peer ID
func (api *API) handleResumeToken(ctx context.Context, rw http.ResponseWriter, r *http.Request) (peerID uuid.UUID, err error) {
peerID = uuid.New()
resumeToken := r.URL.Query().Get("resume_token")
if resumeToken != "" {
peerID, err = api.Options.CoordinatorResumeTokenProvider.VerifyResumeToken(ctx, resumeToken)
// If the token is missing the key ID, it's probably an old token in which
// case we just want to generate a new peer ID.
if xerrors.Is(err, jwtutils.ErrMissingKeyID) {
peerID = uuid.New()
err = nil
} else if err != nil {
httpapi.Write(ctx, rw, http.StatusUnauthorized, codersdk.Response{
Message: workspacesdk.CoordinateAPIInvalidResumeToken,
Detail: err.Error(),
Validations: []codersdk.ValidationError{
{Field: "resume_token", Detail: workspacesdk.CoordinateAPIInvalidResumeToken},
},
})
return peerID, err
} else {
api.Logger.Debug(ctx, "accepted coordinate resume token for peer",
slog.F("peer_id", peerID.String()))
}
}
return peerID, err
}
// @Summary Post workspace agent log source
// @ID post-workspace-agent-log-source
// @Security CoderSessionToken
@@ -1469,6 +1483,80 @@ func (api *API) workspaceAgentsExternalAuthListen(ctx context.Context, rw http.R
}
}
// @Summary User-scoped tailnet RPC connection
// @ID user-scoped-tailnet-rpc-connection
// @Security CoderSessionToken
// @Tags Agents
// @Success 101
// @Router /tailnet [get]
func (api *API) tailnetRPCConn(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
version := "2.0"
qv := r.URL.Query().Get("version")
if qv != "" {
version = qv
}
if err := proto.CurrentVersion.Validate(version); err != nil {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Unknown or unsupported API version",
Validations: []codersdk.ValidationError{
{Field: "version", Detail: err.Error()},
},
})
return
}
peerID, err := api.handleResumeToken(ctx, rw, r)
if err != nil {
// handleResumeToken has already written the response.
return
}
// Used to authorize tunnel request
sshPrep, err := api.HTTPAuth.AuthorizeSQLFilter(r, policy.ActionSSH, rbac.ResourceWorkspace.Type)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error preparing sql filter.",
Detail: err.Error(),
})
return
}
api.WebsocketWaitMutex.Lock()
api.WebsocketWaitGroup.Add(1)
api.WebsocketWaitMutex.Unlock()
defer api.WebsocketWaitGroup.Done()
conn, err := websocket.Accept(rw, r, nil)
if err != nil {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Failed to accept websocket.",
Detail: err.Error(),
})
return
}
ctx, wsNetConn := codersdk.WebsocketNetConn(ctx, conn, websocket.MessageBinary)
defer wsNetConn.Close()
defer conn.Close(websocket.StatusNormalClosure, "")
go httpapi.Heartbeat(ctx, conn)
err = api.TailnetClientService.ServeClient(ctx, version, wsNetConn, tailnet.StreamID{
Name: "client",
ID: peerID,
Auth: tailnet.ClientUserCoordinateeAuth{
Auth: &rbacAuthorizer{
sshPrep: sshPrep,
db: api.Database,
},
},
})
if err != nil && !xerrors.Is(err, io.EOF) && !xerrors.Is(err, context.Canceled) {
_ = conn.Close(websocket.StatusInternalError, err.Error())
return
}
}
// createExternalAuthResponse creates an ExternalAuthResponse based on the
// provider type. This is to support legacy `/workspaceagents/me/gitauth`
// which uses `Username` and `Password`.
+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)
}
}
+1 -35
View File
@@ -916,7 +916,7 @@ func (api *API) convertWorkspaceBuild(
MaxDeadline: codersdk.NewNullTime(build.MaxDeadline, !build.MaxDeadline.IsZero()),
Reason: codersdk.BuildReason(build.Reason),
Resources: apiResources,
Status: convertWorkspaceStatus(apiJob.Status, transition),
Status: codersdk.ConvertWorkspaceStatus(apiJob.Status, transition),
DailyCost: build.DailyCost,
}, nil
}
@@ -946,40 +946,6 @@ func convertWorkspaceResource(resource database.WorkspaceResource, agents []code
}
}
func convertWorkspaceStatus(jobStatus codersdk.ProvisionerJobStatus, transition codersdk.WorkspaceTransition) codersdk.WorkspaceStatus {
switch jobStatus {
case codersdk.ProvisionerJobPending:
return codersdk.WorkspaceStatusPending
case codersdk.ProvisionerJobRunning:
switch transition {
case codersdk.WorkspaceTransitionStart:
return codersdk.WorkspaceStatusStarting
case codersdk.WorkspaceTransitionStop:
return codersdk.WorkspaceStatusStopping
case codersdk.WorkspaceTransitionDelete:
return codersdk.WorkspaceStatusDeleting
}
case codersdk.ProvisionerJobSucceeded:
switch transition {
case codersdk.WorkspaceTransitionStart:
return codersdk.WorkspaceStatusRunning
case codersdk.WorkspaceTransitionStop:
return codersdk.WorkspaceStatusStopped
case codersdk.WorkspaceTransitionDelete:
return codersdk.WorkspaceStatusDeleted
}
case codersdk.ProvisionerJobCanceling:
return codersdk.WorkspaceStatusCanceling
case codersdk.ProvisionerJobCanceled:
return codersdk.WorkspaceStatusCanceled
case codersdk.ProvisionerJobFailed:
return codersdk.WorkspaceStatusFailed
}
// return error status since we should never get here
return codersdk.WorkspaceStatusFailed
}
func (api *API) buildTimings(ctx context.Context, build database.WorkspaceBuild) (codersdk.WorkspaceBuildTimings, error) {
provisionerTimings, err := api.Database.GetProvisionerJobTimingsByJobID(ctx, build.JobID)
if err != nil && !errors.Is(err, sql.ErrNoRows) {
+313
View File
@@ -0,0 +1,313 @@
package coderd
import (
"context"
"fmt"
"sync"
"github.com/google/uuid"
"golang.org/x/xerrors"
"cdr.dev/slog"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/database/pubsub"
"github.com/coder/coder/v2/coderd/rbac"
"github.com/coder/coder/v2/coderd/util/slice"
"github.com/coder/coder/v2/coderd/wspubsub"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/tailnet"
"github.com/coder/coder/v2/tailnet/proto"
)
type UpdatesQuerier interface {
// GetAuthorizedWorkspacesAndAgentsByOwnerID requires a context with an actor set
GetWorkspacesAndAgentsByOwnerID(ctx context.Context, ownerID uuid.UUID) ([]database.GetWorkspacesAndAgentsByOwnerIDRow, error)
GetWorkspaceByAgentID(ctx context.Context, agentID uuid.UUID) (database.Workspace, error)
}
type workspacesByID = map[uuid.UUID]ownedWorkspace
type ownedWorkspace struct {
WorkspaceName string
Status proto.Workspace_Status
Agents []database.AgentIDNamePair
}
// Equal does not compare agents
func (w ownedWorkspace) Equal(other ownedWorkspace) bool {
return w.WorkspaceName == other.WorkspaceName &&
w.Status == other.Status
}
type sub struct {
// ALways contains an actor
ctx context.Context
cancelFn context.CancelFunc
mu sync.RWMutex
userID uuid.UUID
ch chan *proto.WorkspaceUpdate
prev workspacesByID
db UpdatesQuerier
ps pubsub.Pubsub
logger slog.Logger
psCancelFn func()
}
func (s *sub) handleEvent(ctx context.Context, event wspubsub.WorkspaceEvent, err error) {
s.mu.Lock()
defer s.mu.Unlock()
switch event.Kind {
case wspubsub.WorkspaceEventKindStateChange:
case wspubsub.WorkspaceEventKindAgentConnectionUpdate:
case wspubsub.WorkspaceEventKindAgentTimeout:
case wspubsub.WorkspaceEventKindAgentLifecycleUpdate:
default:
if err == nil {
return
} else {
// Always attempt an update if the pubsub lost connection
s.logger.Warn(ctx, "failed to handle workspace event", slog.Error(err))
}
}
// Use context containing actor
rows, err := s.db.GetWorkspacesAndAgentsByOwnerID(s.ctx, s.userID)
if err != nil {
s.logger.Warn(ctx, "failed to get workspaces and agents by owner ID", slog.Error(err))
return
}
latest := convertRows(rows)
out, updated := produceUpdate(s.prev, latest)
if !updated {
return
}
s.prev = latest
select {
case <-s.ctx.Done():
return
case s.ch <- out:
}
}
func (s *sub) start(ctx context.Context) (err error) {
rows, err := s.db.GetWorkspacesAndAgentsByOwnerID(ctx, s.userID)
if err != nil {
return xerrors.Errorf("get workspaces and agents by owner ID: %w", err)
}
latest := convertRows(rows)
initUpdate, _ := produceUpdate(workspacesByID{}, latest)
s.ch <- initUpdate
s.prev = latest
cancel, err := s.ps.SubscribeWithErr(wspubsub.WorkspaceEventChannel(s.userID), wspubsub.HandleWorkspaceEvent(s.handleEvent))
if err != nil {
return xerrors.Errorf("subscribe to workspace event channel: %w", err)
}
s.psCancelFn = cancel
return nil
}
func (s *sub) Close() error {
s.cancelFn()
if s.psCancelFn != nil {
s.psCancelFn()
}
close(s.ch)
return nil
}
func (s *sub) Updates() <-chan *proto.WorkspaceUpdate {
return s.ch
}
var _ tailnet.Subscription = (*sub)(nil)
type updatesProvider struct {
ps pubsub.Pubsub
logger slog.Logger
db UpdatesQuerier
auth rbac.Authorizer
ctx context.Context
cancelFn func()
}
var _ tailnet.WorkspaceUpdatesProvider = (*updatesProvider)(nil)
func NewUpdatesProvider(
logger slog.Logger,
ps pubsub.Pubsub,
db UpdatesQuerier,
auth rbac.Authorizer,
) tailnet.WorkspaceUpdatesProvider {
ctx, cancel := context.WithCancel(context.Background())
out := &updatesProvider{
auth: auth,
db: db,
ps: ps,
logger: logger,
ctx: ctx,
cancelFn: cancel,
}
return out
}
func (u *updatesProvider) Close() error {
u.cancelFn()
return nil
}
// Subscribe subscribes to workspace updates for a user, for the workspaces
// that user is authorized to `ActionRead` on. The provided context must have
// a dbauthz actor set.
func (u *updatesProvider) Subscribe(ctx context.Context, userID uuid.UUID) (tailnet.Subscription, error) {
actor, ok := dbauthz.ActorFromContext(ctx)
if !ok {
return nil, xerrors.Errorf("actor not found in context")
}
ctx, cancel := context.WithCancel(u.ctx)
ctx = dbauthz.As(ctx, actor)
ch := make(chan *proto.WorkspaceUpdate, 1)
sub := &sub{
ctx: ctx,
cancelFn: cancel,
userID: userID,
ch: ch,
db: u.db,
ps: u.ps,
logger: u.logger.Named(fmt.Sprintf("workspace_updates_subscriber_%s", userID)),
prev: workspacesByID{},
}
err := sub.start(ctx)
if err != nil {
_ = sub.Close()
return nil, err
}
return sub, nil
}
func produceUpdate(old, new workspacesByID) (out *proto.WorkspaceUpdate, updated bool) {
out = &proto.WorkspaceUpdate{
UpsertedWorkspaces: []*proto.Workspace{},
UpsertedAgents: []*proto.Agent{},
DeletedWorkspaces: []*proto.Workspace{},
DeletedAgents: []*proto.Agent{},
}
for wsID, newWorkspace := range new {
oldWorkspace, exists := old[wsID]
// Upsert both workspace and agents if the workspace is new
if !exists {
out.UpsertedWorkspaces = append(out.UpsertedWorkspaces, &proto.Workspace{
Id: tailnet.UUIDToByteSlice(wsID),
Name: newWorkspace.WorkspaceName,
Status: newWorkspace.Status,
})
for _, agent := range newWorkspace.Agents {
out.UpsertedAgents = append(out.UpsertedAgents, &proto.Agent{
Id: tailnet.UUIDToByteSlice(agent.ID),
Name: agent.Name,
WorkspaceId: tailnet.UUIDToByteSlice(wsID),
})
}
updated = true
continue
}
// Upsert workspace if the workspace is updated
if !newWorkspace.Equal(oldWorkspace) {
out.UpsertedWorkspaces = append(out.UpsertedWorkspaces, &proto.Workspace{
Id: tailnet.UUIDToByteSlice(wsID),
Name: newWorkspace.WorkspaceName,
Status: newWorkspace.Status,
})
updated = true
}
add, remove := slice.SymmetricDifference(oldWorkspace.Agents, newWorkspace.Agents)
for _, agent := range add {
out.UpsertedAgents = append(out.UpsertedAgents, &proto.Agent{
Id: tailnet.UUIDToByteSlice(agent.ID),
Name: agent.Name,
WorkspaceId: tailnet.UUIDToByteSlice(wsID),
})
updated = true
}
for _, agent := range remove {
out.DeletedAgents = append(out.DeletedAgents, &proto.Agent{
Id: tailnet.UUIDToByteSlice(agent.ID),
Name: agent.Name,
WorkspaceId: tailnet.UUIDToByteSlice(wsID),
})
updated = true
}
}
// Delete workspace and agents if the workspace is deleted
for wsID, oldWorkspace := range old {
if _, exists := new[wsID]; !exists {
out.DeletedWorkspaces = append(out.DeletedWorkspaces, &proto.Workspace{
Id: tailnet.UUIDToByteSlice(wsID),
Name: oldWorkspace.WorkspaceName,
Status: oldWorkspace.Status,
})
for _, agent := range oldWorkspace.Agents {
out.DeletedAgents = append(out.DeletedAgents, &proto.Agent{
Id: tailnet.UUIDToByteSlice(agent.ID),
Name: agent.Name,
WorkspaceId: tailnet.UUIDToByteSlice(wsID),
})
}
updated = true
}
}
return out, updated
}
func convertRows(rows []database.GetWorkspacesAndAgentsByOwnerIDRow) workspacesByID {
out := workspacesByID{}
for _, row := range rows {
agents := []database.AgentIDNamePair{}
for _, agent := range row.Agents {
agents = append(agents, database.AgentIDNamePair{
ID: agent.ID,
Name: agent.Name,
})
}
out[row.ID] = ownedWorkspace{
WorkspaceName: row.Name,
Status: tailnet.WorkspaceStatusToProto(codersdk.ConvertWorkspaceStatus(codersdk.ProvisionerJobStatus(row.JobStatus), codersdk.WorkspaceTransition(row.Transition))),
Agents: agents,
}
}
return out
}
type rbacAuthorizer struct {
sshPrep rbac.PreparedAuthorized
db UpdatesQuerier
}
func (r *rbacAuthorizer) AuthorizeTunnel(ctx context.Context, agentID uuid.UUID) error {
ws, err := r.db.GetWorkspaceByAgentID(ctx, agentID)
if err != nil {
return xerrors.Errorf("get workspace by agent ID: %w", err)
}
// Authorizes against `ActionSSH`
return r.sshPrep.Authorize(ctx, ws.RBACObject())
}
var _ tailnet.TunnelAuthorizer = (*rbacAuthorizer)(nil)
+371
View File
@@ -0,0 +1,371 @@
package coderd_test
import (
"context"
"encoding/json"
"slices"
"strings"
"testing"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
"cdr.dev/slog/sloggers/slogtest"
"github.com/coder/coder/v2/coderd"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/database/pubsub"
"github.com/coder/coder/v2/coderd/rbac"
"github.com/coder/coder/v2/coderd/rbac/policy"
"github.com/coder/coder/v2/coderd/wspubsub"
"github.com/coder/coder/v2/tailnet"
"github.com/coder/coder/v2/tailnet/proto"
"github.com/coder/coder/v2/testutil"
)
func TestWorkspaceUpdates(t *testing.T) {
t.Parallel()
ws1ID := uuid.UUID{0x01}
ws1IDSlice := tailnet.UUIDToByteSlice(ws1ID)
agent1ID := uuid.UUID{0x02}
agent1IDSlice := tailnet.UUIDToByteSlice(agent1ID)
ws2ID := uuid.UUID{0x03}
ws2IDSlice := tailnet.UUIDToByteSlice(ws2ID)
ws3ID := uuid.UUID{0x04}
ws3IDSlice := tailnet.UUIDToByteSlice(ws3ID)
agent2ID := uuid.UUID{0x05}
agent2IDSlice := tailnet.UUIDToByteSlice(agent2ID)
ws4ID := uuid.UUID{0x06}
ws4IDSlice := tailnet.UUIDToByteSlice(ws4ID)
agent3ID := uuid.UUID{0x07}
agent3IDSlice := tailnet.UUIDToByteSlice(agent3ID)
ownerID := uuid.UUID{0x08}
memberRole, err := rbac.RoleByName(rbac.RoleMember())
require.NoError(t, err)
ownerSubject := rbac.Subject{
FriendlyName: "member",
ID: ownerID.String(),
Roles: rbac.Roles{memberRole},
Scope: rbac.ScopeAll,
}
t.Run("Basic", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
db := &mockWorkspaceStore{
orderedRows: []database.GetWorkspacesAndAgentsByOwnerIDRow{
// Gains agent2
{
ID: ws1ID,
Name: "ws1",
JobStatus: database.ProvisionerJobStatusRunning,
Transition: database.WorkspaceTransitionStart,
Agents: []database.AgentIDNamePair{
{
ID: agent1ID,
Name: "agent1",
},
},
},
// Changes status
{
ID: ws2ID,
Name: "ws2",
JobStatus: database.ProvisionerJobStatusRunning,
Transition: database.WorkspaceTransitionStart,
},
// Is deleted
{
ID: ws3ID,
Name: "ws3",
JobStatus: database.ProvisionerJobStatusSucceeded,
Transition: database.WorkspaceTransitionStop,
Agents: []database.AgentIDNamePair{
{
ID: agent3ID,
Name: "agent3",
},
},
},
},
}
ps := &mockPubsub{
cbs: map[string]pubsub.ListenerWithErr{},
}
updateProvider := coderd.NewUpdatesProvider(slogtest.Make(t, nil), ps, db, &mockAuthorizer{})
t.Cleanup(func() {
_ = updateProvider.Close()
})
sub, err := updateProvider.Subscribe(dbauthz.As(ctx, ownerSubject), ownerID)
require.NoError(t, err)
t.Cleanup(func() {
_ = sub.Close()
})
update := testutil.RequireRecvCtx(ctx, t, sub.Updates())
slices.SortFunc(update.UpsertedWorkspaces, func(a, b *proto.Workspace) int {
return strings.Compare(a.Name, b.Name)
})
slices.SortFunc(update.UpsertedAgents, func(a, b *proto.Agent) int {
return strings.Compare(a.Name, b.Name)
})
require.Equal(t, &proto.WorkspaceUpdate{
UpsertedWorkspaces: []*proto.Workspace{
{
Id: ws1IDSlice,
Name: "ws1",
Status: proto.Workspace_STARTING,
},
{
Id: ws2IDSlice,
Name: "ws2",
Status: proto.Workspace_STARTING,
},
{
Id: ws3IDSlice,
Name: "ws3",
Status: proto.Workspace_STOPPED,
},
},
UpsertedAgents: []*proto.Agent{
{
Id: agent1IDSlice,
Name: "agent1",
WorkspaceId: ws1IDSlice,
},
{
Id: agent3IDSlice,
Name: "agent3",
WorkspaceId: ws3IDSlice,
},
},
DeletedWorkspaces: []*proto.Workspace{},
DeletedAgents: []*proto.Agent{},
}, update)
// Update the database
db.orderedRows = []database.GetWorkspacesAndAgentsByOwnerIDRow{
{
ID: ws1ID,
Name: "ws1",
JobStatus: database.ProvisionerJobStatusRunning,
Transition: database.WorkspaceTransitionStart,
Agents: []database.AgentIDNamePair{
{
ID: agent1ID,
Name: "agent1",
},
{
ID: agent2ID,
Name: "agent2",
},
},
},
{
ID: ws2ID,
Name: "ws2",
JobStatus: database.ProvisionerJobStatusRunning,
Transition: database.WorkspaceTransitionStop,
},
{
ID: ws4ID,
Name: "ws4",
JobStatus: database.ProvisionerJobStatusRunning,
Transition: database.WorkspaceTransitionStart,
},
}
publishWorkspaceEvent(t, ps, ownerID, &wspubsub.WorkspaceEvent{
Kind: wspubsub.WorkspaceEventKindStateChange,
WorkspaceID: ws1ID,
})
update = testutil.RequireRecvCtx(ctx, t, sub.Updates())
slices.SortFunc(update.UpsertedWorkspaces, func(a, b *proto.Workspace) int {
return strings.Compare(a.Name, b.Name)
})
require.Equal(t, &proto.WorkspaceUpdate{
UpsertedWorkspaces: []*proto.Workspace{
{
// Changed status
Id: ws2IDSlice,
Name: "ws2",
Status: proto.Workspace_STOPPING,
},
{
// New workspace
Id: ws4IDSlice,
Name: "ws4",
Status: proto.Workspace_STARTING,
},
},
UpsertedAgents: []*proto.Agent{
{
Id: agent2IDSlice,
Name: "agent2",
WorkspaceId: ws1IDSlice,
},
},
DeletedWorkspaces: []*proto.Workspace{
{
Id: ws3IDSlice,
Name: "ws3",
Status: proto.Workspace_STOPPED,
},
},
DeletedAgents: []*proto.Agent{
{
Id: agent3IDSlice,
Name: "agent3",
WorkspaceId: ws3IDSlice,
},
},
}, update)
})
t.Run("Resubscribe", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
db := &mockWorkspaceStore{
orderedRows: []database.GetWorkspacesAndAgentsByOwnerIDRow{
{
ID: ws1ID,
Name: "ws1",
JobStatus: database.ProvisionerJobStatusRunning,
Transition: database.WorkspaceTransitionStart,
Agents: []database.AgentIDNamePair{
{
ID: agent1ID,
Name: "agent1",
},
},
},
},
}
ps := &mockPubsub{
cbs: map[string]pubsub.ListenerWithErr{},
}
updateProvider := coderd.NewUpdatesProvider(slogtest.Make(t, nil), ps, db, &mockAuthorizer{})
t.Cleanup(func() {
_ = updateProvider.Close()
})
sub, err := updateProvider.Subscribe(dbauthz.As(ctx, ownerSubject), ownerID)
require.NoError(t, err)
t.Cleanup(func() {
_ = sub.Close()
})
expected := &proto.WorkspaceUpdate{
UpsertedWorkspaces: []*proto.Workspace{
{
Id: ws1IDSlice,
Name: "ws1",
Status: proto.Workspace_STARTING,
},
},
UpsertedAgents: []*proto.Agent{
{
Id: agent1IDSlice,
Name: "agent1",
WorkspaceId: ws1IDSlice,
},
},
DeletedWorkspaces: []*proto.Workspace{},
DeletedAgents: []*proto.Agent{},
}
update := testutil.RequireRecvCtx(ctx, t, sub.Updates())
slices.SortFunc(update.UpsertedWorkspaces, func(a, b *proto.Workspace) int {
return strings.Compare(a.Name, b.Name)
})
require.Equal(t, expected, update)
resub, err := updateProvider.Subscribe(dbauthz.As(ctx, ownerSubject), ownerID)
require.NoError(t, err)
t.Cleanup(func() {
_ = resub.Close()
})
update = testutil.RequireRecvCtx(ctx, t, resub.Updates())
slices.SortFunc(update.UpsertedWorkspaces, func(a, b *proto.Workspace) int {
return strings.Compare(a.Name, b.Name)
})
require.Equal(t, expected, update)
})
}
func publishWorkspaceEvent(t *testing.T, ps pubsub.Pubsub, ownerID uuid.UUID, event *wspubsub.WorkspaceEvent) {
msg, err := json.Marshal(event)
require.NoError(t, err)
ps.Publish(wspubsub.WorkspaceEventChannel(ownerID), msg)
}
type mockWorkspaceStore struct {
orderedRows []database.GetWorkspacesAndAgentsByOwnerIDRow
}
// GetAuthorizedWorkspacesAndAgentsByOwnerID implements coderd.UpdatesQuerier.
func (m *mockWorkspaceStore) GetWorkspacesAndAgentsByOwnerID(context.Context, uuid.UUID) ([]database.GetWorkspacesAndAgentsByOwnerIDRow, error) {
return m.orderedRows, nil
}
// GetWorkspaceByAgentID implements coderd.UpdatesQuerier.
func (*mockWorkspaceStore) GetWorkspaceByAgentID(context.Context, uuid.UUID) (database.Workspace, error) {
return database.Workspace{}, nil
}
var _ coderd.UpdatesQuerier = (*mockWorkspaceStore)(nil)
type mockPubsub struct {
cbs map[string]pubsub.ListenerWithErr
}
// Close implements pubsub.Pubsub.
func (*mockPubsub) Close() error {
panic("unimplemented")
}
// Publish implements pubsub.Pubsub.
func (m *mockPubsub) Publish(event string, message []byte) error {
cb, ok := m.cbs[event]
if !ok {
return nil
}
cb(context.Background(), message, nil)
return nil
}
func (*mockPubsub) Subscribe(string, pubsub.Listener) (cancel func(), err error) {
panic("unimplemented")
}
func (m *mockPubsub) SubscribeWithErr(event string, listener pubsub.ListenerWithErr) (func(), error) {
m.cbs[event] = listener
return func() {}, nil
}
var _ pubsub.Pubsub = (*mockPubsub)(nil)
type mockAuthorizer struct{}
func (*mockAuthorizer) Authorize(context.Context, rbac.Subject, policy.Action, rbac.Object) error {
return nil
}
// Prepare implements rbac.Authorizer.
func (*mockAuthorizer) Prepare(context.Context, rbac.Subject, policy.Action, string) (rbac.PreparedAuthorized, error) {
return nil, nil
}
var _ rbac.Authorizer = (*mockAuthorizer)(nil)