Add CustomerResourceId to enrich push metrics with tenant ID (NuGet#7936

)

Progress on NuGet/Engineering#3106

src/NuGetGallery.Core/Authentication/MicrosoftClaims.cs

@@ -0,0 +1,15 @@
+// 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.
+
+namespace NuGetGallery.Authentication
+{
+    public static class MicrosoftClaims
+    {
+        private const string ClaimsDomain = "http://schemas.microsoft.com/identity/claims/";
+
+        /// <summary>
+        /// The claim URL for the claim that stores the user's tenant ID, based on the external credential.
+        /// </summary>
+        public const string TenantId = ClaimsDomain + "tenantid";
+    }
+}

src/NuGetGallery.Core/NuGetGallery.Core.csproj

@@ -105,6 +105,7 @@
     <Compile Include="Auditing\UserAuditRecord.cs" />
     <Compile Include="Authentication\AuthenticationExtensions.cs" />
     <Compile Include="Authentication\CredentialTypeInfo.cs" />
+    <Compile Include="Authentication\MicrosoftClaims.cs" />
     <Compile Include="Certificates\CertificateFile.cs" />
     <Compile Include="Cookies\CookieComplianceServiceBase.cs" />
     <Compile Include="Cookies\CookieConsentMessage.cs" />

src/NuGetGallery.Services/Extensions/HttpContextBaseExtensions.cs

