Add auth to V3 monitoring search cursor URLs (#10115)

* Add auth to V3 monitoring search cursor URLs
* Use -clientId for managed identity auth on search cursors
* Allow UseManagedIdentity to be set per search instance

Author: joelverhagen

src/Catalog/AzureBlobCursor.cs

@@ -0,0 +1,45 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Diagnostics;
+using System.Threading;
+using System.Threading.Tasks;
+using Azure;
+using Azure.Storage.Blobs;
+using Azure.Storage.Blobs.Models;
+using Newtonsoft.Json.Linq;
+
+namespace NuGet.Services.Metadata.Catalog
+{
+    public class AzureBlobCursor : ReadCursor
+    {
+        private readonly BlobClient _blobClient;
+
+        public AzureBlobCursor(BlobClient blobClient)
+        {
+            _blobClient = blobClient ?? throw new ArgumentNullException(nameof(blobClient));
+        }
+
+        public override async Task LoadAsync(CancellationToken cancellationToken)
+        {
+            BlobDownloadResult downloadResult;
+            try
+            {
+                downloadResult = await _blobClient.DownloadContentAsync(cancellationToken);
+            }
+            catch (RequestFailedException ex)
+            {
+                Trace.TraceError("AzureBlobCursor.LoadAsync: error {0} {1}", ex.Status, _blobClient.Uri.AbsoluteUri);
+                throw;
+            }
+
+            var json = downloadResult.Content.ToString();
+
+            JObject obj = JObject.Parse(json);
+            Value = obj["value"].ToObject<DateTime>();
+
+            Trace.TraceInformation("AzureBlobCursor.LoadAsync: {0:O} {1}", Value, _blobClient.Uri.AbsoluteUri);
+        }
+    }
+}

src/Ng/Arguments.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using NuGet.Services.Metadata.Catalog.Monitoring;
@@ -99,6 +99,19 @@ public static class Arguments
         /// </summary>
         public const string SearchCursorUriPrefix = "searchCursorUri-";
 
+        /// <summary>
+        /// The argument prefix for the cursor SAS token of a <see cref="SearchEndpoint"/> cursor.
+        /// This is used in conjunction with the <see cref="SearchBaseUriPrefix"/> argument with same suffix for authentication with blob storage.
+        /// </summary>
+        public const string SearchCursorSasValuePrefix = "searchCursorSasValue-";
+
+        /// <summary>
+        /// The argument prefix to enable using a token credential for a <see cref="SearchEndpoint"/> cursor.
+        /// This is used in conjunction with the <see cref="SearchBaseUriPrefix"/> argument with same suffix for authentication with blob storage.
+        /// If <see cref="ClientId"/> is specified, a managed identity credential will be used. Otherwise, a default Azure credential will be used.
+        /// </summary>
+        public const string SearchCursorUseManagedIdentityPrefix = "searchCursorUseManagedIdentity-";
+
         /// <summary>
         /// The argument prefix for the base URL of a <see cref="SearchEndpoint"/>. There should be the same number of
         /// <see cref="SearchBaseUriPrefix"/> parameters passed as <see cref="SearchCursorUriPrefix"/> with the same
@@ -186,4 +199,4 @@ public static class Arguments
         public const string CursorFile = "cursorFile";
         #endregion
     }
-}
+}

src/Ng/CommandHelpers.cs

@@ -8,6 +8,9 @@
 using System.Net;
 using System.Net.Http;
 using System.Security.Cryptography.X509Certificates;
+using Azure;
+using Azure.Core;
+using Azure.Identity;
 using Azure.Storage.Blobs;
 using Azure.Storage.Queues;
 using Microsoft.Extensions.Logging;
