From 9ec0e7ec188d5ef9e41496448da45dee7b259446 Mon Sep 17 00:00:00 2001
From: Roshan Kumar <rok@microsoft.com>
Date: Fri, 9 Mar 2018 19:14:25 +0530
Subject: [PATCH 01/13] msi for native support

---
 .../VstsAzureHelpers_/InitializeFunctions.ps1     | 15 +++++++++++++++
 .../VstsAzureHelpers_/VstsAzureHelpers_.psm1      |  2 +-
 Tasks/Common/VstsAzureHelpers_/module.json        |  3 ++-
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1 b/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
index fa7bc8a74941..0454f4aa3448 100644
--- a/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
+++ b/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
@@ -123,6 +123,21 @@ function Initialize-AzureSubscription {
                 throw (New-Object System.Exception((Get-VstsLocString -Key AZ_ServicePrincipalError), $_.Exception))
             }
 
+            Set-CurrentAzureSubscription -SubscriptionId $Endpoint.Data.SubscriptionId -StorageAccount $StorageAccount
+        }elseif ($Endpoint.Auth.Scheme -eq 'MSI') {
+            $msiUri =  "http://localhost:"+$Endpoint.Data.MsiPort+"/oauth2/token"
+            $response = Invoke-WebRequest -Uri $msiUri -Method GET -Body @{resource= $Endpoint.Url} -Headers @{Metadata="true"}
+            $content =$response.Content | ConvertFrom-Json
+            $access_token = $content.access_token
+            $accountId = "MSI@50342"
+            try {
+                Write-Host "##[command]Add-AzureAccount  -AccessToken $access_token -AccountId $accountId "
+                $null = Add-AzureAccount -AccessToken $access_token -AccountId $accountId
+            } catch {
+                # Provide an additional, custom, credentials-related error message.
+                Write-VstsTaskError -Message $_.Exception.Message
+                throw (New-Object System.Exception((Get-VstsLocString -Key AZ_ManagedServiceIdentityError), $_.Exception))
+            }
             Set-CurrentAzureSubscription -SubscriptionId $Endpoint.Data.SubscriptionId -StorageAccount $StorageAccount
         } elseif ($script:azureModule) {
             # Throw if >=0.9.9 Azure.
diff --git a/Tasks/Common/VstsAzureHelpers_/VstsAzureHelpers_.psm1 b/Tasks/Common/VstsAzureHelpers_/VstsAzureHelpers_.psm1
index 398213f98d85..6e93088482cc 100644
--- a/Tasks/Common/VstsAzureHelpers_/VstsAzureHelpers_.psm1
+++ b/Tasks/Common/VstsAzureHelpers_/VstsAzureHelpers_.psm1
@@ -34,7 +34,7 @@ function Initialize-Azure {
 
         # Determine which modules are preferred.
         $preferredModules = @( )
-        if ($endpoint.Auth.Scheme -eq 'ServicePrincipal') {
+        if (($endpoint.Auth.Scheme -eq 'ServicePrincipal') -or ($endpoint.Auth.Scheme -eq 'MSI')) {
             $preferredModules += 'AzureRM'
         } elseif ($endpoint.Auth.Scheme -eq 'UserNamePassword') {
             $preferredModules += 'Azure'
diff --git a/Tasks/Common/VstsAzureHelpers_/module.json b/Tasks/Common/VstsAzureHelpers_/module.json
index e770dec7b110..72ab66fc78af 100644
--- a/Tasks/Common/VstsAzureHelpers_/module.json
+++ b/Tasks/Common/VstsAzureHelpers_/module.json
@@ -9,6 +9,7 @@
         "AZ_ServicePrincipalAuthNotSupportedAzureVersion0": "Service principal authentication is not supported in version '{0}' of the Azure module.",
         "AZ_UnsupportedAuthScheme0": "Unsupported authentication scheme '{0}' for Azure endpoint.",
         "AZ_AvailableModules": "The list of available {0} modules:",
-        "AZ_InvalidARMEndpoint": "Specified AzureRM endpoint is invalid."
+        "AZ_InvalidARMEndpoint": "Specified AzureRM endpoint is invalid.",
+        "AZ_ManagedServiceIdentityError": "There was an error with the Managed Service Principal. Please configure MSI for vm 'https://docs.microsoft.com/en-us/azure/active-directory/managed-service-identity/qs-configure-portal-windows-vm'"
     }
 }

From 79199a76a690ae8103f01cb254da8a3433c9d5c2 Mon Sep 17 00:00:00 2001
From: Roshan Kumar <rok@microsoft.com>
Date: Tue, 27 Mar 2018 14:58:19 +0530
Subject: [PATCH 02/13] native powershell changes

---
 Tasks/AzurePowerShell/Utility.ps1             |  2 +-
 .../VstsAzureHelpers_/InitializeFunctions.ps1 | 42 ++++++++++++-------
 2 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/Tasks/AzurePowerShell/Utility.ps1 b/Tasks/AzurePowerShell/Utility.ps1
index a7e67876b6fd..51709642fd07 100644
--- a/Tasks/AzurePowerShell/Utility.ps1
+++ b/Tasks/AzurePowerShell/Utility.ps1
@@ -30,7 +30,7 @@ function Update-PSModulePathForHostedAgent {
             $hostedAgentAzureModulePath  =  Get-LatestModule -patternToMatch "^azure_[0-9]+\.[0-9]+\.[0-9]+$"   -patternToExtract "[0-9]+\.[0-9]+\.[0-9]+$" -Classic:$true
         }
 
-        if($authScheme -eq 'ServicePrincipal' -or $authScheme -eq '')
+        if($authScheme -eq 'ServicePrincipal' -or $authScheme -eq '' -or $authScheme -eq 'ManagedServiceIdentity')
         {
             $env:PSModulePath = $hostedAgentAzureModulePath + ";" + $env:PSModulePath
             $env:PSModulePath = $env:PSModulePath.TrimStart(';')
diff --git a/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1 b/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
index 0454f4aa3448..cd5cbcaef394 100644
--- a/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
+++ b/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
@@ -123,21 +123,6 @@ function Initialize-AzureSubscription {
                 throw (New-Object System.Exception((Get-VstsLocString -Key AZ_ServicePrincipalError), $_.Exception))
             }
 
-            Set-CurrentAzureSubscription -SubscriptionId $Endpoint.Data.SubscriptionId -StorageAccount $StorageAccount
-        }elseif ($Endpoint.Auth.Scheme -eq 'MSI') {
-            $msiUri =  "http://localhost:"+$Endpoint.Data.MsiPort+"/oauth2/token"
-            $response = Invoke-WebRequest -Uri $msiUri -Method GET -Body @{resource= $Endpoint.Url} -Headers @{Metadata="true"}
-            $content =$response.Content | ConvertFrom-Json
-            $access_token = $content.access_token
-            $accountId = "MSI@50342"
-            try {
-                Write-Host "##[command]Add-AzureAccount  -AccessToken $access_token -AccountId $accountId "
-                $null = Add-AzureAccount -AccessToken $access_token -AccountId $accountId
-            } catch {
-                # Provide an additional, custom, credentials-related error message.
-                Write-VstsTaskError -Message $_.Exception.Message
-                throw (New-Object System.Exception((Get-VstsLocString -Key AZ_ManagedServiceIdentityError), $_.Exception))
-            }
             Set-CurrentAzureSubscription -SubscriptionId $Endpoint.Data.SubscriptionId -StorageAccount $StorageAccount
         } elseif ($script:azureModule) {
             # Throw if >=0.9.9 Azure.
@@ -163,7 +148,32 @@ function Initialize-AzureSubscription {
 
             Set-CurrentAzureRMSubscription -SubscriptionId $Endpoint.Data.SubscriptionId -TenantId $Endpoint.Auth.Parameters.TenantId
         }
-    } else {
+    } elseif ($Endpoint.Auth.Scheme -eq 'ManagedServiceIdentity') {
+        $accountId = $env:BUILD_BUILDID 
+        if($env:RELEASE_RELEASEID){
+            $accountId = $env:RELEASE_RELEASEID 
+        }
+        $date = Get-Date -Format o
+        $accountId = -join($accountId, "-", $date)
+        $port = 50342
+        if($Endpoint.Data.MsiPort){
+            $port = $Endpoint.Data.MsiPort
+        }
+        $msiUri =  "http://localhost:$port/oauth2/token"
+        $response = Invoke-WebRequest -Uri $msiUri -Method GET -Body @{resource= $Endpoint.Url} -Headers @{Metadata="true"} -UseBasicParsing
+        $content =$response.Content | ConvertFrom-Json
+        $access_token = $content.access_token
+        try {
+            Write-Host "##[command]Add-AzureRmAccount  -AccessToken $access_token -AccountId $accountId "
+            $null = Add-AzureRmAccount -AccessToken $access_token -AccountId $accountId
+        } catch {
+            # Provide an additional, custom, credentials-related error message.
+            Write-VstsTaskError -Message $_.Exception.Message
+            throw (New-Object System.Exception((Get-VstsLocString -Key AZ_ManagedServiceIdentityError), $_.Exception))
+        }
+        
+        Set-CurrentAzureRMSubscription -SubscriptionId $Endpoint.Data.SubscriptionId -TenantId $Endpoint.Auth.Parameters.TenantId
+    }else {
         throw (Get-VstsLocString -Key AZ_UnsupportedAuthScheme0 -ArgumentList $Endpoint.Auth.Scheme)
     }
 }

From b1555f24c58cc73ab54d40583380f119034ffd40 Mon Sep 17 00:00:00 2001
From: Roshan Kumar <rok@microsoft.com>
Date: Thu, 29 Mar 2018 16:16:28 +0530
Subject: [PATCH 03/13] Not printing access token in log

---
 Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1 | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1 b/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
index cd5cbcaef394..f50b8e76c8ff 100644
--- a/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
+++ b/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
@@ -164,7 +164,7 @@ function Initialize-AzureSubscription {
         $content =$response.Content | ConvertFrom-Json
         $access_token = $content.access_token
         try {
-            Write-Host "##[command]Add-AzureRmAccount  -AccessToken $access_token -AccountId $accountId "
+            Write-Host "##[command]Add-AzureRmAccount  -AccessToken ****** -AccountId $accountId "
             $null = Add-AzureRmAccount -AccessToken $access_token -AccountId $accountId
         } catch {
             # Provide an additional, custom, credentials-related error message.

From a28efdb50d8f4d104ea469515cb1d1bf4fb5a459 Mon Sep 17 00:00:00 2001
From: Roshan Kumar <rok@microsoft.com>
Date: Wed, 4 Apr 2018 11:04:09 +0530
Subject: [PATCH 04/13] rename

---
 Tasks/Common/VstsAzureHelpers_/VstsAzureHelpers_.psm1 | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Tasks/Common/VstsAzureHelpers_/VstsAzureHelpers_.psm1 b/Tasks/Common/VstsAzureHelpers_/VstsAzureHelpers_.psm1
index 6e93088482cc..1fa3f0842f64 100644
--- a/Tasks/Common/VstsAzureHelpers_/VstsAzureHelpers_.psm1
+++ b/Tasks/Common/VstsAzureHelpers_/VstsAzureHelpers_.psm1
@@ -34,7 +34,7 @@ function Initialize-Azure {
 
         # Determine which modules are preferred.
         $preferredModules = @( )
-        if (($endpoint.Auth.Scheme -eq 'ServicePrincipal') -or ($endpoint.Auth.Scheme -eq 'MSI')) {
+        if (($endpoint.Auth.Scheme -eq 'ServicePrincipal') -or ($endpoint.Auth.Scheme -eq 'ManagedServiceIdentity')) {
             $preferredModules += 'AzureRM'
         } elseif ($endpoint.Auth.Scheme -eq 'UserNamePassword') {
             $preferredModules += 'Azure'

From b8b3e66ce67e519e6b031761f03ae814d31c5df9 Mon Sep 17 00:00:00 2001
From: Roshan Kumar <rok@microsoft.com>
Date: Wed, 4 Apr 2018 14:08:32 +0530
Subject: [PATCH 05/13] Error messgae changed

---
 .../resources.resjson/en-US/resources.resjson |  3 +-
 ...ureSubscription.ManagedServiceIdentity.ps1 | 54 +++++++++++++++++++
 Tasks/Common/VstsAzureHelpers_/Tests/L0.ts    |  3 ++
 Tasks/Common/VstsAzureHelpers_/module.json    |  2 +-
 4 files changed, 60 insertions(+), 2 deletions(-)
 create mode 100644 Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1

diff --git a/Tasks/Common/VstsAzureHelpers_/Strings/resources.resjson/en-US/resources.resjson b/Tasks/Common/VstsAzureHelpers_/Strings/resources.resjson/en-US/resources.resjson
index abbd1d8ea127..d4641f2f896a 100644
--- a/Tasks/Common/VstsAzureHelpers_/Strings/resources.resjson/en-US/resources.resjson
+++ b/Tasks/Common/VstsAzureHelpers_/Strings/resources.resjson/en-US/resources.resjson
@@ -8,5 +8,6 @@
   "loc.messages.AZ_ServicePrincipalAuthNotSupportedAzureVersion0": "Service principal authentication is not supported in version '{0}' of the Azure module.",
   "loc.messages.AZ_UnsupportedAuthScheme0": "Unsupported authentication scheme '{0}' for Azure endpoint.",
   "loc.messages.AZ_AvailableModules": "The list of available {0} modules:",
-  "loc.messages.AZ_InvalidARMEndpoint": "Specified AzureRM endpoint is invalid."
+  "loc.messages.AZ_InvalidARMEndpoint": "Specified AzureRM endpoint is invalid.",
+  "loc.messages.AZ_ManagedServiceIdentityError": "There was an error with the Managed Service Principal. Please configure MSI for vm 'https://docs.microsoft.com/en-us/azure/active-directory/managed-service-identity/qs-configure-portal-windows-vm'"
 }
\ No newline at end of file
diff --git a/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1 b/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1
new file mode 100644
index 000000000000..1b482749f6f7
--- /dev/null
+++ b/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1
@@ -0,0 +1,54 @@
+[CmdletBinding()]
+param()
+
+# Arrange.
+. $PSScriptRoot\..\..\..\..\Tests\lib\Initialize-Test.ps1
+Microsoft.PowerShell.Core\Import-Module Microsoft.PowerShell.Security
+$module = Microsoft.PowerShell.Core\Import-Module $PSScriptRoot\.. -PassThru
+
+$endpoint = @{
+    Auth = @{
+        Parameters = @{
+            ServicePrincipalId = 'Some service principal ID'
+            ServicePrincipalKey = 'Some service principal key'
+            TenantId = 'Some tenant ID'
+        }
+        Scheme = 'ManagedServiceIdentity'
+    }
+    Data = @{
+        SubscriptionId = 'Some subscription ID'
+        SubscriptionName = 'Some subscription name'
+    }
+}
+
+$content = @"
+	{
+	"Content": 
+           {"access_token" : "Dummy Token" }
+	}
+"@
+
+$variableSets = @(
+    @{ StorageAccount = 'Some storage account' }
+)
+foreach ($variableSet in $variableSets) {
+    Write-Verbose ('-' * 80)
+    Unregister-Mock Add-AzureRMAccount
+    Unregister-Mock Set-CurrentAzureRMSubscription
+    Unregister-Mock Invoke-WebRequest
+    Unregister-Mock Set-UserAgent
+    Register-Mock Add-AzureRMAccount { 'some output' }
+    Register-Mock Set-CurrentAzureRMSubscription
+    Register-Mock Set-UserAgent
+    Register-Mock Invoke-WebRequest { @{Content = $content} } 
+    & $module {
+        $script:azureModule = $null
+        $script:azureRMProfileModule = @{ Version = [version]'1.2.3.4' }
+    }
+    
+    # Act.
+    $result = & $module Initialize-AzureSubscription -Endpoint $endpoint -StorageAccount $variableSet.StorageAccount
+
+    Assert-AreEqual $null $result
+    Assert-WasCalled Set-CurrentAzureRMSubscription -- -SubscriptionId $endpoint.Data.SubscriptionId -TenantId $endpoint.Auth.Parameters.TenantId
+}
\ No newline at end of file
diff --git a/Tasks/Common/VstsAzureHelpers_/Tests/L0.ts b/Tasks/Common/VstsAzureHelpers_/Tests/L0.ts
index ec8b910669a7..1bf2b7b9b098 100644
--- a/Tasks/Common/VstsAzureHelpers_/Tests/L0.ts
+++ b/Tasks/Common/VstsAzureHelpers_/Tests/L0.ts
@@ -63,6 +63,9 @@ describe('Common-VstsAzureHelpers_ Suite', function () {
         it('(Initialize-Azure) throws when service name is null', (done) => {
             psr.run(path.join(__dirname, 'Initialize-Azure.ThrowsWhenServiceNameIsNull.ps1'), done);
         })
+        it('(Initialize-AzureSubscription) manged service identity should pass ', (done) => {
+            psr.run(path.join(__dirname, 'Initialize-AzureSubscription.ManagedServiceIdentity.ps1'), done);
+        })
         it('(Initialize-AzureSubscription) passes values when cert auth', (done) => {
             psr.run(path.join(__dirname, 'Initialize-AzureSubscription.PassesValuesWhenCertAuth.ps1'), done);
         })
diff --git a/Tasks/Common/VstsAzureHelpers_/module.json b/Tasks/Common/VstsAzureHelpers_/module.json
index 72ab66fc78af..8664dbadeadd 100644
--- a/Tasks/Common/VstsAzureHelpers_/module.json
+++ b/Tasks/Common/VstsAzureHelpers_/module.json
@@ -10,6 +10,6 @@
         "AZ_UnsupportedAuthScheme0": "Unsupported authentication scheme '{0}' for Azure endpoint.",
         "AZ_AvailableModules": "The list of available {0} modules:",
         "AZ_InvalidARMEndpoint": "Specified AzureRM endpoint is invalid.",
-        "AZ_ManagedServiceIdentityError": "There was an error with the Managed Service Principal. Please configure MSI for vm 'https://docs.microsoft.com/en-us/azure/active-directory/managed-service-identity/qs-configure-portal-windows-vm'"
+        "AZ_ManagedServiceIdentityError": "There was an error with the Managed Service Principal. Please configure Mnaged Service Identity (MSI) for virtual machine 'https://aka.ms/azure-msi-docs'"
     }
 }

From 220491b18478414464f1148f1dc5b5839e60b884 Mon Sep 17 00:00:00 2001
From: Roshan Kumar <rok@microsoft.com>
Date: Wed, 4 Apr 2018 14:10:43 +0530
Subject: [PATCH 06/13] error message changed

---
 .../Strings/resources.resjson/en-US/resources.resjson           | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Tasks/Common/VstsAzureHelpers_/Strings/resources.resjson/en-US/resources.resjson b/Tasks/Common/VstsAzureHelpers_/Strings/resources.resjson/en-US/resources.resjson
index d4641f2f896a..51d327826e5e 100644
--- a/Tasks/Common/VstsAzureHelpers_/Strings/resources.resjson/en-US/resources.resjson
+++ b/Tasks/Common/VstsAzureHelpers_/Strings/resources.resjson/en-US/resources.resjson
@@ -9,5 +9,5 @@
   "loc.messages.AZ_UnsupportedAuthScheme0": "Unsupported authentication scheme '{0}' for Azure endpoint.",
   "loc.messages.AZ_AvailableModules": "The list of available {0} modules:",
   "loc.messages.AZ_InvalidARMEndpoint": "Specified AzureRM endpoint is invalid.",
-  "loc.messages.AZ_ManagedServiceIdentityError": "There was an error with the Managed Service Principal. Please configure MSI for vm 'https://docs.microsoft.com/en-us/azure/active-directory/managed-service-identity/qs-configure-portal-windows-vm'"
+  "loc.messages.AZ_ManagedServiceIdentityError": "There was an error with the Managed Service Principal. Please configure Mnaged Service Identity (MSI) for virtual machine 'https://aka.ms/azure-msi-docs'"
 }
\ No newline at end of file

From 89aaae120b22e78ea86574739c6477fdfef4bf2f Mon Sep 17 00:00:00 2001
From: Roshan Kumar <rok@microsoft.com>
Date: Wed, 4 Apr 2018 14:11:42 +0530
Subject: [PATCH 07/13] unused code

---
 .../Initialize-AzureSubscription.ManagedServiceIdentity.ps1 | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1 b/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1
index 1b482749f6f7..fb22acd49e5d 100644
--- a/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1
+++ b/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1
@@ -40,11 +40,7 @@ foreach ($variableSet in $variableSets) {
     Register-Mock Add-AzureRMAccount { 'some output' }
     Register-Mock Set-CurrentAzureRMSubscription
     Register-Mock Set-UserAgent
-    Register-Mock Invoke-WebRequest { @{Content = $content} } 
-    & $module {
-        $script:azureModule = $null
-        $script:azureRMProfileModule = @{ Version = [version]'1.2.3.4' }
-    }
+    Register-Mock Invoke-WebRequest { @{Content = $content} }
     
     # Act.
     $result = & $module Initialize-AzureSubscription -Endpoint $endpoint -StorageAccount $variableSet.StorageAccount

From 563c8ddfb5b3a5ddb0fb3469220bb8a26b84e829 Mon Sep 17 00:00:00 2001
From: Roshan Kumar <rok@microsoft.com>
Date: Wed, 4 Apr 2018 14:39:01 +0530
Subject: [PATCH 08/13] error message changed

---
 .../Strings/resources.resjson/en-US/resources.resjson          | 2 +-
 .../Initialize-AzureSubscription.ManagedServiceIdentity.ps1    | 3 ---
 Tasks/Common/VstsAzureHelpers_/module.json                     | 2 +-
 3 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/Tasks/Common/VstsAzureHelpers_/Strings/resources.resjson/en-US/resources.resjson b/Tasks/Common/VstsAzureHelpers_/Strings/resources.resjson/en-US/resources.resjson
index 51d327826e5e..74abde44f7f1 100644
--- a/Tasks/Common/VstsAzureHelpers_/Strings/resources.resjson/en-US/resources.resjson
+++ b/Tasks/Common/VstsAzureHelpers_/Strings/resources.resjson/en-US/resources.resjson
@@ -9,5 +9,5 @@
   "loc.messages.AZ_UnsupportedAuthScheme0": "Unsupported authentication scheme '{0}' for Azure endpoint.",
   "loc.messages.AZ_AvailableModules": "The list of available {0} modules:",
   "loc.messages.AZ_InvalidARMEndpoint": "Specified AzureRM endpoint is invalid.",
-  "loc.messages.AZ_ManagedServiceIdentityError": "There was an error with the Managed Service Principal. Please configure Mnaged Service Identity (MSI) for virtual machine 'https://aka.ms/azure-msi-docs'"
+  "loc.messages.AZ_ManagedServiceIdentityError": "There was an error with the Managed Service Principal. Please configure Managed Service Identity (MSI) for virtual machine 'https://aka.ms/azure-msi-docs'"
 }
\ No newline at end of file
diff --git a/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1 b/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1
index fb22acd49e5d..049b68bc7fc7 100644
--- a/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1
+++ b/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1
@@ -22,10 +22,7 @@ $endpoint = @{
 }
 
 $content = @"
-	{
-	"Content": 
            {"access_token" : "Dummy Token" }
-	}
 "@
 
 $variableSets = @(
diff --git a/Tasks/Common/VstsAzureHelpers_/module.json b/Tasks/Common/VstsAzureHelpers_/module.json
index 8664dbadeadd..c6b555e2cc11 100644
--- a/Tasks/Common/VstsAzureHelpers_/module.json
+++ b/Tasks/Common/VstsAzureHelpers_/module.json
@@ -10,6 +10,6 @@
         "AZ_UnsupportedAuthScheme0": "Unsupported authentication scheme '{0}' for Azure endpoint.",
         "AZ_AvailableModules": "The list of available {0} modules:",
         "AZ_InvalidARMEndpoint": "Specified AzureRM endpoint is invalid.",
-        "AZ_ManagedServiceIdentityError": "There was an error with the Managed Service Principal. Please configure Mnaged Service Identity (MSI) for virtual machine 'https://aka.ms/azure-msi-docs'"
+        "AZ_ManagedServiceIdentityError": "There was an error with the Managed Service Principal. Please configure Managed Service Identity (MSI) for virtual machine 'https://aka.ms/azure-msi-docs'"
     }
 }

From 08aa73aa9492d006016856802d93712888774c59 Mon Sep 17 00:00:00 2001
From: Roshan Kumar <rok@microsoft.com>
Date: Wed, 4 Apr 2018 15:58:43 +0530
Subject: [PATCH 09/13] build failure

---
 .../Initialize-AzureSubscription.ManagedServiceIdentity.ps1     | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1 b/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1
index 049b68bc7fc7..96b9aed57630 100644
--- a/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1
+++ b/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1
@@ -4,6 +4,8 @@ param()
 # Arrange.
 . $PSScriptRoot\..\..\..\..\Tests\lib\Initialize-Test.ps1
 Microsoft.PowerShell.Core\Import-Module Microsoft.PowerShell.Security
+Unregister-Mock Import-Module
+Register-Mock Write-VstsTaskError
 $module = Microsoft.PowerShell.Core\Import-Module $PSScriptRoot\.. -PassThru
 
 $endpoint = @{

From 93d616bb610959635d42f2ce33b8084eace279c1 Mon Sep 17 00:00:00 2001
From: Roshan Kumar <rok@microsoft.com>
Date: Thu, 12 Apr 2018 14:48:20 +0530
Subject: [PATCH 10/13] url changed

---
 Tasks/AzurePowerShell/Utility.ps1             |  2 +-
 .../VstsAzureHelpers_/InitializeFunctions.ps1 | 76 ++++++++++++++++---
 .../resources.resjson/en-US/resources.resjson |  4 +-
 Tasks/Common/VstsAzureHelpers_/module.json    |  4 +-
 4 files changed, 73 insertions(+), 13 deletions(-)

diff --git a/Tasks/AzurePowerShell/Utility.ps1 b/Tasks/AzurePowerShell/Utility.ps1
index 51709642fd07..b385ac9ba042 100644
--- a/Tasks/AzurePowerShell/Utility.ps1
+++ b/Tasks/AzurePowerShell/Utility.ps1
@@ -30,7 +30,7 @@ function Update-PSModulePathForHostedAgent {
             $hostedAgentAzureModulePath  =  Get-LatestModule -patternToMatch "^azure_[0-9]+\.[0-9]+\.[0-9]+$"   -patternToExtract "[0-9]+\.[0-9]+\.[0-9]+$" -Classic:$true
         }
 
-        if($authScheme -eq 'ServicePrincipal' -or $authScheme -eq '' -or $authScheme -eq 'ManagedServiceIdentity')
+        if($authScheme -eq 'ServicePrincipal' -or $authScheme -eq 'ManagedServiceIdentity' -or $authScheme -eq '')
         {
             $env:PSModulePath = $hostedAgentAzureModulePath + ";" + $env:PSModulePath
             $env:PSModulePath = $env:PSModulePath.TrimStart(';')
diff --git a/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1 b/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
index 2d1933be99fc..a1db49171a56 100644
--- a/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
+++ b/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
@@ -168,29 +168,85 @@ function Initialize-AzureSubscription {
         }
         $date = Get-Date -Format o
         $accountId = -join($accountId, "-", $date)
-        $port = 50342
-        if($Endpoint.Data.MsiPort){
-            $port = $Endpoint.Data.MsiPort
-        }
-        $msiUri =  "http://localhost:$port/oauth2/token"
-        $response = Invoke-WebRequest -Uri $msiUri -Method GET -Body @{resource= $Endpoint.Url} -Headers @{Metadata="true"} -UseBasicParsing
-        $content =$response.Content | ConvertFrom-Json
-        $access_token = $content.access_token
+        $access_token = Get-MsiAccessToken $Endpoint 0 0
         try {
             Write-Host "##[command]Add-AzureRmAccount  -AccessToken ****** -AccountId $accountId "
             $null = Add-AzureRmAccount -AccessToken $access_token -AccountId $accountId
         } catch {
             # Provide an additional, custom, credentials-related error message.
             Write-VstsTaskError -Message $_.Exception.Message
-            throw (New-Object System.Exception((Get-VstsLocString -Key AZ_ManagedServiceIdentityError), $_.Exception))
+            throw (New-Object System.Exception((Get-VstsLocString -Key AZ_MsiFailure), $_.Exception))
         }
         
         Set-CurrentAzureRMSubscription -SubscriptionId $Endpoint.Data.SubscriptionId -TenantId $Endpoint.Auth.Parameters.TenantId
     }else {
-        throw (Get-VstsLocString -Key AZ_UnsupportedAuthScheme0 -ArgumentList $Endpoint.Auth.Scheme)
+        throw (Get-VstsLocString-Key AZ_UnsupportedAuthScheme0 -ArgumentList $Endpoint.Auth.Scheme)
+    } 
+}
+
+# Get the Bearer Access Token from the MSI Endpoint
+function Get-MsiAccessToken {
+    [CmdletBinding()]
+    param([Parameter(Mandatory=$true)] $endpoint,
+        [Parameter(Mandatory=$true)] $count,
+        [Parameter(Mandatory=$true)] $timeToWait)
+
+    $msiClientId = "";
+    if($endpoint.Data.msiClientId){
+        $msiClientId  =  "&client_id=" + $endpoint.Data.msiClientId;
+    }
+    $tenantId = $endpoint.Auth.Parameters.TenantId
+
+    # Prepare contents for GET
+    $method = "GET"
+    $apiVersion = "2018-02-01";
+    let msiClientId =  $endpoint.Data.msiClientId ? "&client_id=" + this.msiClientId : 
+    $authUri = "http://169.254.169.254/metadata/identity/oauth2/token?api-version=" + $apiVersion + "&resource=" + $endpoint.Url + $msiClientId;
+    
+    # Call Rest API to fetch AccessToken
+    Write-Verbose "Fetching Access Token For MSI"
+    
+    try
+    {
+        $retryLimit = 5;
+        $waitedTime = 2000 + $timeToWait * 2;
+        $proxyUri = Get-ProxyUri $authUri
+        if ($proxyUri -eq $null)
+        {
+            Write-Verbose "No proxy settings"
+            $response = Invoke-WebRequest -Uri $authUri -Method $method -Headers @{Metadata="true"} -UseBasicParsing
+        }
+        else
+        {
+            Write-Verbose "Using Proxy settings"
+            $response = Invoke-WebRequest -Uri $authUri -Method $method -Headers @{Metadata="true"} -UseDefaultCredentials -Proxy $proxyUri -ProxyUseDefaultCredentials -UseBasicParsing
+        }
+
+        if($response.StatusCode == 209 -or $response.StatusCode == 500)
+        {
+            if(count -lt $retryLimit)
+            {
+                count += 1
+                Get-MsiAccessToken $endpoint count  $waitedTime
+            }
+            else
+            {
+                throw (Get-VstsLocString -Key AZ_MsiAccessTokenFetchFailure -ArgumentList $response.StatusCode $response.StatusDescription)
+            }
+        }
+        
+        $content = $response.Content | ConvertFrom-Json
+        return $content.access_token
+    }
+    catch
+    {
+        $exceptionMessage = $_.Exception.Message.ToString()
+        Write-Verbose "ExceptionMessage: $exceptionMessage (in function: Get-MsiAccessToken)"
+        throw (Get-VstsLocString -Key AZ_MsiAccessNotConfiguredProperlyFailure  -ArgumentList $response.StatusCode $response.StatusDescription)
     }
 }
 
+
 function Set-CurrentAzureSubscription {
     [CmdletBinding()]
     param(
diff --git a/Tasks/Common/VstsAzureHelpers_/Strings/resources.resjson/en-US/resources.resjson b/Tasks/Common/VstsAzureHelpers_/Strings/resources.resjson/en-US/resources.resjson
index 8d6de8e46c86..e7b175ca7567 100644
--- a/Tasks/Common/VstsAzureHelpers_/Strings/resources.resjson/en-US/resources.resjson
+++ b/Tasks/Common/VstsAzureHelpers_/Strings/resources.resjson/en-US/resources.resjson
@@ -9,5 +9,7 @@
   "loc.messages.AZ_UnsupportedAuthScheme0": "Unsupported authentication scheme '{0}' for Azure endpoint.",
   "loc.messages.AZ_AvailableModules": "The list of available {0} modules:",
   "loc.messages.AZ_InvalidARMEndpoint": "Specified AzureRM endpoint is invalid.",
-  "loc.messages.AZ_ManagedServiceIdentityError": "There was an error with the Managed Service Principal. Please configure Managed Service Identity (MSI) for virtual machine 'https://aka.ms/azure-msi-docs'"
+  "loc.messages.AZ_MsiAccessNotConfiguredProperlyFailure": "Could not fetch access token for Managed Service Principal. Please configure Managed Service Identity (MSI) for virtual machine 'https://aka.ms/azure-msi-docs'. Status code: '{0}', status message: {1}",
+  "loc.messages.AZ_MsiAccessTokenFetchFailure": "Could not fetch access token for Managed Service Principal. Status code: '{0}', status message: {1}",
+  "loc.messages.AZ_MsiFailure": "Could not fetch access token for Managed Service Principal. {0}"
 }
\ No newline at end of file
diff --git a/Tasks/Common/VstsAzureHelpers_/module.json b/Tasks/Common/VstsAzureHelpers_/module.json
index 5f288484dfc7..002b72a12346 100644
--- a/Tasks/Common/VstsAzureHelpers_/module.json
+++ b/Tasks/Common/VstsAzureHelpers_/module.json
@@ -10,6 +10,8 @@
         "AZ_UnsupportedAuthScheme0": "Unsupported authentication scheme '{0}' for Azure endpoint.",
         "AZ_AvailableModules": "The list of available {0} modules:",
         "AZ_InvalidARMEndpoint": "Specified AzureRM endpoint is invalid.",
-        "AZ_ManagedServiceIdentityError": "There was an error with the Managed Service Principal. Please configure Managed Service Identity (MSI) for virtual machine 'https://aka.ms/azure-msi-docs'"
+        "AZ_MsiAccessNotConfiguredProperlyFailure": "Could not fetch access token for Managed Service Principal. Please configure Managed Service Identity (MSI) for virtual machine 'https://aka.ms/azure-msi-docs'. Status code: '{0}', status message: {1}",
+        "AZ_MsiAccessTokenFetchFailure": "Could not fetch access token for Managed Service Principal. Status code: '{0}', status message: {1}" ,
+        "AZ_MsiFailure": "Could not fetch access token for Managed Service Principal. {0}"  
     }
 }

From 62756dfc2f46a4b0bc2c11f99a69b04175ca8339 Mon Sep 17 00:00:00 2001
From: Roshan Kumar <rok@microsoft.com>
Date: Fri, 13 Apr 2018 14:35:54 +0530
Subject: [PATCH 11/13] native powershell

---
 .../VstsAzureHelpers_/InitializeFunctions.ps1 | 40 +++++++++++++------
 1 file changed, 28 insertions(+), 12 deletions(-)

diff --git a/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1 b/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
index a1db49171a56..5ef348d368ed 100644
--- a/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
+++ b/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
@@ -184,11 +184,12 @@ function Initialize-AzureSubscription {
     } 
 }
 
-# Get the Bearer Access Token from the MSI Endpoint
+
+# Get the Bearer Access Token from the Endpoint
 function Get-MsiAccessToken {
     [CmdletBinding()]
     param([Parameter(Mandatory=$true)] $endpoint,
-        [Parameter(Mandatory=$true)] $count,
+        [Parameter(Mandatory=$true)] $retryCount,
         [Parameter(Mandatory=$true)] $timeToWait)
 
     $msiClientId = "";
@@ -200,7 +201,6 @@ function Get-MsiAccessToken {
     # Prepare contents for GET
     $method = "GET"
     $apiVersion = "2018-02-01";
-    let msiClientId =  $endpoint.Data.msiClientId ? "&client_id=" + this.msiClientId : 
     $authUri = "http://169.254.169.254/metadata/identity/oauth2/token?api-version=" + $apiVersion + "&resource=" + $endpoint.Url + $msiClientId;
     
     # Call Rest API to fetch AccessToken
@@ -209,7 +209,6 @@ function Get-MsiAccessToken {
     try
     {
         $retryLimit = 5;
-        $waitedTime = 2000 + $timeToWait * 2;
         $proxyUri = Get-ProxyUri $authUri
         if ($proxyUri -eq $null)
         {
@@ -222,27 +221,44 @@ function Get-MsiAccessToken {
             $response = Invoke-WebRequest -Uri $authUri -Method $method -Headers @{Metadata="true"} -UseDefaultCredentials -Proxy $proxyUri -ProxyUseDefaultCredentials -UseBasicParsing
         }
 
-        if($response.StatusCode == 209 -or $response.StatusCode == 500)
+        # Action on the based of response 
+        if(($response.StatusCode -eq 429) -or ($response.StatusCode -eq 500))
         {
-            if(count -lt $retryLimit)
+            if($retryCount -lt $retryLimit)
             {
-                count += 1
-                Get-MsiAccessToken $endpoint count  $waitedTime
+                $retryCount += 1
+                $waitedTime = 2000 + $timeToWait * 2
+                Start-Sleep -m $waitedTime
+                Get-MsiAccessToken $endpoint $retryCount  $waitedTime
             }
             else
             {
-                throw (Get-VstsLocString -Key AZ_MsiAccessTokenFetchFailure -ArgumentList $response.StatusCode $response.StatusDescription)
+                throw (Get-VstsLocString -Key AZ_MsiAccessTokenFetchFailure -ArgumentList $response.StatusCode, $response.StatusDescription)
             }
         }
+        elseif ($response.StatusCode -eq 200)
+        {
+            $accessToken = $response.Content | ConvertFrom-Json
+            return $accessToken.access_token
+        }
+        else
+        {
+            throw (Get-VstsLocString -Key AZ_MsiAccessNotConfiguredProperlyFailure -ArgumentList $response.StatusCode, $response.StatusDescription)
+        }
         
-        $content = $response.Content | ConvertFrom-Json
-        return $content.access_token
     }
     catch
     {
         $exceptionMessage = $_.Exception.Message.ToString()
         Write-Verbose "ExceptionMessage: $exceptionMessage (in function: Get-MsiAccessToken)"
-        throw (Get-VstsLocString -Key AZ_MsiAccessNotConfiguredProperlyFailure  -ArgumentList $response.StatusCode $response.StatusDescription)
+        if($exceptionMessage -match "400")
+        {
+            throw (Get-VstsLocString -Key AZ_MsiAccessNotConfiguredProperlyFailure -ArgumentList $response.StatusCode, $response.StatusDescription)
+        }
+        else
+        {
+            throw $_.Exception
+        }
     }
 }
 

From d059c7d421f0a17e57bce68d97e7fca353a125bf Mon Sep 17 00:00:00 2001
From: Roshan Kumar <rok@microsoft.com>
Date: Fri, 13 Apr 2018 15:13:00 +0530
Subject: [PATCH 12/13] MSI

---
 Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1 | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1 b/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
index 5ef348d368ed..f3869c94d630 100644
--- a/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
+++ b/Tasks/Common/VstsAzureHelpers_/InitializeFunctions.ps1
@@ -180,7 +180,7 @@ function Initialize-AzureSubscription {
         
         Set-CurrentAzureRMSubscription -SubscriptionId $Endpoint.Data.SubscriptionId -TenantId $Endpoint.Auth.Parameters.TenantId
     }else {
-        throw (Get-VstsLocString-Key AZ_UnsupportedAuthScheme0 -ArgumentList $Endpoint.Auth.Scheme)
+        throw (Get-VstsLocString -Key AZ_UnsupportedAuthScheme0 -ArgumentList $Endpoint.Auth.Scheme)
     } 
 }
 

From 7cca967fba41aa4f8abafa5124d3ed415c9781af Mon Sep 17 00:00:00 2001
From: Roshan Kumar <rok@microsoft.com>
Date: Fri, 13 Apr 2018 17:35:52 +0530
Subject: [PATCH 13/13] response modifieed

---
 ...nitialize-AzureSubscription.ManagedServiceIdentity.ps1 | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1 b/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1
index 96b9aed57630..f68bf4c033dd 100644
--- a/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1
+++ b/Tasks/Common/VstsAzureHelpers_/Tests/Initialize-AzureSubscription.ManagedServiceIdentity.ps1
@@ -27,6 +27,12 @@ $content = @"
            {"access_token" : "Dummy Token" }
 "@
 
+$response = @{
+    Content = $content
+    StatusCode = 200
+    StatusDescription = 'OK'
+};
+
 $variableSets = @(
     @{ StorageAccount = 'Some storage account' }
 )
@@ -39,7 +45,7 @@ foreach ($variableSet in $variableSets) {
     Register-Mock Add-AzureRMAccount { 'some output' }
     Register-Mock Set-CurrentAzureRMSubscription
     Register-Mock Set-UserAgent
-    Register-Mock Invoke-WebRequest { @{Content = $content} }
+    Register-Mock Invoke-WebRequest { $response }
     
     # Act.
     $result = & $module Initialize-AzureSubscription -Endpoint $endpoint -StorageAccount $variableSet.StorageAccount