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

Progress on NuGet/Engineering#3106

Author: joelverhagen

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();
+    }
+}