Skip to content

Commit 5ee6ad5

Browse files
sreyasreya
committed
add test for tailnet_resume jwt
1 parent 6b9a3e4 commit 5ee6ad5

File tree

2 files changed

+136
-65
lines changed

2 files changed

+136
-65
lines changed

‎coderd/workspaceagents.go

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ import (
3232
"github.com/coder/coder/v2/coderd/externalauth"
3333
"github.com/coder/coder/v2/coderd/httpapi"
3434
"github.com/coder/coder/v2/coderd/httpmw"
35+
"github.com/coder/coder/v2/coderd/jwtutils"
3536
"github.com/coder/coder/v2/coderd/rbac/policy"
3637
"github.com/coder/coder/v2/codersdk"
3738
"github.com/coder/coder/v2/codersdk/agentsdk"
@@ -854,7 +855,11 @@ func (api *API) workspaceAgentClientCoordinate(rw http.ResponseWriter, r *http.R
854855
if resumeToken != "" {
855856
var err error
856857
peerID, err = api.Options.CoordinatorResumeTokenProvider.VerifyResumeToken(ctx, resumeToken)
857-
if err != nil {
858+
// If the token is missing the key ID, it's probably an old token in which
859+
// case we just want to generate a new peer ID.
860+
if xerrors.Is(err, jwtutils.ErrMissingKeyID) {
861+
peerID = uuid.New()
862+
} else if err != nil {
858863
httpapi.Write(ctx, rw, http.StatusUnauthorized, codersdk.Response{
859864
Message: workspacesdk.CoordinateAPIInvalidResumeToken,
860865
Detail: err.Error(),
@@ -863,9 +868,10 @@ func (api *API) workspaceAgentClientCoordinate(rw http.ResponseWriter, r *http.R
863868
},
864869
})
865870
return
871+
} else {
872+
api.Logger.Debug(ctx, "accepted coordinate resume token for peer",
873+
slog.F("peer_id", peerID.String()))
866874
}
867-
api.Logger.Debug(ctx, "accepted coordinate resume token for peer",
868-
slog.F("peer_id", peerID.String()))
869875
}
870876

871877
api.WebsocketWaitMutex.Lock()

‎coderd/workspaceagents_test.go

Lines changed: 127 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313
"testing"
1414
"time"
1515

16+
"github.com/go-jose/go-jose/v4/jwt"
1617
"github.com/google/uuid"
1718
"github.com/stretchr/testify/assert"
1819
"github.com/stretchr/testify/require"
@@ -37,6 +38,7 @@ import (
3738
"github.com/coder/coder/v2/coderd/database/dbtime"
3839
"github.com/coder/coder/v2/coderd/database/pubsub"
3940
"github.com/coder/coder/v2/coderd/externalauth"
41+
"github.com/coder/coder/v2/coderd/jwtutils"
4042
"github.com/coder/coder/v2/codersdk"
4143
"github.com/coder/coder/v2/codersdk/agentsdk"
4244
"github.com/coder/coder/v2/codersdk/workspacesdk"
@@ -555,73 +557,136 @@ func (r *resumeTokenRecordingProvider) VerifyResumeToken(ctx context.Context, to
555557
func TestWorkspaceAgentClientCoordinate_ResumeToken(t *testing.T) {
556558
t.Parallel()
557559

558-
logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
559-
clock := quartz.NewMock(t)
560-
resumeTokenSigningKey, err := tailnet.GenerateResumeTokenSigningKey()
561-
mgr := cryptokeys.StaticKey{
562-
ID: uuid.New().String(),
563-
Key: resumeTokenSigningKey[:],
564-
}
565-
require.NoError(t, err)
566-
resumeTokenProvider := newResumeTokenRecordingProvider(
567-
t,
568-
tailnet.NewResumeTokenKeyProvider(mgr, clock, time.Hour),
569-
)
570-
client, closer, api := coderdtest.NewWithAPI(t, &coderdtest.Options{
571-
Coordinator: tailnet.NewCoordinator(logger),
572-
CoordinatorResumeTokenProvider: resumeTokenProvider,
560+
t.Run("OK", func(t *testing.T) {
561+
t.Parallel()
562+
563+
logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
564+
clock := quartz.NewMock(t)
565+
resumeTokenSigningKey, err := tailnet.GenerateResumeTokenSigningKey()
566+
mgr := cryptokeys.StaticKey{
567+
ID: uuid.New().String(),
568+
Key: resumeTokenSigningKey[:],
569+
}
570+
require.NoError(t, err)
571+
resumeTokenProvider := newResumeTokenRecordingProvider(
572+
t,
573+
tailnet.NewResumeTokenKeyProvider(mgr, clock, time.Hour),
574+
)
575+
client, closer, api := coderdtest.NewWithAPI(t, &coderdtest.Options{
576+
Coordinator: tailnet.NewCoordinator(logger),
577+
CoordinatorResumeTokenProvider: resumeTokenProvider,
578+
})
579+
defer closer.Close()
580+
user := coderdtest.CreateFirstUser(t, client)
581+
582+
// Create a workspace with an agent. No need to connect it since clients can
583+
// still connect to the coordinator while the agent isn't connected.
584+
r := dbfake.WorkspaceBuild(t, api.Database, database.Workspace{
585+
OrganizationID: user.OrganizationID,
586+
OwnerID: user.UserID,
587+
}).WithAgent().Do()
588+
agentTokenUUID, err := uuid.Parse(r.AgentToken)
589+
require.NoError(t, err)
590+
ctx := testutil.Context(t, testutil.WaitLong)
591+
agentAndBuild, err := api.Database.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentTokenUUID) //nolint
592+
require.NoError(t, err)
593+
594+
// Connect with no resume token, and ensure that the peer ID is set to a
595+
// random value.
596+
originalResumeToken, err := connectToCoordinatorAndFetchResumeToken(ctx, logger, client, agentAndBuild.WorkspaceAgent.ID, "")
597+
require.NoError(t, err)
598+
originalPeerID := testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.generateCalls)
599+
require.NotEqual(t, originalPeerID, uuid.Nil)
600+
601+
// Connect with a valid resume token, and ensure that the peer ID is set to
602+
// the stored value.
603+
clock.Advance(time.Second)
604+
newResumeToken, err := connectToCoordinatorAndFetchResumeToken(ctx, logger, client, agentAndBuild.WorkspaceAgent.ID, originalResumeToken)
605+
require.NoError(t, err)
606+
verifiedToken := testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.verifyCalls)
607+
require.Equal(t, originalResumeToken, verifiedToken)
608+
newPeerID := testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.generateCalls)
609+
require.Equal(t, originalPeerID, newPeerID)
610+
require.NotEqual(t, originalResumeToken, newResumeToken)
611+
612+
// Connect with an invalid resume token, and ensure that the request is
613+
// rejected.
614+
clock.Advance(time.Second)
615+
_, err = connectToCoordinatorAndFetchResumeToken(ctx, logger, client, agentAndBuild.WorkspaceAgent.ID, "invalid")
616+
require.Error(t, err)
617+
var sdkErr *codersdk.Error
618+
require.ErrorAs(t, err, &sdkErr)
619+
require.Equal(t, http.StatusUnauthorized, sdkErr.StatusCode())
620+
require.Len(t, sdkErr.Validations, 1)
621+
require.Equal(t, "resume_token", sdkErr.Validations[0].Field)
622+
verifiedToken = testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.verifyCalls)
623+
require.Equal(t, "invalid", verifiedToken)
624+
625+
select {
626+
case <-resumeTokenProvider.generateCalls:
627+
t.Fatal("unexpected peer ID in channel")
628+
default:
629+
}
573630
})
574-
defer closer.Close()
575-
user := coderdtest.CreateFirstUser(t, client)
576631

577-
// Create a workspace with an agent. No need to connect it since clients can
578-
// still connect to the coordinator while the agent isn't connected.
579-
r := dbfake.WorkspaceBuild(t, api.Database, database.Workspace{
580-
OrganizationID: user.OrganizationID,
581-
OwnerID: user.UserID,
582-
}).WithAgent().Do()
583-
agentTokenUUID, err := uuid.Parse(r.AgentToken)
584-
require.NoError(t, err)
585-
ctx := testutil.Context(t, testutil.WaitLong)
586-
agentAndBuild, err := api.Database.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentTokenUUID) //nolint
587-
require.NoError(t, err)
632+
t.Run("BadJWT", func(t *testing.T) {
633+
t.Parallel()
588634

589-
// Connect with no resume token, and ensure that the peer ID is set to a
590-
// random value.
591-
originalResumeToken, err := connectToCoordinatorAndFetchResumeToken(ctx, logger, client, agentAndBuild.WorkspaceAgent.ID, "")
592-
require.NoError(t, err)
593-
originalPeerID := testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.generateCalls)
594-
require.NotEqual(t, originalPeerID, uuid.Nil)
635+
logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
636+
clock := quartz.NewMock(t)
637+
resumeTokenSigningKey, err := tailnet.GenerateResumeTokenSigningKey()
638+
mgr := cryptokeys.StaticKey{
639+
ID: uuid.New().String(),
640+
Key: resumeTokenSigningKey[:],
641+
}
642+
require.NoError(t, err)
643+
resumeTokenProvider := newResumeTokenRecordingProvider(
644+
t,
645+
tailnet.NewResumeTokenKeyProvider(mgr, clock, time.Hour),
646+
)
647+
client, closer, api := coderdtest.NewWithAPI(t, &coderdtest.Options{
648+
Coordinator: tailnet.NewCoordinator(logger),
649+
CoordinatorResumeTokenProvider: resumeTokenProvider,
650+
})
651+
defer closer.Close()
652+
user := coderdtest.CreateFirstUser(t, client)
595653

596-
// Connect with a valid resume token, and ensure that the peer ID is set to
597-
// the stored value.
598-
clock.Advance(time.Second)
599-
newResumeToken, err := connectToCoordinatorAndFetchResumeToken(ctx, logger, client, agentAndBuild.WorkspaceAgent.ID, originalResumeToken)
600-
require.NoError(t, err)
601-
verifiedToken := testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.verifyCalls)
602-
require.Equal(t, originalResumeToken, verifiedToken)
603-
newPeerID := testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.generateCalls)
604-
require.Equal(t, originalPeerID, newPeerID)
605-
require.NotEqual(t, originalResumeToken, newResumeToken)
606-
607-
// Connect with an invalid resume token, and ensure that the request is
608-
// rejected.
609-
clock.Advance(time.Second)
610-
_, err = connectToCoordinatorAndFetchResumeToken(ctx, logger, client, agentAndBuild.WorkspaceAgent.ID, "invalid")
611-
require.Error(t, err)
612-
var sdkErr *codersdk.Error
613-
require.ErrorAs(t, err, &sdkErr)
614-
require.Equal(t, http.StatusUnauthorized, sdkErr.StatusCode())
615-
require.Len(t, sdkErr.Validations, 1)
616-
require.Equal(t, "resume_token", sdkErr.Validations[0].Field)
617-
verifiedToken = testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.verifyCalls)
618-
require.Equal(t, "invalid", verifiedToken)
654+
// Create a workspace with an agent. No need to connect it since clients can
655+
// still connect to the coordinator while the agent isn't connected.
656+
r := dbfake.WorkspaceBuild(t, api.Database, database.Workspace{
657+
OrganizationID: user.OrganizationID,
658+
OwnerID: user.UserID,
659+
}).WithAgent().Do()
660+
agentTokenUUID, err := uuid.Parse(r.AgentToken)
661+
require.NoError(t, err)
662+
ctx := testutil.Context(t, testutil.WaitLong)
663+
agentAndBuild, err := api.Database.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentTokenUUID) //nolint
664+
require.NoError(t, err)
619665

