custom DiscoveryResponseGenerator and helpers

src/Core/Utilities/CoreHelpers.cs

@@ -646,5 +646,30 @@ namespace Bit.Core.Utilities
             }
             return null;
         }
+
+        public static Dictionary<string, object> AdjustIdentityServerConfig(Dictionary<string, object> configDict,
+            string publicServiceUri, string internalServiceUri)
+        {
+            var dictReplace = new Dictionary<string, object>();
+            foreach (var item in configDict)
+            {
+                var change = item.Key.EndsWith("_endpoint") || item.Key.EndsWith("_iframe");
+                if (change && item.Value is string val)
+                {
+                    var uri = new Uri(val);
+                    dictReplace.Add(item.Key, string.Concat(publicServiceUri, uri.LocalPath));
+                }
+                else if (item.Key == "jwks_uri" && item.Value is string jwksVal)
+                {
+                    var uri = new Uri(jwksVal);
+                    dictReplace.Add(item.Key, string.Concat(internalServiceUri, uri.LocalPath));
+                }
+            }
+            foreach (var replace in dictReplace)
+            {
+                configDict[replace.Key] = replace.Value;
+            }
+            return configDict;
+        }
     }
 }

src/Identity/Startup.cs

@@ -12,10 +12,7 @@ using Microsoft.Extensions.Hosting;
 using Microsoft.IdentityModel.Logging;
 using System.IdentityModel.Tokens.Jwt;
 using System.Threading.Tasks;
-using IdentityServer4.Stores;
+using Bit.Identity.Utilities;
-using Bit.Core.IdentityServer;
-using IdentityServer4.Services;
-using Microsoft.IdentityModel.Protocols.OpenIdConnect;
 using IdentityServer4.Extensions;
 
 namespace Bit.Identity
@@ -112,7 +109,7 @@ namespace Bit.Identity
                 });
 
             // IdentityServer
-            AddCustomIdentityServerServices(services, Environment, globalSettings);
+            services.AddCustomIdentityServerServices(Environment, globalSettings);
 
             // Identity
             services.AddCustomIdentityServices(globalSettings);
@@ -148,6 +145,12 @@ namespace Bit.Identity
 
             if (globalSettings.SelfHosted)
             {
+                var uri = new Uri(globalSettings.BaseServiceUri.Identity);
+                app.Use(async (ctx, next) =>
+                {
+                    ctx.SetIdentityServerOrigin($"{uri.Scheme}://{uri.Host}");
+                    await next();
+                });
                 app.UsePathBase("/identity");
                 app.UseForwardedHeaders(globalSettings);
             }
@@ -192,42 +195,5 @@ namespace Bit.Identity
             // Log startup
             logger.LogInformation(Constants.BypassFiltersEventId, globalSettings.ProjectName + " started.");
         }
-
-        public static IIdentityServerBuilder AddCustomIdentityServerServices(IServiceCollection services,
-            IWebHostEnvironment env, GlobalSettings globalSettings)
-        {
-            services.AddSingleton<StaticClientStore>();
-            services.AddTransient<IAuthorizationCodeStore, AuthorizationCodeStore>();
-
-            var issuerUri = new Uri(globalSettings.BaseServiceUri.InternalIdentity);
-            var identityServerBuilder = services
-                .AddIdentityServer(options =>
-                {
-                    options.Endpoints.EnableIntrospectionEndpoint = false;
-                    options.Endpoints.EnableEndSessionEndpoint = false;
-                    options.Endpoints.EnableUserInfoEndpoint = false;
-                    options.Endpoints.EnableCheckSessionEndpoint = false;
-                    options.Endpoints.EnableTokenRevocationEndpoint = false;
-                    options.IssuerUri = $"{issuerUri.Scheme}://{issuerUri.Host}";
-                    options.Caching.ClientStoreExpiration = new TimeSpan(0, 5, 0);
-                    if (env.IsDevelopment())
-                    {
-                        options.Authentication.CookieSameSiteMode = Microsoft.AspNetCore.Http.SameSiteMode.Unspecified;
-                    }
-                })
-                .AddInMemoryCaching()
-                .AddInMemoryApiResources(ApiResources.GetApiResources())
-                .AddInMemoryApiScopes(ApiScopes.GetApiScopes())
-                .AddClientStoreCache<ClientStore>()
-                .AddCustomTokenRequestValidator<CustomTokenRequestValidator>()
-                .AddProfileService<ProfileService>()
-                .AddResourceOwnerValidator<ResourceOwnerPasswordValidator>()
-                .AddPersistedGrantStore<PersistedGrantStore>()
-                .AddClientStore<ClientStore>()
-                .AddIdentityServerCertificate(env, globalSettings);
-
-            services.AddTransient<ICorsPolicyService, CustomCorsPolicyService>();
-            return identityServerBuilder;
-        }
     }
 }

