MWI: Use specific lock targets when locking out bots

lib/auth/auth.go

@@ -3211,6 +3211,8 @@ func generateCert(ctx context.Context, a *Server, req certRequest, caType types.
 		mfaVerified:          req.mfaVerified,
 		activeAccessRequests: req.activeRequests,
 		deviceID:             req.deviceExtensions.DeviceID,
+		botInstanceID:        req.botInstanceID,
+		joinToken:            req.joinToken,
 	}); err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -3675,6 +3677,10 @@ type verifyLocksForUserCertsReq struct {
 	// deviceID is the trusted device ID.
 	// Eg: tlsca.Identity.DeviceExtensions.DeviceID
 	deviceID string
+	// botInstanceID is the bot instance UUID, set only for bots.
+	botInstanceID string
+	// joinMethod is the join token name, set only for non-token bots.
+	joinToken string
 }
 
 // verifyLocksForUserCerts verifies if any locks are in place before issuing new
@@ -3694,6 +3700,12 @@ func (a *Server) verifyLocksForUserCerts(req verifyLocksForUserCertsReq) error {
 	lockTargets = append(lockTargets,
 		services.AccessRequestsToLockTargets(req.activeAccessRequests)...,
 	)
+	if req.botInstanceID != "" {
+		lockTargets = append(lockTargets, types.LockTarget{BotInstanceID: req.botInstanceID})
+	}
+	if req.joinToken != "" {
+		lockTargets = append(lockTargets, types.LockTarget{JoinToken: req.joinToken})
+	}
 
 	return trace.Wrap(a.checkLockInForce(lockingMode, lockTargets))
 }

lib/auth/bot.go

