From 6755be0796705828c8f912f8de62f0e96fdefd60 Mon Sep 17 00:00:00 2001
From: maliming <malimings@gmail.com>
Date: Wed, 6 Mar 2024 10:48:04 +0800
Subject: [PATCH] Add `AbpRefreshingPrincipalOptions`.

---
 .../AspNetCore/AbpIdentityAspNetCoreModule.cs   | 14 +++++++++-----
 .../AspNetCore/AbpRefreshingPrincipalOptions.cs | 17 +++++++++++++++++
 .../AbpSecurityStampValidatorCallback.cs        | 17 ++++++++++++++++-
 .../SecurityStampValidatorOptionsExtensions.cs  |  4 ++--
 4 files changed, 44 insertions(+), 8 deletions(-)
 create mode 100644 modules/identity/src/Volo.Abp.Identity.AspNetCore/Volo/Abp/Identity/AspNetCore/AbpRefreshingPrincipalOptions.cs

diff --git a/modules/identity/src/Volo.Abp.Identity.AspNetCore/Volo/Abp/Identity/AspNetCore/AbpIdentityAspNetCoreModule.cs b/modules/identity/src/Volo.Abp.Identity.AspNetCore/Volo/Abp/Identity/AspNetCore/AbpIdentityAspNetCoreModule.cs
index 9b38e2e25a..d261068adb 100644
--- a/modules/identity/src/Volo.Abp.Identity.AspNetCore/Volo/Abp/Identity/AspNetCore/AbpIdentityAspNetCoreModule.cs
+++ b/modules/identity/src/Volo.Abp.Identity.AspNetCore/Volo/Abp/Identity/AspNetCore/AbpIdentityAspNetCoreModule.cs
@@ -1,5 +1,7 @@
-using Microsoft.AspNetCore.Identity;
+using System;
+using Microsoft.AspNetCore.Identity;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Options;
 using Volo.Abp.Modularity;
 using static Volo.Abp.Identity.AspNetCore.AbpSecurityStampValidatorCallback;
 
@@ -45,9 +47,11 @@ public class AbpIdentityAspNetCoreModule : AbpModule
 
     public override void PostConfigureServices(ServiceConfigurationContext context)
     {
-        Configure<SecurityStampValidatorOptions>(options =>
-        {
-            options.UpdatePrincipal();
-        });
+        context.Services.AddOptions<SecurityStampValidatorOptions>()
+            .Configure<IServiceProvider>((securityStampValidatorOptions, serviceProvider) =>
+            {
+                var abpRefreshingPrincipalOptions = serviceProvider.GetRequiredService<IOptions<AbpRefreshingPrincipalOptions>>().Value;
+                securityStampValidatorOptions.UpdatePrincipal(abpRefreshingPrincipalOptions);
+            });
     }
 }
diff --git a/modules/identity/src/Volo.Abp.Identity.AspNetCore/Volo/Abp/Identity/AspNetCore/AbpRefreshingPrincipalOptions.cs b/modules/identity/src/Volo.Abp.Identity.AspNetCore/Volo/Abp/Identity/AspNetCore/AbpRefreshingPrincipalOptions.cs
new file mode 100644
index 0000000000..55e9007d6c
--- /dev/null
+++ b/modules/identity/src/Volo.Abp.Identity.AspNetCore/Volo/Abp/Identity/AspNetCore/AbpRefreshingPrincipalOptions.cs
@@ -0,0 +1,17 @@
+using System.Collections.Generic;
+using Volo.Abp.Security.Claims;
+
+namespace Volo.Abp.Identity.AspNetCore;
+
+public class AbpRefreshingPrincipalOptions
+{
+    public List<string> CurrentPrincipalKeepClaimTypes { get; set; }
+
+    public AbpRefreshingPrincipalOptions()
+    {
+        CurrentPrincipalKeepClaimTypes = new List<string>
+        {
+            AbpClaimTypes.SessionId
+        };
+    }
+}
diff --git a/modules/identity/src/Volo.Abp.Identity.AspNetCore/Volo/Abp/Identity/AspNetCore/AbpSecurityStampValidatorCallback.cs b/modules/identity/src/Volo.Abp.Identity.AspNetCore/Volo/Abp/Identity/AspNetCore/AbpSecurityStampValidatorCallback.cs
index 63bed70dfd..7108d30a5b 100644
--- a/modules/identity/src/Volo.Abp.Identity.AspNetCore/Volo/Abp/Identity/AspNetCore/AbpSecurityStampValidatorCallback.cs
+++ b/modules/identity/src/Volo.Abp.Identity.AspNetCore/Volo/Abp/Identity/AspNetCore/AbpSecurityStampValidatorCallback.cs
@@ -1,5 +1,7 @@
 using System.Linq;
