chore: Relax the GKE WIF principal:// restriction

sasl-plain-access-token/segmentio/saslplainoauthmechanism/README.md

@@ -9,7 +9,6 @@ It allows you to use Authorization Tokens with SASL/Plain in [segmentio/kafka-go
 saslplainoauthmechanism supports the following Application Default credential types:
 
 1. [GKE Workload Identity Federation](https://cloud.google.com/kubernetes-engine/docs/concepts/workload-identity).
-    1. Note: Native Workload Identity Federation principals (`principal://` format) are currently unsupported, you must [link your Kubernetes Service Account](https://cloud.google.com/kubernetes-engine/docs/how-to/workload-identity#kubernetes-sa-to-iam) to a Google Service Account.
 2. [Metadata Server Credentials](https://cloud.google.com/docs/authentication/application-default-credentials#attached-sa) - Such as Google Compute Engine, Cloud Run, etc.
 3. [gcloud CLI Application Default Credentials](https://cloud.google.com/docs/authentication/application-default-credentials#personal). Specifically the following subset:
     1. `user_credentials` (`gcloud auth application-default login`).

sasl-plain-access-token/segmentio/saslplainoauthmechanism/principal_email.go

@@ -47,7 +47,7 @@ func getADCPrincipalEmail(creds *google.Credentials) (string, error) {
 		if err != nil {
 			return "", fmt.Errorf("detected GCE, but unable to get Service Account email for default: %w", err)
 		}
-		if validatePrincipalEmail(email) != nil {
+		if err := validatePrincipalEmail(email); err != nil {
 			return "", err
 		}
 		return email, nil
@@ -62,15 +62,25 @@ func getADCPrincipalEmail(creds *google.Credentials) (string, error) {
 	return "", errors.New("unable to determine principal email, did not detect Metadata Server or JSON Credentials")
 }
 
-// Checks that the Service Account in use is a Google Service Account, and not an unsupported type
+// Checks that the principal email in use is valid format for a 1P or 3P identity
 func validatePrincipalEmail(email string) error {
-	// TODO: Relax principal:// restriction once b/385138184 is resolved.
+
 	// There are two responses that the Metadata server will return for http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/email in GKE
+	// when NOT using the GA impersonation approach (https://cloud.google.com/kubernetes-engine/docs/how-to/workload-identity#kubernetes-sa-to-iam)
 	// If the KSA is **not** annotated with 'iam.gke.io/return-principal-id-as-email: "true"', then it will return the Workload Identity Pool name
 	// If it **is** annotated with 'iam.gke.io/return-principal-id-as-email: "true"' then it will return the full principal identifier
-	if strings.HasSuffix(email, ".svc.id.goog") || strings.HasPrefix(email, "principal://iam.googleapis.com") {
-		return errors.New("GMK SASL PLAIN OAuth cannot be used with direct Workload Identity Federation - you must configure KSA --> GSA impersonation, see: https://cloud.google.com/kubernetes-engine/docs/how-to/workload-identity#kubernetes-sa-to-iam")
+	if strings.HasSuffix(email, ".svc.id.goog") {
+		return errors.New("to use  GKE Workload Identity Federation you must annotate your Kubernetes Service Account with 'iam.gke.io/return-principal-id-as-email: \"true\"'")
+	}
+
+	// Identities are either an email (service account, user email, or group) or a Workload/Workforce identity principal
+	// All Workload/Workforce identity pools start with 'principal://iam.googleapis.com/' so we assume valid if our principal
+	// string begins with this
+	if strings.HasPrefix(email, "principal://iam.googleapis.com/") {
+		return nil
 	}
+
+	// Else check it is a valid email identifier
 	_, err := mail.ParseAddress(email)
 	if err != nil {
 		return fmt.Errorf("invalid email address '%s': %w", email, err)

sasl-plain-access-token/segmentio/saslplainoauthmechanism/principal_email_test.go

@@ -46,19 +46,40 @@ func TestValidatePrincipalEmail(t *testing.T) {
 			email:     "human-principal@example.com",
 			expectErr: false,
 		},
+		{
+			name:      "GKE WIF Principal - Valid",
+			email:     "principal://iam.googleapis.com/projects/1234567891011/locations/global/workloadIdentityPools/my-project.svc.id.goog/subject/ns/custom-namespace/sa/gmktest",
+			expectErr: false,
+		},
+		{
+			name:      "Standard Workload Identity Principal - Valid",
+			email:     "principal://iam.googleapis.com/projects/12345678/locations/global/workloadIdentityPools/my-pool/subject/my-subject",
+			expectErr: false,
+		},
+		{
+			name:      "Standard Workforce Identity Principal - Valid",
+			email:     "principal://iam.googleapis.com/locations/global/workforcePools/altostrat-contractors/subject/raha@altostrat.com",
+			expectErr: false,
+		},
+
 		{
 			name:      "Empty email - Invalid",
 			email:     "",
 			expectErr: true,
 		},
 		{
-			name:      "Workload Identity Federation Pool - No iam.gke.io/return-principal-id-as-email KSA annotation - Invalid",
-			email:     "my-project.svc.id.goog",
+			name:      "Non RFC Email - Invalid",
+			email:     "human[]principaluser@example.com",
+			expectErr: true,
+		},
+		{
+			name:      "Bad  WIF Principal Bad Prefix - Invalid",
+			email:     "principal://iammmmmmmm.googleapis.com/projects/1234567891011/locations/global/workloadIdentityPools/my-project.svc.id.goog/subject/ns/custom-namespace/sa/gmktest",
 			expectErr: true,
 		},
 		{
-			name:      "Workload Identity Federation Principal - Invalid",
-			email:     "principal://iam.googleapis.com/projects/123456789/locations/global/workloadIdentityPools/my-project.svc.id.goog/subject/ns/default/sa/gmktest",
+			name:      "Workload Identity Federation Pool - No iam.gke.io/return-principal-id-as-email KSA annotation - Invalid",
+			email:     "my-project.svc.id.goog",
 			expectErr: true,
 		},
 	}