@@ -28,6 +31,7 @@ namespace Ng
 {
     public static class CommandHelpers
     {
+        private const string DefaultStorageSuffix = "core.windows.net";
         private static readonly int DefaultKeyVaultSecretCachingTimeout = 60 * 60 * 6; // 6 hours;
         private static readonly HashSet<string> NotInjectedKeys = new HashSet<string>(StringComparer.OrdinalIgnoreCase)
         {
@@ -211,16 +215,14 @@ private static CatalogStorageFactory CreateStorageFactoryImpl(
 
             if (Arguments.AzureStorageType.Equals(storageType, StringComparison.InvariantCultureIgnoreCase))
             {
-                var storageAccountName = arguments.GetOrThrow<string>(argumentNameMap[Arguments.StorageAccountName]);
                 var storageContainer = arguments.GetOrThrow<string>(argumentNameMap[Arguments.StorageContainer]);
                 var storagePath = arguments.GetOrDefault<string>(argumentNameMap[Arguments.StoragePath]);
-                var storageSuffix = arguments.GetOrDefault(argumentNameMap[Arguments.StorageSuffix], "core.windows.net");
                 var storageOperationMaxExecutionTime = MaxExecutionTime(arguments.GetOrDefault<int>(argumentNameMap[Arguments.StorageOperationMaxExecutionTimeInSeconds]));
                 var storageServerTimeout = MaxExecutionTime(arguments.GetOrDefault<int>(argumentNameMap[Arguments.StorageServerTimeoutInSeconds]));
                 var storageUseServerSideCopy = arguments.GetOrDefault<bool>(argumentNameMap[Arguments.StorageUseServerSideCopy]);
                 var storageInitializeContainer = arguments.GetOrDefault(argumentNameMap[Arguments.StorageInitializeContainer], defaultValue: true);
 
-                BlobServiceClient account = GetBlobServiceClient(storageAccountName, storageSuffix, arguments, argumentNameMap);
+                BlobServiceClient account = GetBlobServiceClient(arguments, argumentNameMap);
 
                 return new CatalogAzureStorageFactory(
                     account,
@@ -283,19 +285,28 @@ public static Func<HttpMessageHandler> GetHttpMessageHandlerFactory(
 
         public static EndpointConfiguration GetEndpointConfiguration(IDictionary<string, string> arguments)
         {
+            var clientId = arguments.GetOrDefault<string>(Arguments.ClientId);
+
             var registrationCursorUri = arguments.GetOrThrow<Uri>(Arguments.RegistrationCursorUri);
             var flatContainerCursorUri = arguments.GetOrThrow<Uri>(Arguments.FlatContainerCursorUri);
 
-            var instanceNameToSearchBaseUri = GetSuffixToUri(arguments, Arguments.SearchBaseUriPrefix);
-            var instanceNameToSearchCursorUri = GetSuffixToUri(arguments, Arguments.SearchCursorUriPrefix);
+            var instanceNameToSearchBaseUri = GetSuffixToValue<Uri>(arguments, Arguments.SearchBaseUriPrefix);
+            var instanceNameToSearchCursorUri = GetSuffixToValue<Uri>(arguments, Arguments.SearchCursorUriPrefix);
+            var instanceNameToSearchCursorSasValue = GetSuffixToValue<string>(arguments, Arguments.SearchCursorSasValuePrefix);
+            var instanceNameToSearchCursorUseManagedIdentity = GetSuffixToValue<bool>(arguments, Arguments.SearchCursorUseManagedIdentityPrefix);
             var instanceNameToSearchConfig = new Dictionary<string, SearchEndpointConfiguration>();
+
             foreach (var pair in instanceNameToSearchBaseUri)
             {
                 var instanceName = pair.Key;
 
                 // Find all cursors with an instance name starting with the search base URI instance name. We do this
                 // because there may be multiple potential cursors representing the state of a search service.
-                var matchingCursors = instanceNameToSearchCursorUri.Keys.Where(x => x.StartsWith(instanceName)).ToList();
+                var matchingCursors = instanceNameToSearchCursorUri
+                    .Keys
+                    .Where(x => x.StartsWith(instanceName))
+                    .OrderBy(x => x)
+                    .ToList();
 
                 if (!matchingCursors.Any())
                 {
@@ -304,9 +315,51 @@ public static EndpointConfiguration GetEndpointConfiguration(IDictionary<string,
                         $"-{Arguments.SearchCursorUriPrefix}{instanceName}* arguments.");
                 }
 
-                instanceNameToSearchConfig[instanceName] = new SearchEndpointConfiguration(
-                    matchingCursors.Select(x => instanceNameToSearchCursorUri[x]).ToList(),
-                    pair.Value);
+                var cursors = new List<SearchCursorConfiguration>();
+
+                foreach (var suffix in matchingCursors)
+                {
+                    var cursorUri = instanceNameToSearchCursorUri[suffix];
+                    SearchCursorCredentialType credentialType;
+
+                    BlobClient blobClient = null;
+                    if (instanceNameToSearchCursorUseManagedIdentity.TryGetValue(suffix, out var useManagedIdentity)
+                        && useManagedIdentity)
+                    {
+                        TokenCredential credential;
+                        if (string.IsNullOrEmpty(clientId))
+                        {
+                            credential = new DefaultAzureCredential();
+                            credentialType = SearchCursorCredentialType.DefaultAzureCredential;
+                        }
+                        else
+                        {
+                            credential = new ManagedIdentityCredential(clientId);
+                            credentialType = SearchCursorCredentialType.ManagedIdentityCredential;
+                        }
+
+                        blobClient = new BlobClient(cursorUri, credential);
+                    }
+                    else if (instanceNameToSearchCursorSasValue.TryGetValue(suffix, out var sas))
+                    {
+                        if (sas.StartsWith("?"))
+                        {
+                            // workaround for https://github.com/Azure/azure-sdk-for-net/issues/44373
+                            sas = sas.Substring(1);
+                        }
+
+                        blobClient = new BlobClient(cursorUri, new AzureSasCredential(sas));
+                        credentialType = SearchCursorCredentialType.AzureSasCredential;
+                    }
+                    else
+                    {
+                        credentialType = SearchCursorCredentialType.Anonymous;
+                    }
+
+                    cursors.Add(new SearchCursorConfiguration(cursorUri, blobClient, credentialType));
+                }
+
+                instanceNameToSearchConfig[instanceName] = new SearchEndpointConfiguration(cursors, pair.Value);
 
                 foreach (var key in matchingCursors)
                 {
@@ -329,13 +382,13 @@ public static EndpointConfiguration GetEndpointConfiguration(IDictionary<string,
                 instanceNameToSearchConfig);
         }
 
-        private static Dictionary<string, Uri> GetSuffixToUri(IDictionary<string, string> arguments, string prefix)
+        private static Dictionary<string, T> GetSuffixToValue<T>(IDictionary<string, string> arguments, string prefix)
         {
-            var suffixToUri = new Dictionary<string, Uri>();
+            var suffixToUri = new Dictionary<string, T>();
             foreach (var key in arguments.Keys.Where(x => x.StartsWith(prefix)))
             {
                 var suffix = key.Substring(prefix.Length);
-                suffixToUri[suffix] = arguments.GetOrThrow<Uri>(key);
+                suffixToUri[suffix] = arguments.GetOrThrow<T>(key);
             }
 
             return suffixToUri;
@@ -354,10 +407,8 @@ public static IStorageQueue<T> CreateStorageQueue<T>(IDictionary<string, string>
 
             if (Arguments.AzureStorageType.Equals(storageType, StringComparison.InvariantCultureIgnoreCase))
             {
-                var storageAccountName = arguments.GetOrThrow<string>(Arguments.StorageAccountName);
                 var storageQueueName = arguments.GetOrDefault<string>(Arguments.StorageQueueName);
-
-                QueueServiceClient account = GetQueueServiceClient(storageAccountName, endpointSuffix: null, arguments, ArgumentNames);
+                QueueServiceClient account = GetQueueServiceClient(arguments, ArgumentNames);
                 return new StorageQueue<T>(new AzureStorageQueue(account, storageQueueName),
                     new JsonMessageSerializer<T>(JsonSerializerUtility.SerializerSettings), version);
             }
@@ -367,8 +418,34 @@ public static IStorageQueue<T> CreateStorageQueue<T>(IDictionary<string, string>
             }
         }
 
-        private static BlobServiceClient GetBlobServiceClient(string storageAccountName, string endpointSuffix, IDictionary<string, string> arguments, IDictionary<string, string> argumentNameMap)
+        private static BlobServiceClient GetBlobServiceClient(
+            IDictionary<string, string> arguments,
+            IDictionary<string, string> argumentNameMap)
+        {
+            string connectionString = GetConnectionString(arguments, argumentNameMap, "BlobEndpoint", "blob");
+            return new BlobServiceClient(connectionString);
+        }
+
+        private static QueueServiceClient GetQueueServiceClient(
+            IDictionary<string, string> arguments,
+            IDictionary<string, string> argumentNameMap)
+        {
+            string connectionString = GetConnectionString(arguments, argumentNameMap, "QueueEndpoint", "queue");
+            return new QueueServiceClient(connectionString, new QueueClientOptions
+            {
+                // We use base64 encoding for compatibility with the older SDK
+                MessageEncoding = QueueMessageEncoding.Base64,
+            });
+        }
+
+        private static string GetConnectionString(
+            IDictionary<string, string> arguments,
+            IDictionary<string, string> argumentNameMap,
+            string endpointKey,
+            string endpointDomain)
         {
+            var storageAccountName = arguments.GetOrThrow<string>(argumentNameMap[Arguments.StorageAccountName]);
+            var storageSuffix = arguments.GetOrDefault(argumentNameMap[Arguments.StorageSuffix], DefaultStorageSuffix);
             var storageKeyValue = arguments.GetOrDefault<string>(argumentNameMap[Arguments.StorageKeyValue]);
 
             string connectionString;
@@ -382,34 +459,14 @@ private static BlobServiceClient GetBlobServiceClient(string storageAccountName,
                     storageSasValue = storageSasValue.Substring(1);
                 }
 
-                connectionString = $"BlobEndpoint=https://{storageAccountName}.blob.{endpointSuffix}/;SharedAccessSignature={storageSasValue}";
-            }
-            else
-            {
-                connectionString = $"DefaultEndpointsProtocol=https;AccountName={storageAccountName};AccountKey={storageKeyValue};EndpointSuffix={endpointSuffix}";
-            }
-
-            return new BlobServiceClient(connectionString);
-        }
-
-        private static QueueServiceClient GetQueueServiceClient(string storageAccountName, string endpointSuffix, IDictionary<string, string> arguments, IDictionary<string, string> argumentNameMap)
-        {
-            var storageKeyValue = arguments.GetOrDefault<string>(argumentNameMap[Arguments.StorageKeyValue]);
-
-            string connectionString;
-
-            if (string.IsNullOrEmpty(storageKeyValue))
-            {
-                var storageSasValue = arguments.GetOrThrow<string>(argumentNameMap[Arguments.StorageSasValue]);
-                connectionString = $"BlobEndpoint=https://{storageAccountName}.blob.{endpointSuffix}/;SharedAccessSignature={storageSasValue}";
+                connectionString = $"{endpointKey}=https://{storageAccountName}.{endpointDomain}.{storageSuffix}/;SharedAccessSignature={storageSasValue}";
             }
             else
             {
-                connectionString = $"DefaultEndpointsProtocol=https;AccountName={storageAccountName};AccountKey={storageKeyValue};EndpointSuffix={endpointSuffix}";
+                connectionString = $"DefaultEndpointsProtocol=https;AccountName={storageAccountName};AccountKey={storageKeyValue};EndpointSuffix={storageSuffix}";
             }
 
-            return new QueueServiceClient(connectionString);
+            return connectionString;
         }
-
     }
 }

src/NuGet.Services.Configuration/DictionaryExtensions.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
@@ -35,7 +35,10 @@ public static class DictionaryExtensions
         /// <exception cref="ArgumentException">Thrown when the value associated with the key in the dictionary is null or empty.</exception>
         public static T GetOrThrow<T>(this IDictionary<string, string> dictionary, string key)
         {
-            var value = dictionary[key];
+            if (!dictionary.TryGetValue(key, out var value))
+            {
+                throw new KeyNotFoundException($"Key {key} was not found in the dictionary.");
+            }
 
             if (string.IsNullOrEmpty(value))
             {

src/NuGet.Services.Metadata.Catalog.Monitoring/Utility/ContainerBuilderExtensions.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
@@ -77,7 +77,7 @@ private static void RegisterSearchEndpoints(this ContainerBuilder builder, Endpo
                 builder
                     .Register(c => new SearchEndpoint(
                         pair.Key,
-                        pair.Value.CursorUris,
+                        pair.Value.Cursors,
                         pair.Value.BaseUri,
                         c.Resolve<Func<HttpMessageHandler>>()))
                     .As<IEndpoint>()
@@ -240,4 +240,4 @@ private static void RegisterResourceProvider<TProvider>(this ContainerBuilder bu
                 .Keyed<INuGetResourceProvider>(type);
         }
     }
-}
+}

src/NuGet.Services.Metadata.Catalog.Monitoring/Validation/Test/Endpoint/SearchCursorConfiguration.cs

@@ -0,0 +1,30 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using Azure.Storage.Blobs;
+
+namespace NuGet.Services.Metadata.Catalog.Monitoring
+{
+    public enum SearchCursorCredentialType
+    {
+        Anonymous = 1,
+        AzureSasCredential,
+        DefaultAzureCredential,
+        ManagedIdentityCredential,
+    }
+
+    public sealed class SearchCursorConfiguration
+    {
+        public SearchCursorConfiguration(Uri cursorUri, BlobClient blobClient, SearchCursorCredentialType credentialType)
+        {
+            CursorUri = cursorUri;
+            BlobClient = blobClient;
+            CredentialType = credentialType;
+        }
+
+        public Uri CursorUri { get; }
+        public BlobClient BlobClient { get; }
+        public SearchCursorCredentialType CredentialType { get; }
+    }
+}