Generate smaller OIDC keys for unit tests

- significantly increases unit test performance by moving from 4096 -> 256 bit keys - preserves 4096 bit keys for all non-testing scenarios
Azure · Aug 31, 2024 · d6f776e · d6f776e
1 parent 710eb2b
commit d6f776e
Show file tree

Hide file tree

Showing 8 changed files with 58 additions and 20 deletions.
diff --git a/pkg/cluster/deploybaseresources_test.go b/pkg/cluster/deploybaseresources_test.go
@@ -1421,6 +1421,7 @@ func TestCreateOIDC(t *testing.T) {
 			},
 		},
 	}
+	testOIDCKeyBitSize := 256
 
 	for _, tt := range []struct {
 		name                              string
@@ -1481,6 +1482,7 @@ func TestCreateOIDC(t *testing.T) {
 			},
 			mocks: func(blob *mock_azblob.MockManager, menv *mock_env.MockInterface, azblobClient *mock_azblob.MockAZBlobClient) {
 				menv.EXPECT().FeatureIsSet(env.FeatureRequireOIDCStorageWebEndpoint).Return(false)
+				menv.EXPECT().OIDCKeyBitSize().Return(testOIDCKeyBitSize)
 				menv.EXPECT().OIDCEndpoint().Return(afdEndpoint)
 				menv.EXPECT().OIDCStorageAccountName().Return(oidcStorageAccountName)
 				menv.EXPECT().Environment().Return(&azureclient.PublicCloud)
@@ -1508,6 +1510,7 @@ func TestCreateOIDC(t *testing.T) {
 			mocks: func(blob *mock_azblob.MockManager, menv *mock_env.MockInterface, azblobClient *mock_azblob.MockAZBlobClient) {
 				menv.EXPECT().FeatureIsSet(env.FeatureRequireOIDCStorageWebEndpoint).Return(true)
 				menv.EXPECT().ResourceGroup().Return(resourceGroupName)
+				menv.EXPECT().OIDCKeyBitSize().Return(testOIDCKeyBitSize)
 				menv.EXPECT().OIDCStorageAccountName().AnyTimes().Return(oidcStorageAccountName)
 				blob.EXPECT().GetContainerProperties(gomock.Any(), resourceGroupName, oidcStorageAccountName, oidcbuilder.WebContainer).Return(containerProperties, nil)
 				menv.EXPECT().Environment().Return(&azureclient.PublicCloud)
@@ -1557,6 +1560,7 @@ func TestCreateOIDC(t *testing.T) {
 			},
 			mocks: func(blob *mock_azblob.MockManager, menv *mock_env.MockInterface, azblobClient *mock_azblob.MockAZBlobClient) {
 				menv.EXPECT().FeatureIsSet(env.FeatureRequireOIDCStorageWebEndpoint).Return(false)
+				menv.EXPECT().OIDCKeyBitSize().Return(testOIDCKeyBitSize)
 				menv.EXPECT().OIDCEndpoint().Return(afdEndpoint)
 				menv.EXPECT().OIDCStorageAccountName().Return(oidcStorageAccountName)
 				menv.EXPECT().Environment().Return(&azureclient.PublicCloud)
@@ -1581,6 +1585,7 @@ func TestCreateOIDC(t *testing.T) {
 			},
 			mocks: func(blob *mock_azblob.MockManager, menv *mock_env.MockInterface, azblobClient *mock_azblob.MockAZBlobClient) {
 				menv.EXPECT().FeatureIsSet(env.FeatureRequireOIDCStorageWebEndpoint).Return(false)
+				menv.EXPECT().OIDCKeyBitSize().Return(testOIDCKeyBitSize)
 				menv.EXPECT().OIDCEndpoint().Return(afdEndpoint)
 				menv.EXPECT().OIDCStorageAccountName().Return(oidcStorageAccountName)
 				menv.EXPECT().Environment().Return(&azureclient.PublicCloud)

diff --git a/pkg/env/env.go b/pkg/env/env.go
@@ -96,6 +96,7 @@ type Interface interface {
 	ACRDomain() string
 	OIDCStorageAccountName() string
 	OIDCEndpoint() string
+	OIDCKeyBitSize() int
 	AROOperatorImage() string
 	LiveConfig() liveconfig.Manager
 

diff --git a/pkg/env/prod.go b/pkg/env/prod.go
@@ -271,6 +271,10 @@ func (p *prod) OIDCEndpoint() string {
 	return fmt.Sprintf("https://%s/", os.Getenv("OIDC_AFD_ENDPOINT"))
 }
 
+func (p *prod) OIDCKeyBitSize() int {
+	return 4096
+}
+
 func (p *prod) AROOperatorImage() string {
 	return fmt.Sprintf("%s/aro:%s", p.acrDomain, version.GitCommit)
 }

diff --git a/pkg/util/mocks/env/env.go b/pkg/util/mocks/env/env.go
diff --git a/pkg/util/oidcbuilder/jwks.go b/pkg/util/oidcbuilder/jwks.go
@@ -15,13 +15,13 @@ import (
 
 	"github.com/pkg/errors"
 	"gopkg.in/go-jose/go-jose.v2"
-)
 
-func CreateKeyPair() (encPrivateKey []byte, encPublicKey []byte, err error) {
-	bitSize := 4096
+	"github.com/Azure/ARO-RP/pkg/env"
+)
 
+func CreateKeyPair(env env.Interface) (encPrivateKey []byte, encPublicKey []byte, err error) {
 	// Generate RSA keypair
-	privateKey, err := rsa.GenerateKey(rand.Reader, bitSize)
+	privateKey, err := rsa.GenerateKey(rand.Reader, env.OIDCKeyBitSize())
 	if err != nil {
 		return nil, nil, errors.Wrapf(err, "failed to generate private key")
 	}

diff --git a/pkg/util/oidcbuilder/jwks_test.go b/pkg/util/oidcbuilder/jwks_test.go
@@ -15,7 +15,7 @@ import (
 // Licensed under the Apache License 2.0.
 
 func TestKeyIDFromPublicKey(t *testing.T) {
-	privateKey, err := rsa.GenerateKey(rand.Reader, 4096)
+	privateKey, err := rsa.GenerateKey(rand.Reader, 256)
 	if err != nil {
 		t.Fatal(err)
 	}

diff --git a/pkg/util/oidcbuilder/oidcbuilder.go b/pkg/util/oidcbuilder/oidcbuilder.go
@@ -28,7 +28,7 @@ type OIDCBuilder struct {
 }
 
 func NewOIDCBuilder(env env.Interface, oidcEndpoint string, directoryName string) (*OIDCBuilder, error) {
-	privateKey, publicKey, err := CreateKeyPair()
+	privateKey, publicKey, err := CreateKeyPair(env)
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/util/oidcbuilder/oidcbuilder_test.go b/pkg/util/oidcbuilder/oidcbuilder_test.go
@@ -36,20 +36,7 @@ func TestEnsureOIDCDocs(t *testing.T) {
 	blobContainerURL := "fakeBlobContainerURL"
 	endpointURL := "fakeEndPointURL"
 
-	priKey, pubKey, err := CreateKeyPair()
-	if err != nil {
-		t.Fatal(err)
-	}
-	privateKey, err := rsa.GenerateKey(rand.Reader, 4096)
-	if err != nil {
-		t.Fatal(err)
-	}
-	pubKeyBytes := x509.MarshalPKCS1PublicKey(&privateKey.PublicKey)
-	incorrectlyEncodedPublicKey := pem.EncodeToMemory(&pem.Block{
-		Type:    "PUBLIC KEY",
-		Headers: nil,
-		Bytes:   pubKeyBytes,
-	})
+	priKey, pubKey, incorrectlyEncodedPublicKey := getTestKeyData(t)
 
 	nonRSAPrivateKey, _ := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
 	nonRSAPubKeyBytes, err := x509.MarshalPKIXPublicKey(&nonRSAPrivateKey.PublicKey)
@@ -182,6 +169,33 @@ func TestEnsureOIDCDocs(t *testing.T) {
 	}
 }
 
+func getTestKeyData(t *testing.T) ([]byte, []byte, []byte) {
+	t.Helper()
+
+	testKeyBitSize := 256
+
+	controller := gomock.NewController(t)
+	defer controller.Finish()
+
+	env := mock_env.NewMockInterface(controller)
+	env.EXPECT().OIDCKeyBitSize().Return(testKeyBitSize)
+	priKey, pubKey, err := CreateKeyPair(env)
+	if err != nil {
+		t.Fatal(err)
+	}
+	privateKey, err := rsa.GenerateKey(rand.Reader, testKeyBitSize)
+	if err != nil {
+		t.Fatal(err)
+	}
+	pubKeyBytes := x509.MarshalPKCS1PublicKey(&privateKey.PublicKey)
+	incorrectlyEncodedPublicKey := pem.EncodeToMemory(&pem.Block{
+		Type:    "PUBLIC KEY",
+		Headers: nil,
+		Bytes:   pubKeyBytes,
+	})
+	return priKey, pubKey, incorrectlyEncodedPublicKey
+}
+
 type fakeReadCloser struct {
 	io.Reader
 }