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:
Generated
+19
@@ -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": [
|
||||
|
||||
Generated
+17
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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`.
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user