From a3038d8b0a57b2cd1ff259c975c675dd1572b2ce Mon Sep 17 00:00:00 2001
From: Andrei Alistar <andrei.alistar@genisoft.eu>
Date: Thu, 1 Aug 2024 12:01:21 +0300
Subject: [PATCH] AccessToken security group retrieval strategy.

---
 .../Authentication/EntraUserClaimConstants.cs | 13 +++++++++++
 .../SecurityGroupRetrievalStrategies.cs       |  5 +++++
 .../Interfaces/IUserClaimsProviderService.cs  |  9 ++++++++
 .../Middleware/CallContextMiddleware.cs       | 22 ++++++++++++++-----
 .../EntraUserClaimsProviderService.cs         | 13 ++++++++++-
 .../Security/NoOpUserClaimsProviderService.cs |  8 +++----
 6 files changed, 58 insertions(+), 12 deletions(-)
 create mode 100644 src/dotnet/Common/Constants/Authentication/EntraUserClaimConstants.cs

diff --git a/src/dotnet/Common/Constants/Authentication/EntraUserClaimConstants.cs b/src/dotnet/Common/Constants/Authentication/EntraUserClaimConstants.cs
new file mode 100644
index 0000000000..289b77f959
--- /dev/null
+++ b/src/dotnet/Common/Constants/Authentication/EntraUserClaimConstants.cs
@@ -0,0 +1,13 @@
+namespace FoundationaLLM.Common.Constants.Authentication
+{
+    /// <summary>
+    /// Additional claim constants that are not included in <see cref="Microsoft.Identity.Web.ClaimConstants"></see>.
+    /// </summary>
+    public static class EntraUserClaimConstants
+    {
+        /// <summary>
+        /// Groups claim.
+        /// </summary>
+        public const string Groups = "groups";
+    }
+}
diff --git a/src/dotnet/Common/Constants/Instance/SecurityGroupRetrievalStrategies.cs b/src/dotnet/Common/Constants/Instance/SecurityGroupRetrievalStrategies.cs
index 65cf0c09df..518cec3d81 100644
--- a/src/dotnet/Common/Constants/Instance/SecurityGroupRetrievalStrategies.cs
+++ b/src/dotnet/Common/Constants/Instance/SecurityGroupRetrievalStrategies.cs
@@ -14,5 +14,10 @@ public static class SecurityGroupRetrievalStrategies
         /// Identity management service.
         /// </summary>
         public const string IdentityManagementService = "IdentityManagementService";
+
+        /// <summary>
+        /// Access token groups claim.
+        /// </summary>
+        public const string AccessToken = "AccessToken";
     }
 }
diff --git a/src/dotnet/Common/Interfaces/IUserClaimsProviderService.cs b/src/dotnet/Common/Interfaces/IUserClaimsProviderService.cs
index b248cf9666..d8ed34e4d2 100644
--- a/src/dotnet/Common/Interfaces/IUserClaimsProviderService.cs
+++ b/src/dotnet/Common/Interfaces/IUserClaimsProviderService.cs
@@ -28,5 +28,14 @@ public interface IUserClaimsProviderService
         /// </summary>
         /// <param name="userPrincipal">The <see cref="ClaimsPrincipal"/> object providing details about the security principal.</param>
         bool IsServicePrincipal(ClaimsPrincipal userPrincipal);
+
+        /// <summary>
+        /// Returns a list of security group identifiers from the provided
+        /// <see cref="ClaimsPrincipal"/>.
+        /// </summary>
+        /// <param name="userPrincipal">The principal that provides multiple
+        /// claims-based identities.</param>
+        /// <returns></returns>
+        List<string>? GetSecurityGroupIds(ClaimsPrincipal? userPrincipal);
     }
 }
