add param to Add-AzVhd, and New-AzHostGroup

src/Compute/Compute.AlcWrapper/PSPageBlobClient.cs

@@ -17,6 +17,7 @@
 using System.Text;
 using Azure.Storage.Blobs.Specialized;
 using System.IO;
+using Azure.Core;
 
 namespace Microsoft.Azure.Commands.Compute
 {
@@ -28,9 +29,9 @@ internal PSPageBlobClient(PageBlobClient pageblobclient)
         {
             _pageBlobClient = pageblobclient;
         }
-        public PSPageBlobClient(Uri blobUri)
+        public PSPageBlobClient(Uri blobUri, TokenCredential tokenCredential = null )
         {
-            _pageBlobClient = new PageBlobClient(blobUri, null);
+            _pageBlobClient = tokenCredential == null ? new PageBlobClient(blobUri, null) : new PageBlobClient(blobUri, tokenCredential);
         }
 
         public Uri Uri { get { return _pageBlobClient.Uri; } }

src/Compute/Compute/ChangeLog.md

@@ -20,6 +20,7 @@
 
 -->
 ## Upcoming Release
+* Add `-DataAccessAuthMode` parameter to Add-AzVhd DirectUploadToManagedDisk parameter set. 
 
 ## Version 4.27.0
 * Edited `New-AzVm` cmdlet internal logic to use the `PlatformFaultDomain` value in the `PSVirtualMachine` object passed to it in the new virtual machine.

src/Compute/Compute/Common/ComputeTokenCredential.cs

@@ -0,0 +1,73 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using Azure.Core;
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Azure.Commands.Common.Authentication;
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
+
+namespace Microsoft.Azure.Commands.Compute.Common
+{
+    public class ComputeTokenCredential : TokenCredential
+    {
+        public IAccessToken accessToken { get; set; }
+        public ComputeTokenCredential(IAzureContext DefaultContext, string customAudience)
+        {
+
+            if (DefaultContext == null || DefaultContext.Account == null)
+            {
+                throw new InvalidOperationException();
+            }
+
+            accessToken = AzureSession.Instance.AuthenticationFactory.Authenticate(
+                   DefaultContext.Account,
+                   EnsureCustomAudienceSet(DefaultContext.Environment, customAudience),
+                   DefaultContext.Tenant.Id,
+                   null,
+                   ShowDialog.Never,
+                   null,
+                   customAudience);
+
+        }
+
+        public override AccessToken GetToken(TokenRequestContext requestContext, CancellationToken cancellationToken)
+        {
+            AccessToken token;
+            accessToken.AuthorizeRequest((tokenType, tokenValue) =>
+            {
+                token = new AccessToken(tokenValue, DateTimeOffset.UtcNow);
+            });
+            return token;
+        }
+
+        public override ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken)
+        {
+            return new ValueTask<AccessToken>(this.GetToken(requestContext, cancellationToken));
+        }
+
+        private IAzureEnvironment EnsureCustomAudienceSet(IAzureEnvironment environment, string customAudience)
+        {
+            if (environment != null)
+            {
+                if (!environment.IsPropertySet(customAudience))
+                {
+                    environment.SetProperty(customAudience, customAudience);
+                }
+            }
+            return environment;
+        }
+    }
+}

src/Compute/Compute/Compute.csproj

@@ -1,4 +1,4 @@
-<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
 
   <PropertyGroup>
     <PsModuleName>Compute</PsModuleName>

src/Compute/Compute/StorageServices/AddAzureVhdCommand.cs

@@ -36,6 +36,8 @@
 using Microsoft.Samples.HyperV.Storage;
 using Microsoft.Samples.HyperV.Common;
 using System.Threading;
+using Azure.Core;
+using Microsoft.Azure.Commands.Compute.Common;
 
 
 namespace Microsoft.Azure.Commands.Compute.StorageServices
@@ -178,6 +180,14 @@ public class AddAzureVhdCommand : ComputeClientBaseCmdlet
             HelpMessage = "Skips the resizing of VHD")]
         public SwitchParameter SkipResizing { get; set; }
 
+        [Parameter(
+            Mandatory = false,
+            ValueFromPipelineByPropertyName = true,
+            ParameterSetName = DirectUploadToManagedDiskSet,
+            HelpMessage = "Additional authentication requirements when exporting or uploading to a disk or snapshot. Possible options are: \"AzureActiveDirectory\" and \"None\".")]
+        [PSArgumentCompleter("AzureActiveDirectory", "None")]
+        public string DataAccessAuthMode { get; set; }
+
         [Parameter(Mandatory = false, HelpMessage = "Run cmdlet in the background")]
         public SwitchParameter AsJob { get; set; }
 
@@ -228,13 +238,19 @@ public override void ExecuteCmdlet()
 
                         // 3-4: UPLOAD                  
                         WriteVerbose("Preparing for Upload");
