Merge #150721

150721: pgwire: decouple jwt authentication and authorization logic r=souravcrl a=shriramters

Previously, the logic for JWT authorization (extracting group claims and synchronizing roles) was located within the `Authenticator` behavior. This was a necessary design when the JWT token was only exchanged from the `AuthConn` within the authenticator.

This was inadequate because it bundled authentication and authorization logic, violating the separation of concerns intended by the `AuthBehaviors` framework and creating an inconsistency with other methods like `AuthLDAP`. A recent change (#149415) made the token available via closure capture to all behaviors, removing the original constraint and making this refactoring possible.

To address this, this patch decouples the logic. The `Authenticator` is now solely responsible for validating the token (authentication). The authorization logic has been moved to a new `Authorizer` behavior, which aligns the implementation with the framework's design and improves code clarity and maintainability.

Fixes: #150720

Release note: None

Co-authored-by: Shriram Ravindranathan <[email protected]>

Author: craig[bot]

pkg/security/provisioning/settings.go

@@ -24,18 +24,28 @@ const (
 
 	baseCounterPrefix = "auth.provisioning."
 	ldapCounterPrefix = baseCounterPrefix + "ldap."
+	jwtCounterPrefix  = baseCounterPrefix + "jwt."
 
 	beginLDAPProvisionCounterName   = ldapCounterPrefix + "begin"
 	provisionLDAPSuccessCounterName = ldapCounterPrefix + "success"
 	enableLDAPProvisionCounterName  = ldapCounterPrefix + "enable"
 
+	beginJWTProvisionCounterName   = jwtCounterPrefix + "begin"
+	provisionJWTSuccessCounterName = jwtCounterPrefix + "success"
+	enableJWTProvisionCounterName  = jwtCounterPrefix + "enable"
+
 	provisionedUserLoginSuccessCounterName = baseCounterPrefix + "login_success"
 )
 
 var (
-	BeginLDAPProvisionUseCounter       = telemetry.GetCounterOnce(beginLDAPProvisionCounterName)
-	ProvisionLDAPSuccessCounter        = telemetry.GetCounterOnce(provisionLDAPSuccessCounterName)
-	enableLDAPProvisionCounter         = telemetry.GetCounterOnce(enableLDAPProvisionCounterName)
+	BeginLDAPProvisionUseCounter = telemetry.GetCounterOnce(beginLDAPProvisionCounterName)
+	ProvisionLDAPSuccessCounter  = telemetry.GetCounterOnce(provisionLDAPSuccessCounterName)
+	enableLDAPProvisionCounter   = telemetry.GetCounterOnce(enableLDAPProvisionCounterName)
+
+	BeginJWTProvisionUseCounter = telemetry.GetCounterOnce(beginJWTProvisionCounterName)
+	ProvisionJWTSuccessCounter  = telemetry.GetCounterOnce(provisionJWTSuccessCounterName)
+	enableJWTProvisionCounter   = telemetry.GetCounterOnce(enableJWTProvisionCounterName)
+
 	ProvisionedUserLoginSuccessCounter = telemetry.GetCounterOnce(provisionedUserLoginSuccessCounterName)
 )
 
@@ -99,5 +109,10 @@ func ClusterProvisioningConfig(settings *cluster.Settings) UserProvisioningConfi
 			telemetry.Inc(enableLDAPProvisionCounter)
 		}
 	})
+	jwtProvisioningEnabled.SetOnChange(&settings.SV, func(_ context.Context) {
+		if jwtProvisioningEnabled.Get(&settings.SV) {
+			telemetry.Inc(enableJWTProvisionCounter)
+		}
+	})
 	return clusterProvisioningConfig{settings: settings}
 }

pkg/sql/pgwire/auth_methods.go

@@ -860,6 +860,9 @@ func authJwtToken(
 	}
 
 	b.SetProvisioner(func(ctx context.Context) error {
+		c.LogAuthInfof(ctx, "Starting JWT provisioning; attempting to verify token")
+		telemetry.Inc(provisioning.BeginJWTProvisionUseCounter)
+
 		if validationErr != nil {
 			return validationErr
 		}
@@ -893,6 +896,8 @@ func authJwtToken(
 			c.LogAuthFailed(ctx, eventpb.AuthFailReason_PROVISIONING_ERROR, err)
 			return err
 		}
+
+		telemetry.Inc(provisioning.ProvisionJWTSuccessCounter)
 		return nil
 	})
 
@@ -905,6 +910,8 @@ func authJwtToken(
 			return err
 		}
 
+		c.LogAuthInfof(ctx, "JWT Provided; attempting to validate token for authentication")
+
 		// Validate the token and, if there's an error, log it with the correct reason.
 		if detailedErrors, authError := jwtVerifier.ValidateJWTLogin(
 			sctx, execCfg.Settings, user, []byte(token), identMap); authError != nil {
@@ -916,6 +923,12 @@ func authJwtToken(
 			return authError
 		}
 
+		return nil
+	})
+
+	b.SetAuthorizer(func(ctx context.Context, systemIdentity string, clientConnection bool) error {
+		c.LogAuthInfof(ctx, "JWT authentication succeeded; attempting authorization")
+
 		// Ask the CCL verifier for groups (nil slice means feature disabled).
 		groups, err := jwtVerifier.ExtractGroups(ctx, execCfg.Settings, []byte(token))
 		if err != nil {
@@ -959,6 +972,7 @@ func authJwtToken(
 
 		return nil
 	})
+
 	return b, nil
 }