+using System.Security.Principal;
 using System.Threading.Tasks;
+using JetBrains.Annotations;
 using Microsoft.AspNetCore.Identity;
 
 namespace Volo.Abp.Identity.AspNetCore;
@@ -17,8 +19,9 @@ public class AbpSecurityStampValidatorCallback
         /// This is needed to preserve claims such as idp, auth_time, amr.
         /// </summary>
         /// <param name="context">The context.</param>
+        /// <param name="refreshingPrincipalOptions">The AbpRefreshingPrincipalOptions.</param>
         /// <returns></returns>
-        public static Task UpdatePrincipal(SecurityStampRefreshingPrincipalContext context)
+        public static Task UpdatePrincipal(SecurityStampRefreshingPrincipalContext context, AbpRefreshingPrincipalOptions refreshingPrincipalOptions)
         {
             var newClaimTypes = context.NewPrincipal.Claims.Select(x => x.Type).ToArray();
             var currentClaimsToKeep = context.CurrentPrincipal.Claims.Where(x => !newClaimTypes.Contains(x.Type)).ToArray();
@@ -26,6 +29,18 @@ public class AbpSecurityStampValidatorCallback
             var id = context.NewPrincipal.Identities.First();
             id.AddClaims(currentClaimsToKeep);
 
+            if (refreshingPrincipalOptions.CurrentPrincipalKeepClaimTypes.Any())
+            {
+                foreach (var claimType in refreshingPrincipalOptions.CurrentPrincipalKeepClaimTypes)
+                {
+                    var sessionIdClaim = context.CurrentPrincipal.Claims.FirstOrDefault(x => x.Type == claimType);
+                    if (sessionIdClaim != null)
+                    {
+                        id.AddOrReplace(sessionIdClaim);
+                    }
+                }
+            }
+
             return Task.CompletedTask;
         }
     }
diff --git a/modules/identity/src/Volo.Abp.Identity.AspNetCore/Volo/Abp/Identity/AspNetCore/SecurityStampValidatorOptionsExtensions.cs b/modules/identity/src/Volo.Abp.Identity.AspNetCore/Volo/Abp/Identity/AspNetCore/SecurityStampValidatorOptionsExtensions.cs
index d1c51fbf41..009e170310 100644
--- a/modules/identity/src/Volo.Abp.Identity.AspNetCore/Volo/Abp/Identity/AspNetCore/SecurityStampValidatorOptionsExtensions.cs
+++ b/modules/identity/src/Volo.Abp.Identity.AspNetCore/Volo/Abp/Identity/AspNetCore/SecurityStampValidatorOptionsExtensions.cs
@@ -5,12 +5,12 @@ namespace Volo.Abp.Identity.AspNetCore;
 
 public static class SecurityStampValidatorOptionsExtensions
 {
-    public static SecurityStampValidatorOptions UpdatePrincipal(this SecurityStampValidatorOptions options)
+    public static SecurityStampValidatorOptions UpdatePrincipal(this SecurityStampValidatorOptions options, AbpRefreshingPrincipalOptions abpRefreshingPrincipalOptions)
     {
         var previousOnRefreshingPrincipal = options.OnRefreshingPrincipal;
         options.OnRefreshingPrincipal = async context =>
         {
-            await SecurityStampValidatorCallback.UpdatePrincipal(context);
+            await SecurityStampValidatorCallback.UpdatePrincipal(context, abpRefreshingPrincipalOptions);
             if(previousOnRefreshingPrincipal != null)
             {
                 await previousOnRefreshingPrincipal.Invoke(context);