@@ -11,6 +11,8 @@ namespace NuGetGallery
 {
     public static class HttpContextBaseExtensions
     {
+        public static HttpContextBase GetCurrent() => new HttpContextWrapper(HttpContext.Current);
+
         public static User GetCurrentUser(this HttpContextBase httpContext)
         {
             return httpContext.GetOwinContext().GetCurrentUser();

src/NuGetGallery.Services/Extensions/IOwinContextExtensions.cs

@@ -2,11 +2,12 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Linq;
 using System.Security.Claims;
-using System.Web;
 using System.Web.Mvc;
 using Microsoft.Owin;
 using NuGet.Services.Entities;
+using NuGetGallery.Authentication;
 
 namespace NuGetGallery
 {
@@ -64,9 +65,27 @@ private static User LoadUser(IOwinContext context)
 
                 if (!String.IsNullOrEmpty(userName))
                 {
-                    return DependencyResolver.Current
+                    var user = DependencyResolver.Current
                         .GetService<IUserService>()
                         .FindByUsername(userName);
+
+                    // Try to add the tenant ID information as an additional claim since we have the full user record
+                    // and the associated credentials.
+                    if (user != null && principal.Identity is ClaimsIdentity identity)
+                    {
+                        // From the schema, it is possible to have multiple credentials. Prefer the latest one.
+                        var externalCredential = user
+                            .Credentials
+                            .OrderByDescending(x => x.Created)
+                            .FirstOrDefault(c => c.IsExternal() && c.TenantId != null);
+
+                        if (externalCredential != null)
+                        {
+                            identity.TryAddClaim(MicrosoftClaims.TenantId, externalCredential.TenantId);
+                        }
+                    }
+
+                    return user;
                 }
             }
             return null; // No user logged in, or credentials could not be resolved

src/NuGetGallery.Services/Extensions/PrincipalExtensions.cs

@@ -213,6 +213,17 @@ public static bool TryRemoveClaim(this IIdentity identity, string claimType)
 
             return false;
         }
+
+        /// <summary>
+        /// Get the tenant ID from the claims, if available. If no such claim exists, null is returned.
+        /// </summary>
+        /// <param name="identity">The identity to look for claims in.</param>
+        /// <returns>The tenant ID or null.</returns>
+        public static string GetTenantIdOrNull(this IIdentity identity)
+        {
+            var claimsIdentity = identity as ClaimsIdentity;
+            return claimsIdentity?.FindFirst(MicrosoftClaims.TenantId)?.Value;
+        }
 
         /// <summary>
         /// Try to add a new default claim to the identity. It will not replace an existing claim.

src/NuGetGallery.Services/Telemetry/TelemetryService.cs

@@ -18,7 +18,7 @@ namespace NuGetGallery
 {
     public class TelemetryService : ITelemetryService, IFeatureFlagTelemetryService
     {
-        internal class Events
+        public class Events
         {
             public const string ODataQueryFilter = "ODataQueryFilter";
             public const string ODataCustomQuery = "ODataCustomQuery";

src/NuGetGallery/App_Start/DefaultDependenciesModule.cs

@@ -525,6 +525,7 @@ internal static ApplicationInsightsConfiguration ConfigureApplicationInsights(
             telemetryConfiguration.TelemetryInitializers.Add(new ClientInformationTelemetryEnricher());
             telemetryConfiguration.TelemetryInitializers.Add(new KnownOperationNameEnricher());
             telemetryConfiguration.TelemetryInitializers.Add(new AzureWebAppTelemetryInitializer());
+            telemetryConfiguration.TelemetryInitializers.Add(new CustomerResourceIdEnricher());
 
             // Add processors
             telemetryConfiguration.TelemetryProcessorChainBuilder.Use(next =>

src/NuGetGallery/NuGetGallery.csproj

@@ -587,6 +587,7 @@
     <Compile Include="Services\ValidationService.cs" />
     <Compile Include="Telemetry\ClientInformationTelemetryEnricher.cs" />
     <Compile Include="Telemetry\ClientTelemetryPIIProcessor.cs" />
+    <Compile Include="Telemetry\CustomerResourceIdEnricher.cs" />
     <Compile Include="Telemetry\KnownOperationNameEnricher.cs" />
     <Compile Include="Telemetry\QuietExceptionLogger.cs" />
     <Compile Include="Infrastructure\PasswordValidationAttribute.cs" />

src/NuGetGallery/Telemetry/ClientInformationTelemetryEnricher.cs

@@ -46,9 +46,6 @@ public void Initialize(ITelemetry telemetry)
             }
         }
 
-        protected virtual HttpContextBase GetHttpContext()
-        {
-            return new HttpContextWrapper(HttpContext.Current);
-        }
+        protected virtual HttpContextBase GetHttpContext() => HttpContextBaseExtensions.GetCurrent();
     }
 }

src/NuGetGallery/Telemetry/CustomerResourceIdEnricher.cs

@@ -0,0 +1,52 @@
+// 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.Collections.Generic;
+using System.Web;
+using Microsoft.ApplicationInsights.Channel;
+using Microsoft.ApplicationInsights.DataContracts;
+using Microsoft.ApplicationInsights.Extensibility;
+
+namespace NuGetGallery
+{
+    public class CustomerResourceIdEnricher : ITelemetryInitializer
+    {
+        private const string CustomerResourceId = "CustomerResourceId";
+        private const string Prefix = "/tenants/";
+        private static readonly string Empty = Prefix + Guid.Empty.ToString();
+
+        private static readonly HashSet<string> CustomMetricNames = new HashSet<string>
+        {
+            TelemetryService.Events.PackagePush,
+            TelemetryService.Events.PackagePushFailure,
+        };
+
+        public void Initialize(ITelemetry telemetry)
+        {
+            var metric = telemetry as MetricTelemetry;
+            if (metric == null)
+            {
+                return;
+            }
+
+            if (!CustomMetricNames.Contains(metric.Name))
+            {
+                return;
+            }
+
+            var itemTelemetry = telemetry as ISupportProperties;
+            if (itemTelemetry == null)
+            {
+                return;
+            }
+
+            var httpContext = GetHttpContext();
+            var tenantId = httpContext?.User?.Identity.GetTenantIdOrNull();
+            var customerResourceId = tenantId != null ? Prefix + tenantId : Empty;
+            itemTelemetry.Properties[CustomerResourceId] = customerResourceId;
+        }
+
+        protected virtual HttpContextBase GetHttpContext() => HttpContextBaseExtensions.GetCurrent();
+    }
+}

tests/NuGetGallery.Facts/App_Start/DefaultDependenciesModuleFacts.cs

@@ -154,6 +154,7 @@ private Action<ITelemetryInitializer>[] GetTelemetryInitializerInspectors(string
                     ti => ti.GetType().Equals(typeof(ClientInformationTelemetryEnricher)),
                     ti => ti.GetType().Equals(typeof(KnownOperationNameEnricher)),
                     ti => ti.GetType().Equals(typeof(AzureWebAppTelemetryInitializer)),
+                    ti => ti.GetType().Equals(typeof(CustomerResourceIdEnricher)),
 
                     // Registered by applicationinsights.config
                     ti => ti.GetType().Equals(typeof(HttpDependenciesParsingTelemetryInitializer)),

tests/NuGetGallery.Facts/NuGetGallery.Facts.csproj

@@ -110,6 +110,7 @@
     <Compile Include="Services\SearchSideBySideServiceFacts.cs" />
     <Compile Include="Services\PackageUpdateServiceFacts.cs" />
     <Compile Include="Services\StatusServiceFacts.cs" />
+    <Compile Include="Telemetry\CustomerResourceIdEnricherFacts.cs" />
     <Compile Include="TestData\TestDataResourceUtility.cs" />
     <Compile Include="UsernameValidationRegex.cs" />
     <Compile Include="Extensions\NumberExtensionsFacts.cs" />

tests/NuGetGallery.Facts/Telemetry/CustomerResourceIdEnricherFacts.cs

@@ -0,0 +1,128 @@
+// 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.Collections.Generic;
+using System.Linq;
+using System.Security.Claims;
+using System.Security.Principal;
+using System.Web;
+using Microsoft.ApplicationInsights.Channel;
+using Microsoft.ApplicationInsights.DataContracts;
+using Moq;
+using NuGetGallery.Authentication;
+using Xunit;
+
+namespace NuGetGallery.Telemetry
+{
+    public class CustomerResourceIdEnricherFacts
+    {
+        private class TestableCustomerResourceIdEnricher : CustomerResourceIdEnricher
+        {
+            private readonly HttpContextBase _httpContextBase;
+
+            public TestableCustomerResourceIdEnricher(HttpContextBase httpContextBase)
+            {
+                _httpContextBase = httpContextBase;
+            }
+
+            protected override HttpContextBase GetHttpContext()
+            {
+                return _httpContextBase;
+            }
+        }
+
+        [Theory]
+        [InlineData(typeof(RequestTelemetry), 0)]
+        [InlineData(typeof(DependencyTelemetry), 0)]
+        [InlineData(typeof(TraceTelemetry), 0)]
+        [InlineData(typeof(ExceptionTelemetry), 0)]
+        [InlineData(typeof(MetricTelemetry), 1)]
+        public void EnrichesOnlyMetricTelemetry(Type telemetryType, int addedProperties)
+        {
+            // Arrange
+            var telemetry = (ITelemetry)Activator.CreateInstance(telemetryType);
+            if (telemetry is MetricTelemetry metric)
+            {
+                metric.Name = "PackagePush";
+            }
+
+            var itemTelemetry = telemetry as ISupportProperties;
+            itemTelemetry.Properties.Add("Test", "blala");
+
+            var enricher = CreateTestEnricher(new Dictionary<string, string> { { MicrosoftClaims.TenantId, "tenant-id" } });
+
+            // Act
+            enricher.Initialize(telemetry);
+
+            // Assert
+            Assert.Equal(1 + addedProperties, itemTelemetry.Properties.Count);
+        }
+
+        [Fact]
+        public void DoesNotEnrichMetricNotInAllowList()
+        {
+            // Arrange
+            var telemetry = new MetricTelemetry { Name = "SomethingElse" };
+
+            var enricher = CreateTestEnricher(new Dictionary<string, string> { { MicrosoftClaims.TenantId, "tenant-id" } });
+
+            // Act
+            enricher.Initialize(telemetry);
+
+            // Assert
+            Assert.Empty(telemetry.Properties);
+        }
+
+        [Theory]
+        [MemberData(nameof(MetricNames))]
+        public void EnrichesTelemetryWithTenantId(string name)
+        {
+            // Arrange
+            var telemetry = new MetricTelemetry { Name = name };
+
+            var enricher = CreateTestEnricher(new Dictionary<string, string> { { MicrosoftClaims.TenantId, "tenant-id" } });
+
+            // Act
+            enricher.Initialize(telemetry);
+
+            // Assert
+            Assert.Equal("/tenants/tenant-id", telemetry.Properties["CustomerResourceId"]);
+        }
+
+        [Theory]
+        [MemberData(nameof(MetricNames))]
+        public void EnrichesTelemetryWithEmptyWhenTenantIdIsNotInClaims(string name)
+        {
+            // Arrange
+            var telemetry = new MetricTelemetry { Name = name };
+
+            var enricher = CreateTestEnricher(new Dictionary<string, string>());
+
+            // Act
+            enricher.Initialize(telemetry);
+
+            // Assert
+            Assert.Equal("/tenants/00000000-0000-0000-0000-000000000000", telemetry.Properties["CustomerResourceId"]);
+        }
+
+        private TestableCustomerResourceIdEnricher CreateTestEnricher(IReadOnlyDictionary<string, string> claims)
+        {
+            var claimsIdentity = new ClaimsIdentity(claims.Select(x => new Claim(x.Key, x.Value)));
+
+            var principal = new Mock<IPrincipal>();
+            principal.Setup(x => x.Identity).Returns(claimsIdentity);
+
+            var httpContext = new Mock<HttpContextBase>(MockBehavior.Strict);
+            httpContext.SetupGet(c => c.User).Returns(principal.Object);
+
+            return new TestableCustomerResourceIdEnricher(httpContext.Object);
+        }
+
+        public static IEnumerable<object[]> MetricNames => new[]
+        {
+            new object[] { "PackagePush" },
+            new object[] { "PackagePushFailure" },
+        };
+    }
+}