From 6d4c849442a359d039d96a302968ff3fc4e8ea65 Mon Sep 17 00:00:00 2001
From: Heba Elayoty <hebaelayoty@gmail.com>
Date: Fri, 20 Jan 2023 13:52:55 -0800
Subject: [PATCH] Implement credential for client auth

Signed-off-by: Heba Elayoty <hebaelayoty@gmail.com>
---
 pkg/auth/auth.go          | 31 ++++++++++++++++++++++++++++++-
 pkg/client/client_apis.go | 22 +++++++++++++++++-----
 pkg/provider/aci.go       |  5 ++++-
 3 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/pkg/auth/auth.go b/pkg/auth/auth.go
index fe38014c..712ed96b 100644
--- a/pkg/auth/auth.go
+++ b/pkg/auth/auth.go
@@ -11,8 +11,10 @@ import (
 	"strings"
 	"unicode/utf16"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	_ "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/cloud"
+	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
 	"github.com/Azure/go-autorest/autorest"
 	"github.com/Azure/go-autorest/autorest/adal"
 	"github.com/dimchansky/utfbom"
@@ -36,6 +38,34 @@ type Config struct {
 	Authorizer    autorest.Authorizer
 }
 
+// GetMSICredential retrieve MSI credential
+func (c *Config) GetMSICredential(ctx context.Context) (*azidentity.ManagedIdentityCredential, error) {
+	log.G(ctx).Debug("getting token using user identity")
+	opts := &azidentity.ManagedIdentityCredentialOptions{ID: azidentity.ClientID(c.AuthConfig.ClientID)}
+	msiCredential, err := azidentity.NewManagedIdentityCredential(opts)
+	if err != nil {
+		return nil, err
+	}
+
+	return msiCredential, nil
+}
+
+// GetSPCredential retrieve SP credential
+func (c *Config) GetSPCredential(ctx context.Context) (*azidentity.ClientSecretCredential, error) {
+	log.G(ctx).Debug("getting token using service principal")
+	opts := &azidentity.ClientSecretCredentialOptions{
+		ClientOptions: azcore.ClientOptions{
+			Cloud: c.Cloud,
+		},
+	}
+	spCredential, err := azidentity.NewClientSecretCredential(c.AuthConfig.TenantID, c.AuthConfig.ClientID, c.AuthConfig.ClientSecret, opts)
+	if err != nil {
+		return nil, err
+	}
+
+	return spCredential, nil
+}
+
 // getAuthorizer return autorest authorizer.
 func (c *Config) getAuthorizer(ctx context.Context, resource string) (autorest.Authorizer, error) {
 	var auth autorest.Authorizer
@@ -66,7 +96,6 @@ func (c *Config) getAuthorizer(ctx context.Context, resource string) (autorest.A
 			return nil, err
 		}
 	}
-
 	auth = autorest.NewBearerAuthorizer(token)
 	return auth, err
 }
diff --git a/pkg/client/client_apis.go b/pkg/client/client_apis.go
index cd8bfb98..cee0a554 100644
--- a/pkg/client/client_apis.go
+++ b/pkg/client/client_apis.go
@@ -10,7 +10,6 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
-	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
 	azaciv2 "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerinstance/armcontainerinstance/v2"
 	"github.com/pkg/errors"
 	"github.com/virtual-kubelet/azure-aci/pkg/auth"
@@ -40,10 +39,23 @@ type AzClientsAPIs struct {
 
 func NewAzClientsAPIs(ctx context.Context, azConfig auth.Config) (*AzClientsAPIs, error) {
 	obj := AzClientsAPIs{}
-	cred, err := azidentity.NewDefaultAzureCredential(nil)
+
+	// get credential
+	var credential azcore.TokenCredential
+	var err error
+	isUserIdentity := len(azConfig.AuthConfig.ClientID) == 0
+
+	if isUserIdentity {
+		credential, err = azConfig.GetMSICredential(ctx)
+
+	} else {
+		credential, err = azConfig.GetSPCredential(ctx)
+	}
 	if err != nil {
 		return nil, errors.Wrap(err, "an error has occurred while creating getting credential ")
 	}
+
+	// set aci user agent
 	ua := os.Getenv("ACI_EXTRA_USER_AGENT")
 	options := arm.ClientOptions{
 		ClientOptions: azcore.ClientOptions{
@@ -54,17 +66,17 @@ func NewAzClientsAPIs(ctx context.Context, azConfig auth.Config) (*AzClientsAPIs
 		},
 	}
 
-	cClient, err := azaciv2.NewContainersClient(azConfig.AuthConfig.SubscriptionID, cred, &options)
+	cClient, err := azaciv2.NewContainersClient(azConfig.AuthConfig.SubscriptionID, credential, &options)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create container client ")
 	}
 
-	cgClient, err := azaciv2.NewContainerGroupsClient(azConfig.AuthConfig.SubscriptionID, cred, &options)
+	cgClient, err := azaciv2.NewContainerGroupsClient(azConfig.AuthConfig.SubscriptionID, credential, &options)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create container group client ")
 	}
 
-	lClient, err := azaciv2.NewLocationClient(azConfig.AuthConfig.SubscriptionID, cred, &options)
+	lClient, err := azaciv2.NewLocationClient(azConfig.AuthConfig.SubscriptionID, credential, &options)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create location client ")
 	}
diff --git a/pkg/provider/aci.go b/pkg/provider/aci.go
index c9316ab8..b9e322b7 100644
--- a/pkg/provider/aci.go
+++ b/pkg/provider/aci.go
@@ -902,13 +902,16 @@ func (p *ACIProvider) verifyContainer(container *v1.Container) error {
 
 //this method is used for both initConainers and containers
 func (p *ACIProvider) getCommand(container v1.Container) []*string {
-	var command, args []*string
+	command := make([]*string, len(container.Command))
 	for c := range container.Command {
 		command[c] = &container.Command[c]
 	}
+
+	args := make([]*string, len(container.Command))
 	for a := range container.Args {
 		args[a] = &container.Args[a]
 	}
+
 	return append(command, args...)
 }