@@ -25,6 +25,7 @@ import (
25
25
"github.com/coder/coder/v2/coderd/database/dbrollup"
26
26
"github.com/coder/coder/v2/coderd/database/dbtestutil"
27
27
"github.com/coder/coder/v2/coderd/database/dbtime"
28
+ "github.com/coder/coder/v2/coderd/provisionerdserver"
28
29
"github.com/coder/coder/v2/codersdk"
29
30
"github.com/coder/coder/v2/provisionerd/proto"
30
31
"github.com/coder/coder/v2/provisionersdk"
@@ -490,3 +491,71 @@ func containsProvisionerDaemon(daemons []database.ProvisionerDaemon, name string
490
491
return d .Name == name
491
492
})
492
493
}
494
+
495
+ func TestExpireOldAPIKeys (t * testing.T ) {
496
+ t .Parallel ()
497
+ if ! dbtestutil .WillUsePostgres () {
498
+ t .Skip ("only implemented in postgres" )
499
+ }
500
+
501
+ // Given: a number of workspaces and API keys owned by a regular user and the prebuilds system user.
502
+ var (
503
+ ctx = testutil .Context (t , testutil .WaitShort )
504
+ now = dbtime .Now ()
505
+ db , _ = dbtestutil .NewDB (t , dbtestutil .WithDumpOnFailure ())
506
+ org = dbgen .Organization (t , db , database.Organization {})
507
+ user = dbgen .User (t , db , database.User {})
508
+ tpl = dbgen .Template (t , db , database.Template {OrganizationID : org .ID , CreatedBy : user .ID })
509
+ userWs = dbgen .Workspace (t , db , database.WorkspaceTable {
510
+ OwnerID : user .ID ,
511
+ TemplateID : tpl .ID ,
512
+ })
513
+ prebuildsWs = dbgen .Workspace (t , db , database.WorkspaceTable {
514
+ OwnerID : database .PrebuildsSystemUserID ,
515
+ TemplateID : tpl .ID ,
516
+ })
517
+ createAPIKey = func (userID uuid.UUID , name string ) database.APIKey {
518
+ k , _ := dbgen .APIKey (t , db , database.APIKey {UserID : userID , TokenName : name , ExpiresAt : now .Add (time .Hour )}, func (iap * database.InsertAPIKeyParams ) {
519
+ iap .TokenName = name
520
+ })
521
+ return k
522
+ }
523
+ assertKeyActive = func (kid string ) {
524
+ k , err := db .GetAPIKeyByID (ctx , kid )
525
+ require .NoError (t , err )
526
+ assert .True (t , k .ExpiresAt .After (now ))
527
+ }
528
+ assertKeyExpired = func (kid string ) {
529
+ k , err := db .GetAPIKeyByID (ctx , kid )
530
+ require .NoError (t , err )
531
+ assert .True (t , k .ExpiresAt .Equal (now ))
532
+ }
533
+ unnamedUserAPIKey = createAPIKey (user .ID , "" )
534
+ unnamedPrebuildsAPIKey = createAPIKey (database .PrebuildsSystemUserID , "" )
535
+ namedUserAPIKey = createAPIKey (user .ID , "my-token" )
536
+ namedPrebuildsAPIKey = createAPIKey (database .PrebuildsSystemUserID , "also-my-token" )
537
+ userWorkspaceAPIKey1 = createAPIKey (user .ID , provisionerdserver .WorkspaceSessionTokenName (user .ID , userWs .ID ))
538
+ userWorkspaceAPIKey2 = createAPIKey (user .ID , provisionerdserver .WorkspaceSessionTokenName (user .ID , prebuildsWs .ID ))
539
+ prebuildsWorkspaceAPIKey1 = createAPIKey (database .PrebuildsSystemUserID , provisionerdserver .WorkspaceSessionTokenName (database .PrebuildsSystemUserID , prebuildsWs .ID ))
540
+ prebuildsWorkspaceAPIKey2 = createAPIKey (database .PrebuildsSystemUserID , provisionerdserver .WorkspaceSessionTokenName (database .PrebuildsSystemUserID , userWs .ID ))
541
+ )
542
+
543
+ // When: we call ExpirePrebuildsAPIKeys
544
+ err := db .ExpirePrebuildsAPIKeys (ctx , now )
545
+ // Then: no errors is reported.
546
+ require .NoError (t , err )
547
+
548
+ // We do not touch user API keys.
549
+ assertKeyActive (unnamedUserAPIKey .ID )
550
+ assertKeyActive (namedUserAPIKey .ID )
551
+ assertKeyActive (userWorkspaceAPIKey1 .ID )
552
+ assertKeyActive (userWorkspaceAPIKey2 .ID )
553
+ // Unnamed prebuilds API keys get expired.
554
+ assertKeyExpired (unnamedPrebuildsAPIKey .ID )
555
+ // API keys for workspaces still owned by prebuilds user remain active until claimed.
556
+ assertKeyActive (prebuildsWorkspaceAPIKey1 .ID )
557
+ // API keys for workspaces no longer owned by prebuilds user get expired.
558
+ assertKeyExpired (prebuildsWorkspaceAPIKey2 .ID )
559
+ // Out of an abundance of caution, we do not expire explicitly named prebuilds API keys.
560
+ assertKeyActive (namedPrebuildsAPIKey .ID )
561
+ }
0 commit comments