@@ -121,7 +121,13 @@ func (a *Server) legacyValidateGenerationLabel(ctx context.Context, username str
 
 	// The current generations must match to continue:
 	if currentIdentityGeneration != currentUserGeneration {
-		if err := a.tryLockBotDueToGenerationMismatch(ctx, user.GetName()); err != nil {
+		if err := a.tryLockBotDueToGenerationMismatch(
+			ctx,
+			certReq.botName,
+			certReq.botInstanceID,
+			certReq.joinToken,
+			certReq.renewable,
+		); err != nil {
 			a.logger.WarnContext(ctx, "Failed to lock bot when a generation mismatch was detected",
 				"error", err,
 				"bot", user.GetName(),
@@ -223,19 +229,44 @@ func (a *Server) commitLegacyGenerationCounterToBotUser(ctx context.Context, use
 
 // tryLockBotDueToGenerationMismatch creates a lock for the given bot user and
 // emits a `RenewableCertificateGenerationMismatch` audit event.
-func (a *Server) tryLockBotDueToGenerationMismatch(ctx context.Context, username string) error {
-	// TODO: In the future, consider only locking the current join method / token.
+func (a *Server) tryLockBotDueToGenerationMismatch(
+	ctx context.Context, botName, botInstanceID, joinTokenName string, renewable bool,
+) error {
+	var spec types.LockSpecV2
+	if renewable {
+		// Renewable implies `token` joining. These are one-time use secrets
+		// and will not be embedded in the TLS identity, so we can't target
+		// the join token and should instead rely on the bot instance ID. As
+		// there is a 1:1 relationship between bot instance and "token"-type
+		// token, this should be functionally equivalent.
+		spec = types.LockSpecV2{
+			Target: types.LockTarget{
+				BotInstanceID: botInstanceID,
+			},
+			Message: fmt.Sprintf(
+				"The bot instance %s/%s has been locked due to a certificate "+
+					"generation mismatch, possibly indicating a stolen "+
+					"certificate.",
+				botName, botInstanceID,
+			),
+			CreatedAt: a.clock.Now(),
+		}
+	} else {
+		spec = types.LockSpecV2{
+			Target: types.LockTarget{
+				JoinToken: joinTokenName,
+			},
+			Message: fmt.Sprintf(
+				"Bot joins via the token %q have been locked due to a "+
+					"certificate generation mismatch by %s/%s, possibly "+
+					"indicating a stolen certificate.",
+				joinTokenName, botName, botInstanceID,
+			),
+			CreatedAt: a.clock.Now(),
+		}
+	}
 
-	// Lock the bot user indefinitely.
-	lock, err := types.NewLock(uuid.New().String(), types.LockSpecV2{
-		Target: types.LockTarget{
-			User: username,
-		},
-		Message: fmt.Sprintf(
-			"The bot user %q has been locked due to a certificate generation mismatch, possibly indicating a stolen certificate.",
-			username,
-		),
-	})
+	lock, err := types.NewLock(uuid.New().String(), spec)
 	if err != nil {
 		return trace.Wrap(err)
 	}
@@ -360,7 +391,12 @@ func (a *Server) updateBotInstance(
 			// If the incoming identity has a nonzero generation, validate it
 			// using the legacy check. This will increment the counter on the
 			// request automatically
-			if err := a.legacyValidateGenerationLabel(ctx, username, req, uint64(currentIdentityGeneration)); err != nil {
+			if err := a.legacyValidateGenerationLabel(
+				ctx,
+				username,
+				req,
+				uint64(currentIdentityGeneration),
+			); err != nil {
 				return trace.Wrap(err)
 			}
 
@@ -409,7 +445,7 @@ func (a *Server) updateBotInstance(
 		// Generation counter enforcement depends on the type of cert and join
 		// method (if any - token renewals technically have no join method.)
 		if shouldEnforceGenerationCounter(req.renewable, authRecord.JoinMethod) {
-			if err := a.tryLockBotDueToGenerationMismatch(ctx, username); err != nil {
+			if err := a.tryLockBotDueToGenerationMismatch(ctx, botName, botInstanceID, req.joinToken, req.renewable); err != nil {
 				log.WarnContext(ctx, "Failed to lock bot when a generation mismatch was detected", "error", err)
 			}
 
@@ -572,6 +608,14 @@ func (a *Server) generateInitialBotCerts(
 		joinAttributes: joinAttrs,
 	}
 
+	// Set the join token cert field for non-renewable identities. This is used
+	// for lock targeting; token name lock targets are particularly useful for
+	// token-joined bots and it's a secret value, so we don't bother setting it.
+	// (The renewable flag implies token joining.)
+	if !renewable {
+		certReq.joinToken = initialAuth.JoinToken
+	}
+
 	if existingInstanceID == "" {
 		// If no existing instance ID is known, create a new one.
 		uuid, err := uuid.NewRandom()
@@ -614,14 +658,6 @@ func (a *Server) generateInitialBotCerts(
 		}
 	}
 
-	// Set the join token cert field for non-renewable identities. This is used
-	// for lock targeting; token name lock targets are particularly useful for
-	// token-joined bots and it's a secret value, so we don't bother setting it.
-	// (The renewable flag implies token joining.)
-	if !renewable {
-		certReq.joinToken = initialAuth.JoinToken
-	}
-
 	certs, err := a.generateUserCert(ctx, certReq)
 	if err != nil {
 		return nil, "", trace.Wrap(err)

lib/auth/bot_test.go

@@ -549,7 +549,11 @@ func TestRegisterBotCertificateGenerationStolen(t *testing.T) {
 	require.NoError(t, err)
 
 	// Renew the certs once (e.g. this is the actual bot process)
-	_, certsReal, err := renewBotCerts(ctx, srv, result.Certs.TLS, bot.Status.UserName, result.PrivateKey)
+	renewedClient, certsReal, err := renewBotCerts(ctx, srv, result.Certs.TLS, bot.Status.UserName, result.PrivateKey)
+	require.NoError(t, err)
+
+	// This client should be able to ping.
+	_, err = renewedClient.Ping(ctx)
 	require.NoError(t, err)
 
 	// Check the generation, it should be 2.
@@ -564,12 +568,16 @@ func TestRegisterBotCertificateGenerationStolen(t *testing.T) {
 	require.Error(t, err)
 	require.True(t, trace.IsAccessDenied(err))
 
-	// The user should now be locked.
+	// The bot instance should now be locked.
 	locks, err := srv.Auth().GetLocks(ctx, true, types.LockTarget{
-		User: "bot-test",
+		BotInstanceID: impersonatedIdent.BotInstanceID,
 	})
 	require.NoError(t, err)
 	require.NotEmpty(t, locks)
+
+	// The original client should now be locked out.
+	_, err = renewedClient.Ping(ctx)
+	require.ErrorContains(t, err, "access denied")
 }
 
 // TestRegisterBotCertificateExtensions ensures bot cert extensions are present.