From 5003614b0c5ff976614f3056d7132385b039b2b5 Mon Sep 17 00:00:00 2001
From: maliming <malimings@gmail.com>
Date: Thu, 17 Apr 2025 10:20:48 +0800
Subject: [PATCH] Try to get tenant's `EditionId` in
 `EditionFeatureManagementProvider`.

---
 .../Features/EditionFeatureValueProvider.cs   |  4 +-
 .../EditionFeatureManagementProvider.cs       | 50 ++++++++++++++++---
 2 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/framework/src/Volo.Abp.Features/Volo/Abp/Features/EditionFeatureValueProvider.cs b/framework/src/Volo.Abp.Features/Volo/Abp/Features/EditionFeatureValueProvider.cs
index e252bc31f0..cdab8e74c8 100644
--- a/framework/src/Volo.Abp.Features/Volo/Abp/Features/EditionFeatureValueProvider.cs
+++ b/framework/src/Volo.Abp.Features/Volo/Abp/Features/EditionFeatureValueProvider.cs
@@ -32,7 +32,7 @@ public class EditionFeatureValueProvider : FeatureValueProvider
 
     public async override Task<string?> GetOrNullAsync(FeatureDefinition feature)
     {
-        var editionId = await GetEditionIdAsync();
+        var editionId = await FindEditionIdAsync();
         if (editionId == null)
         {
             return null;
@@ -41,7 +41,7 @@ public class EditionFeatureValueProvider : FeatureValueProvider
         return await FeatureStore.GetOrNullAsync(feature.Name, Name, editionId.Value.ToString());
     }
 
-    protected virtual async Task<Guid?> GetEditionIdAsync()
+    protected virtual async Task<Guid?> FindEditionIdAsync()
     {
         var editionId = PrincipalAccessor.Principal?.FindEditionId();
         if (editionId != null)
diff --git a/modules/feature-management/src/Volo.Abp.FeatureManagement.Domain/Volo/Abp/FeatureManagement/EditionFeatureManagementProvider.cs b/modules/feature-management/src/Volo.Abp.FeatureManagement.Domain/Volo/Abp/FeatureManagement/EditionFeatureManagementProvider.cs
index 82c1c98cb9..f9de4cbd1f 100644
--- a/modules/feature-management/src/Volo.Abp.FeatureManagement.Domain/Volo/Abp/FeatureManagement/EditionFeatureManagementProvider.cs
+++ b/modules/feature-management/src/Volo.Abp.FeatureManagement.Domain/Volo/Abp/FeatureManagement/EditionFeatureManagementProvider.cs
@@ -1,7 +1,9 @@
-using System.Security.Principal;
+using System;
+using System.Security.Principal;
 using System.Threading.Tasks;
 using Volo.Abp.DependencyInjection;
 using Volo.Abp.Features;
+using Volo.Abp.MultiTenancy;
 using Volo.Abp.Security.Claims;
 
 namespace Volo.Abp.FeatureManagement;
@@ -11,22 +13,58 @@ public class EditionFeatureManagementProvider : FeatureManagementProvider, ITran
     public override string Name => EditionFeatureValueProvider.ProviderName;
 
     protected ICurrentPrincipalAccessor PrincipalAccessor { get; }
+    protected ITenantStore TenantStore { get; }
+    protected ICurrentTenant CurrentTenant { get; }
+    protected string CurrentCompatibleProviderName { get; set; }
 
     public EditionFeatureManagementProvider(
         IFeatureManagementStore store,
-        ICurrentPrincipalAccessor principalAccessor)
+        ICurrentPrincipalAccessor principalAccessor,
+        ITenantStore tenantStore,
+        ICurrentTenant currentTenant)
         : base(store)
     {
         PrincipalAccessor = principalAccessor;
+        TenantStore = tenantStore;
+        CurrentTenant = currentTenant;
     }
 
-    protected override Task<string> NormalizeProviderKeyAsync(string providerKey)
+    public override bool Compatible(string providerName)
     {
-        if (providerKey != null)
+        CurrentCompatibleProviderName = providerName;
+        return providerName == TenantFeatureValueProvider.ProviderName || base.Compatible(providerName);
+    }
+
+    protected async override Task<string> NormalizeProviderKeyAsync(string providerKey)
+    {
+        return (await FindEditionIdAsync(providerKey))?.ToString();
+    }
+
+    protected virtual async Task<Guid?> FindEditionIdAsync(string providerKey)
+    {
+        if (Guid.TryParse(providerKey, out var parsedEditionOrTenantId))
+        {
+            if (CurrentCompatibleProviderName == TenantFeatureValueProvider.ProviderName)
+            {
+                var tenant = await TenantStore.FindAsync(parsedEditionOrTenantId);
+                if (tenant != null)
+                {
+                    return tenant?.EditionId;
+                }
+            }
+
+            return parsedEditionOrTenantId;
+        }
+
+        if (CurrentTenant.Id.HasValue)
         {
-            return Task.FromResult(providerKey);
+            var tenant = await TenantStore.FindAsync(CurrentTenant.GetId());
+            if (tenant != null)
+            {
+                return tenant?.EditionId;
+            }
         }
 
-        return Task.FromResult(PrincipalAccessor.Principal?.FindEditionId()?.ToString("N"));
+        return PrincipalAccessor.Principal?.FindEditionId();
     }
 }