🐛 Fix setup MFA breaking due to grafana client ssm parameter (#1005)

Author: deifyed

Diff for: pkg/cognito/mfa.go

@@ -14,6 +14,8 @@ import (
 	"runtime"
 	"strings"
 
+	"github.com/aws/aws-sdk-go/aws/request"
+
 	"github.com/google/uuid"
 	qrcode "github.com/skip2/go-qrcode"
 
@@ -32,11 +34,13 @@ const (
 	// MFAOutputFormatQRCode represents showing the MFA details as a QR code
 	MFAOutputFormatQRCode = "qrcode"
 	// MFAOutputFormatText represents showing the MFA details as text
-	MFAOutputFormatText             = "text"
-	defaultOneTimePasswordType      = "TOTP"
-	defaultOneTimePasswordDigits    = 6
-	defaultOneTimePasswordAlgorithm = "SHA1"
-	defaultOneTimePasswordInterval  = 30
+	MFAOutputFormatText                  = "text"
+	defaultOneTimePasswordType           = "TOTP"
+	defaultOneTimePasswordDigits         = 6
+	defaultOneTimePasswordAlgorithm      = "SHA1"
+	defaultOneTimePasswordInterval       = 30
+	defaultRelevantUserPoolClientKeyword = "argocd"
+	defaultQRCodePixelSize               = 256
 )
 
 type userPoolClient struct {
@@ -127,24 +131,29 @@ func getRelevantUserPoolID(ctx context.Context, provider cognitoidentityprovider
 	return "", fmt.Errorf("no user pool found for cluster %s", cluster.Metadata.Name)
 }
 
-func getRelevantUserPoolClient(ctx context.Context, provider cognitoidentityprovideriface.CognitoIdentityProviderAPI, userPoolClientID string) (
+func getRelevantUserPoolClient(ctx context.Context, provider userpoolClientsLister, userPoolID string) (
 	cognitoidentityprovider.UserPoolClientDescription,
 	error,
 ) {
 	var nextToken *string
 
 	for {
 		listUserPoolsResult, err := provider.ListUserPoolClientsWithContext(ctx, &cognitoidentityprovider.ListUserPoolClientsInput{
-			MaxResults: aws.Int64(1),
+			MaxResults: aws.Int64(maximumMaximumUserPoolResults),
 			NextToken:  nextToken,
-			UserPoolId: aws.String(userPoolClientID),
+			UserPoolId: aws.String(userPoolID),
 		})
 		if err != nil {
 			return cognitoidentityprovider.UserPoolClientDescription{}, fmt.Errorf("listing user pools: %w", err)
 		}
 
 		for _, client := range listUserPoolsResult.UserPoolClients {
-			return *client, nil
+			// This should be any of the clients provisioned by okctl, but due to inconsistent naming of the Grafana client
+			// secret SSM parameter and the situation regarding golden path we'll settle on picking the ArgoCD client for
+			// now. This breaks MFA for environments without ArgoCD.
+			if strings.Contains(*client.ClientName, defaultRelevantUserPoolClientKeyword) {
+				return *client, nil
+			}
 		}
 
 		if listUserPoolsResult.NextToken == nil {
@@ -154,7 +163,7 @@ func getRelevantUserPoolClient(ctx context.Context, provider cognitoidentityprov
 		nextToken = listUserPoolsResult.NextToken
 	}
 
-	return cognitoidentityprovider.UserPoolClientDescription{}, fmt.Errorf("no clients found for user pool %s", userPoolClientID)
+	return cognitoidentityprovider.UserPoolClientDescription{}, fmt.Errorf("no clients found for user pool %s", userPoolID)
 }
 
 func getCognitoClientSecretForClient(ctx context.Context, provider ssmiface.SSMAPI, clientName string) (string, error) {
@@ -226,7 +235,7 @@ func generateDeviceSecretQRCode(cluster v1alpha1.Cluster, userEmail string, secr
 		defaultOneTimePasswordInterval,
 	)
 
-	err := qrcode.WriteFile(qrCodeURI, qrcode.Medium, 256, qrCodePath)
+	err := qrcode.WriteFile(qrCodeURI, qrcode.Medium, defaultQRCodePixelSize, qrCodePath)
 	if err != nil {
 		return "", fmt.Errorf("writing QR code: %w", err)
 	}
@@ -253,3 +262,11 @@ func openbrowser(url string) {
 		log.Fatal(err)
 	}
 }
+
+type userpoolClientsLister interface {
+	ListUserPoolClientsWithContext(
+		context.Context,
+		*cognitoidentityprovider.ListUserPoolClientsInput,
+		...request.Option,
+	) (*cognitoidentityprovider.ListUserPoolClientsOutput, error)
+}

Diff for: pkg/cognito/mfa_test.go

@@ -0,0 +1,82 @@
+package cognito
+
+import (
+	"context"
+	"testing"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/request"
+	"github.com/aws/aws-sdk-go/service/cognitoidentityprovider"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetRelevantUserPoolClient(t *testing.T) {
+	testCases := []struct {
+		name         string
+		withClients  []string
+		expectClient string
+	}{
+		{
+			name: "Should return the argocd client in an expected okctl setup",
+			withClients: []string{
+				"okctl-mock-cluster-argocd",
+				"okctl-mock-cluster-grafana",
+			},
+			expectClient: "okctl-mock-cluster-argocd",
+		},
+		{
+			name: "Should return the argocd client in a strange setup with a gazzilion clients",
+			withClients: []string{
+				"okctl-mock-cluster-client1",
+				"okctl-mock-cluster-grafana",
+				"okctl-mock-cluster-client2",
+				"okctl-mock-cluster-client3",
+				"okctl-mock-cluster-argocd",
+				"okctl-mock-cluster-client4",
+				"okctl-mock-cluster-client5",
+			},
+			expectClient: "okctl-mock-cluster-argocd",
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			result, err := getRelevantUserPoolClient(
+				context.Background(),
+				&mockCognitoIdentityProviderAPI{clients: tc.withClients},
+				"mock-userpool-id",
+			)
+			assert.NoError(t, err)
+
+			assert.Equal(t, tc.expectClient, *result.ClientName)
+		})
+	}
+}
+
+type mockCognitoIdentityProviderAPI struct {
+	clients []string
+}
+
+func (m mockCognitoIdentityProviderAPI) ListUserPoolClientsWithContext(
+	_ aws.Context,
+	_ *cognitoidentityprovider.ListUserPoolClientsInput,
+	_ ...request.Option,
+) (*cognitoidentityprovider.ListUserPoolClientsOutput, error) {
+	clients := make([]*cognitoidentityprovider.UserPoolClientDescription, len(m.clients))
+
+	for index, name := range m.clients {
+		clients[index] = &cognitoidentityprovider.UserPoolClientDescription{
+			ClientId:   aws.String(name),
+			ClientName: aws.String(name),
+			UserPoolId: aws.String("mock-user-pool-id"),
+		}
+	}
+
+	return &cognitoidentityprovider.ListUserPoolClientsOutput{
+		UserPoolClients: clients,
+	}, nil
+}