diff --git a/src/dotnet/Common/Middleware/CallContextMiddleware.cs b/src/dotnet/Common/Middleware/CallContextMiddleware.cs
index 317b014974..eb1acf3bf4 100644
--- a/src/dotnet/Common/Middleware/CallContextMiddleware.cs
+++ b/src/dotnet/Common/Middleware/CallContextMiddleware.cs
@@ -45,14 +45,24 @@ public async Task InvokeAsync(
                 // Extract from ClaimsPrincipal if available:
                 callContext.CurrentUserIdentity = claimsProviderService.GetUserIdentity(context.User);
 
+                // We are only expanding group membership for User objects
+                // Service Principal permissions must be assigned directly and not over group membership.
                 if (callContext.CurrentUserIdentity != null
-                    && !claimsProviderService.IsServicePrincipal(context.User)
-                    && instanceSettings.Value.SecurityGroupRetrievalStrategy == SecurityGroupRetrievalStrategies.IdentityManagementService)
+                    && !claimsProviderService.IsServicePrincipal(context.User))
                 {
-                    // We are only expanding group membership for User objects
-                    // Service Principal permissions must be assigned directly and not over group membership.
-                    callContext.CurrentUserIdentity.GroupIds = await identityManagementService.GetGroupsForPrincipal(
-                        callContext.CurrentUserIdentity.UserId!);
+                    switch(instanceSettings.Value.SecurityGroupRetrievalStrategy)
+                    {
+                        case SecurityGroupRetrievalStrategies.IdentityManagementService:
+                            callContext.CurrentUserIdentity.GroupIds = await identityManagementService.GetGroupsForPrincipal(
+                                callContext.CurrentUserIdentity.UserId!);
+                            break;
+                        case SecurityGroupRetrievalStrategies.AccessToken:
+                            callContext.CurrentUserIdentity.GroupIds = claimsProviderService.GetSecurityGroupIds(context.User) ?? [];
+                            break;
+                        case SecurityGroupRetrievalStrategies.None:
+                        default:
+                            break;
+                    }
                 }
             }
             else
diff --git a/src/dotnet/Common/Services/Security/EntraUserClaimsProviderService.cs b/src/dotnet/Common/Services/Security/EntraUserClaimsProviderService.cs
index a164126301..278c2ca7ad 100644
--- a/src/dotnet/Common/Services/Security/EntraUserClaimsProviderService.cs
+++ b/src/dotnet/Common/Services/Security/EntraUserClaimsProviderService.cs
@@ -1,4 +1,5 @@
-using FoundationaLLM.Common.Interfaces;
+using FoundationaLLM.Common.Constants.Authentication;
+using FoundationaLLM.Common.Interfaces;
 using FoundationaLLM.Common.Models.Authentication;
 using Microsoft.Identity.Web;
 using System.Security.Claims;
@@ -27,6 +28,16 @@ public class EntraUserClaimsProviderService : IUserClaimsProviderService
             };
         }
 
+        /// <inheritdoc/>
+        public List<string>? GetSecurityGroupIds(ClaimsPrincipal? userPrincipal)
+        {
+            if (userPrincipal == null)
+            {
+                return null;
+            }
+            return userPrincipal.Claims.Where(c => c.Type == EntraUserClaimConstants.Groups).Select(x => x?.Value).ToList()!;
+        }
+
         /// <inheritdoc/>
         public bool IsServicePrincipal(ClaimsPrincipal userPrincipal) =>
             // Service Principal tokens do not have a "scp" claim
diff --git a/src/dotnet/Common/Services/Security/NoOpUserClaimsProviderService.cs b/src/dotnet/Common/Services/Security/NoOpUserClaimsProviderService.cs
index 1a4e566f75..00733ce808 100644
--- a/src/dotnet/Common/Services/Security/NoOpUserClaimsProviderService.cs
+++ b/src/dotnet/Common/Services/Security/NoOpUserClaimsProviderService.cs
@@ -1,11 +1,6 @@
 using FoundationaLLM.Common.Interfaces;
 using FoundationaLLM.Common.Models.Authentication;
-using System;
-using System.Collections.Generic;
-using System.Linq;
 using System.Security.Claims;
-using System.Text;
-using System.Threading.Tasks;
 
 namespace FoundationaLLM.Common.Services.Security
 {
@@ -18,6 +13,9 @@ public class NoOpUserClaimsProviderService : IUserClaimsProviderService
         /// <inheritdoc/>
         public UnifiedUserIdentity? GetUserIdentity(ClaimsPrincipal? userPrincipal) => null;
 
+        /// <inheritdoc/>
+        public List<string>? GetSecurityGroupIds(ClaimsPrincipal? userPrincipal) => null;
+
         /// <inheritdoc/>
         public bool IsServicePrincipal(ClaimsPrincipal userPrincipal) => false;
     }