-                        PSPageBlobClient managedDisk = new PSPageBlobClient(sasUri);
+                        ComputeTokenCredential tokenCredential = null;
+                        if (this.DataAccessAuthMode == "AzureActiveDirectory")
+                        {
+                            // get token 
+                            tokenCredential = new ComputeTokenCredential(DefaultContext, "https://disk.compute.azure.com/");
+                        }
+                        PSPageBlobClient managedDisk = new PSPageBlobClient(sasUri, tokenCredential);
                         DiskUploadCreator diskUploadCreator = new DiskUploadCreator();
                         var uploadContext = diskUploadCreator.Create(this.LocalFilePath, managedDisk, false);
                         var synchronizer = new DiskSynchronizer(uploadContext, this.NumberOfUploaderThreads ?? DefaultNumberOfUploaderThreads);
 
                         WriteVerbose("Uploading");
-                        if (synchronizer.Synchronize())
+                        if (synchronizer.Synchronize(tokenCredential))
                         {
                             var result = new VhdUploadContext { LocalFilePath = this.LocalFilePath, DestinationUri = sasUri };
                             WriteObject(result);
@@ -378,7 +394,8 @@ private PSDisk CreateDiskConfig()
                 EncryptionSettingsCollection = null,
                 Encryption = null,
                 NetworkAccessPolicy = null,
-                DiskAccessId = null
+                DiskAccessId = null,
+                DataAccessAuthMode = this.IsParameterBound(c => c.DataAccessAuthMode) ? this.DataAccessAuthMode : null
             };
             return vDisk;
         }

src/Compute/Compute/Sync/Upload/DiskSynchronizer.cs

@@ -23,6 +23,7 @@
 using Microsoft.WindowsAzure.Commands.Sync.Upload;
 using Microsoft.WindowsAzure.Commands.Sync;
 using System.Threading.Tasks;
+using Microsoft.Azure.Commands.Compute.Common;
 
 namespace Microsoft.Azure.Commands.Compute.Sync.Upload
 {
@@ -47,7 +48,7 @@ public DiskSynchronizer(UploadContextDisk context, int maxParallelism)
             this.maxParallelism = maxParallelism;
         }
 
-        public bool Synchronize()
+        public bool Synchronize(ComputeTokenCredential tokenCredential)
         {
             var uploadStatus = new ProgressStatus(alreadyUploadedData, alreadyUploadedData + dataToUpload, new ComputeStats());
 
@@ -57,7 +58,7 @@ public bool Synchronize()
                 Task<LoopResult> task = Task<LoopResult>.Factory.StartNew(() =>
                 {
                     return Threading.Parallel.ForEach(dataWithRanges,
-                                                  () => new PSPageBlobClient(pageBlob.Uri),
+                                                  () => new PSPageBlobClient(pageBlob.Uri, tokenCredential),
                                                   (dwr, b) =>
                                                   {
                                                       using (dwr)

src/Compute/Compute/help/Add-AzVhd.md

@@ -24,8 +24,8 @@ Add-AzVhd [-ResourceGroupName] <String> [-Destination] <Uri> [-LocalFilePath] <F
 ```
 Add-AzVhd [-ResourceGroupName] <String> [-LocalFilePath] <FileInfo> -DiskName <String> [-Location] <String>
  [-DiskSku <String>] [-DiskZone <String[]>] [-DiskHyperVGeneration <String>]
- [-DiskOsType <OperatingSystemTypes>] [[-NumberOfUploaderThreads] <Int32>] [-AsJob]
- [-DefaultProfile <IAzureContextContainer>] [<CommonParameters>]
+ [-DiskOsType <OperatingSystemTypes>] [[-NumberOfUploaderThreads] <Int32>] [-DataAccessAuthMode <String>]
+ [-AsJob] [-DefaultProfile <IAzureContextContainer>] [<CommonParameters>]
 ```
 
 ## DESCRIPTION
@@ -127,6 +127,21 @@ Accept pipeline input: True (ByPropertyName)
 Accept wildcard characters: False
 ```
 
+### -DataAccessAuthMode
+Additional authentication requirements when exporting or uploading to a disk or snapshot. Possible options are: "AzureActiveDirectory" and "None".
+
+```yaml
+Type: System.String
+Parameter Sets: DirectUploadToManagedDiskSet
+Aliases:
+
+Required: False
+Position: Named
+Default value: None
+Accept pipeline input: True (ByPropertyName)
+Accept wildcard characters: False
+```
+
 ### -DefaultProfile
 The credentials, account, tenant, and subscription used for communication with azure.
 
@@ -299,19 +314,7 @@ Specifies the name of the resource group of the virtual machine.
 
 ```yaml
 Type: System.String
-Parameter Sets: DefaultParameterSet
-Aliases:
-
-Required: False
-Position: 0
-Default value: None
-Accept pipeline input: True (ByPropertyName)
-Accept wildcard characters: False
-```
-
-```yaml
-Type: System.String
-Parameter Sets: DirectUploadToManagedDiskSet
+Parameter Sets: (All)
 Aliases:
 
 Required: True