620-
select {
621-
case <-resumeTokenProvider.generateCalls:
622-
t.Fatal("unexpected peer ID in channel")
623-
default:
624-
}
666+
// Connect with no resume token, and ensure that the peer ID is set to a
667+
// random value.
668+
originalResumeToken, err := connectToCoordinatorAndFetchResumeToken(ctx, logger, client, agentAndBuild.WorkspaceAgent.ID, "")
669+
require.NoError(t, err)
670+
originalPeerID := testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.generateCalls)
671+
require.NotEqual(t, originalPeerID, uuid.Nil)
672+
673+
// Connect with an outdated token, and ensure that the peer ID is set to a
674+
// random value. We don't want to fail requests just because
675+
// a user got unlucky during a deployment upgrade.
676+
outdatedToken := generateBadJWT(t, jwtutils.RegisteredClaims{
677+
Subject: originalPeerID.String(),
678+
Expiry: jwt.NewNumericDate(clock.Now().Add(time.Minute)),
679+
})
680+
681+
clock.Advance(time.Second)
682+
newResumeToken, err := connectToCoordinatorAndFetchResumeToken(ctx, logger, client, agentAndBuild.WorkspaceAgent.ID, outdatedToken)
683+
require.NoError(t, err)
684+
verifiedToken := testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.verifyCalls)
685+
require.Equal(t, outdatedToken, verifiedToken)
686+
newPeerID := testutil.RequireRecvCtx(ctx, t, resumeTokenProvider.generateCalls)
687+
require.NotEqual(t, originalPeerID, newPeerID)
688+
require.NotEqual(t, originalResumeToken, newResumeToken)
689+
})
625690
}
626691

627692
// connectToCoordinatorAndFetchResumeToken connects to the tailnet coordinator

0 commit comments

Comments
 (0)