From 964f4bbd8d034fa997550878d37f49425a47dc7e Mon Sep 17 00:00:00 2001
From: Igor Beliakov <demtis.register@gmail.com>
Date: Thu, 17 Nov 2022 17:42:05 +0100
Subject: [PATCH] feat(AzureDNS): add a test for federated SPT

Signed-off-by: Igor Beliakov <demtis.register@gmail.com>
---
 pkg/issuer/acme/dns/azuredns/azuredns_test.go | 122 ++++++++++++++++++
 1 file changed, 122 insertions(+)

diff --git a/pkg/issuer/acme/dns/azuredns/azuredns_test.go b/pkg/issuer/acme/dns/azuredns/azuredns_test.go
index 50be54855..387e704ef 100644
--- a/pkg/issuer/acme/dns/azuredns/azuredns_test.go
+++ b/pkg/issuer/acme/dns/azuredns/azuredns_test.go
@@ -9,10 +9,16 @@ this directory.
 package azuredns
 
 import (
+	"encoding/json"
+	"io"
+	"net/http"
+	"net/http/httptest"
 	"os"
 	"testing"
 	"time"
 
+	"github.com/Azure/go-autorest/autorest/adal"
+	"github.com/Azure/go-autorest/autorest/azure"
 	v1 "github.com/cert-manager/cert-manager/pkg/apis/acme/v1"
 	"github.com/cert-manager/cert-manager/pkg/issuer/acme/dns/util"
 	"github.com/stretchr/testify/assert"
@@ -77,3 +83,119 @@ func TestInvalidAzureDns(t *testing.T) {
 	_, err := NewDNSProviderCredentials("invalid env", "cid", "secret", "", "", "", "", util.RecursiveNameservers, false, &v1.AzureManagedIdentity{})
 	assert.Error(t, err)
 }
+
+func populateFederatedToken(t *testing.T, filename string, content string) {
+	t.Helper()
+
+	f, err := os.Create(filename)
+	if err != nil {
+		assert.FailNow(t, err.Error())
+	}
+
+	if _, err := io.WriteString(f, content); err != nil {
+		assert.FailNow(t, err.Error())
+	}
+
+	if err := f.Close(); err != nil {
+		assert.FailNow(t, err.Error())
+	}
+}
+
+func TestGetAuthorizationFederatedSPT(t *testing.T) {
+	// Create a file that will be used to store a federated token
+	f, err := os.CreateTemp("", "")
+	if err != nil {
+		assert.FailNow(t, err.Error())
+	}
+	defer os.Remove(f.Name())
+
+	// Close the file to simplify logic within populateFederatedToken helper
+	if err := f.Close(); err != nil {
+		assert.FailNow(t, err.Error())
+	}
+
+	// The initial federated token is never used, so we don't care about the value yet
+	// Though, it's a requirement from adal to have a non-empty value set
+	populateFederatedToken(t, f.Name(), "random-jwt")
+
+	// Prepare environment variables adal will rely on. Skip changes for some envs if they are already defined (=live environment)
+	// Envs themselves are described here: https://azure.github.io/azure-workload-identity/docs/installation/mutating-admission-webhook.html
+	if os.Getenv("AZURE_TENANT_ID") == "" {
+		t.Setenv("AZURE_TENANT_ID", "fakeTenantID")
+	}
+
+	if os.Getenv("AZURE_CLIENT_ID") == "" {
+		t.Setenv("AZURE_CLIENT_ID", "fakeClientID")
+	}
+
+	t.Setenv("AZURE_FEDERATED_TOKEN_FILE", f.Name())
+
+	t.Run("token refresh", func(t *testing.T) {
+		// Basically, we want one token to be exchanged for the other (key and value respectively)
+		tokens := map[string]string{
+			"initialFederatedToken":   "initialAccessToken",
+			"refreshedFederatedToken": "refreshedAccessToken",
+		}
+
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if err := r.ParseForm(); err != nil {
+				assert.FailNow(t, err.Error())
+			}
+
+			w.Header().Set("Content-Type", "application/json")
+			receivedFederatedToken := r.FormValue("client_assertion")
+			accessToken := adal.Token{AccessToken: tokens[receivedFederatedToken]}
+
+			if err := json.NewEncoder(w).Encode(accessToken); err != nil {
+				assert.FailNow(t, err.Error())
+			}
+
+			// Expected format: http://<server>/<tenant-ID>/oauth2/token?api-version=1.0
+			assert.Contains(t, r.RequestURI, os.Getenv("AZURE_TENANT_ID"), "URI should contain the tenant ID exposed through env variable")
+
+			assert.Equal(t, os.Getenv("AZURE_CLIENT_ID"), r.FormValue("client_id"), "client_id should match the value exposed through env variable")
+		}))
+		defer ts.Close()
+
+		ambient := true
+		env := azure.Environment{ActiveDirectoryEndpoint: ts.URL, ResourceManagerEndpoint: ts.URL}
+		managedIdentity := &v1.AzureManagedIdentity{ClientID: ""}
+
+		spt, err := getAuthorization(env, "", "", "", "", ambient, managedIdentity)
+		assert.NoError(t, err)
+
+		for federatedToken, accessToken := range tokens {
+			populateFederatedToken(t, f.Name(), federatedToken)
+			assert.NoError(t, spt.Refresh(), "Token refresh failed")
+			assert.Equal(t, accessToken, spt.Token().AccessToken, "Access token should have been set to a value returned by the webserver")
+		}
+	})
+
+	t.Run("clientID overrides through managedIdentity section", func(t *testing.T) {
+		managedIdentity := &v1.AzureManagedIdentity{ClientID: "anotherClientID"}
+
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if err := r.ParseForm(); err != nil {
+				assert.FailNow(t, err.Error())
+			}
+
+			w.Header().Set("Content-Type", "application/json")
+			accessToken := adal.Token{AccessToken: "abc"}
+
+			if err := json.NewEncoder(w).Encode(accessToken); err != nil {
+				assert.FailNow(t, err.Error())
+			}
+
+			assert.Equal(t, managedIdentity.ClientID, r.FormValue("client_id"), "client_id should match the value passed through managedIdentity section")
+		}))
+		defer ts.Close()
+
+		ambient := true
+		env := azure.Environment{ActiveDirectoryEndpoint: ts.URL, ResourceManagerEndpoint: ts.URL}
+
+		spt, err := getAuthorization(env, "", "", "", "", ambient, managedIdentity)
+		assert.NoError(t, err)
+
+		assert.NoError(t, spt.Refresh(), "Token refresh failed")
+	})
+}