src/Identity/Utilities/DiscoveryResponseGenerator.cs

@@ -0,0 +1,39 @@
+using System.Collections.Generic;
+using System.Threading.Tasks;
+using Bit.Core;
+using Bit.Core.Utilities;
+using IdentityServer4.Configuration;
+using IdentityServer4.Services;
+using IdentityServer4.Stores;
+using IdentityServer4.Validation;
+using Microsoft.Extensions.Logging;
+
+namespace Bit.Identity.Utilities
+{
+    public class DiscoveryResponseGenerator : IdentityServer4.ResponseHandling.DiscoveryResponseGenerator
+    {
+        private readonly GlobalSettings _globalSettings;
+
+        public DiscoveryResponseGenerator(
+            IdentityServerOptions options,
+            IResourceStore resourceStore,
+            IKeyMaterialService keys,
+            ExtensionGrantValidator extensionGrants,
+            ISecretsListParser secretParsers,
+            IResourceOwnerPasswordValidator resourceOwnerValidator,
+            ILogger<DiscoveryResponseGenerator> logger,
+            GlobalSettings globalSettings)
+            : base(options, resourceStore, keys, extensionGrants, secretParsers, resourceOwnerValidator, logger)
+        {
+            _globalSettings = globalSettings;
+        }
+
+        public override async Task<Dictionary<string, object>> CreateDiscoveryDocumentAsync(
+            string baseUrl, string issuerUri)
+        {
+            var dict = await base.CreateDiscoveryDocumentAsync(baseUrl, issuerUri);
+            return CoreHelpers.AdjustIdentityServerConfig(dict, _globalSettings.BaseServiceUri.Identity,
+                _globalSettings.BaseServiceUri.InternalIdentity);
+        }
+    }
+}

src/Identity/Utilities/ServiceCollectionExtensions.cs

@@ -0,0 +1,58 @@
+using System;
+using Bit.Core;
+using Bit.Core.IdentityServer;
+using Bit.Core.Utilities;
+using IdentityServer4.ResponseHandling;
+using IdentityServer4.Services;
+using IdentityServer4.Stores;
+using Microsoft.AspNetCore.Hosting;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
+
+namespace Bit.Identity.Utilities
+{
+    public static class ServiceCollectionExtensions
+    {
+        public static IIdentityServerBuilder AddCustomIdentityServerServices(this IServiceCollection services,
+            IWebHostEnvironment env, GlobalSettings globalSettings)
+        {
+            if (globalSettings.SelfHosted)
+            {
+                services.AddTransient<IDiscoveryResponseGenerator, DiscoveryResponseGenerator>();
+            }
+
+            services.AddSingleton<StaticClientStore>();
+            services.AddTransient<IAuthorizationCodeStore, AuthorizationCodeStore>();
+
+            var issuerUri = new Uri(globalSettings.BaseServiceUri.InternalIdentity);
+            var identityServerBuilder = services
+                .AddIdentityServer(options =>
+                {
+                    options.Endpoints.EnableIntrospectionEndpoint = false;
+                    options.Endpoints.EnableEndSessionEndpoint = false;
+                    options.Endpoints.EnableUserInfoEndpoint = false;
+                    options.Endpoints.EnableCheckSessionEndpoint = false;
+                    options.Endpoints.EnableTokenRevocationEndpoint = false;
+                    options.IssuerUri = $"{issuerUri.Scheme}://{issuerUri.Host}";
+                    options.Caching.ClientStoreExpiration = new TimeSpan(0, 5, 0);
+                    if (env.IsDevelopment())
+                    {
+                        options.Authentication.CookieSameSiteMode = Microsoft.AspNetCore.Http.SameSiteMode.Unspecified;
+                    }
+                })
+                .AddInMemoryCaching()
+                .AddInMemoryApiResources(ApiResources.GetApiResources())
+                .AddInMemoryApiScopes(ApiScopes.GetApiScopes())
+                .AddClientStoreCache<ClientStore>()
+                .AddCustomTokenRequestValidator<CustomTokenRequestValidator>()
+                .AddProfileService<ProfileService>()
+                .AddResourceOwnerValidator<ResourceOwnerPasswordValidator>()
+                .AddPersistedGrantStore<PersistedGrantStore>()
+                .AddClientStore<ClientStore>()
+                .AddIdentityServerCertificate(env, globalSettings);
+
+            services.AddTransient<ICorsPolicyService, CustomCorsPolicyService>();
+            return identityServerBuilder;
+        }
+    }
+}