IdentityServer4 Tutorial - Part 3: Store Refresh Token in Cosmos DB

By default refresh tokens are stored in memory. In this tutorial we will add an IPersistedGrantStore implementation to store refresh tokens in Cosmos DB. Cosmos DB provides 5 APIs. We will use SQL API with Version 3.0+ of the Azure Cosmos DB .NET SDK. The work is based on IdentityServer4 Tutorial - Part 2: Resource Owner Password Grant Type.

Create an Azure Cosmos DB

Follow this document to create an Azure Cosmos DB. When create your container, use ‘/id’ as partition key.

Nuget Update

Install nuget package Microsoft.Azure.Cosmos. By the time the tutorial is written, the package is still under preview. You need to tick ‘Include prerelease’ to see the package in Package Manager.

Define an Azure Cosmos Item

Refresh tokens are stored as a PersistedGrant object. But we can’t save it directly to Azure Cosmos DB, because an Azure Cosmos Item requires an id property. Let’s add a PersistedGrantItem class as below.

public class PersistedGrantItem
{
    [JsonProperty(PropertyName = "id")]
    public string Key;

    [JsonProperty(PropertyName = "type")]
    public string Type;

    [JsonProperty(PropertyName = "subject_id")]
    public string SubjectId;

    [JsonProperty(PropertyName = "client_id")]
    public string ClientId;

    [JsonProperty(PropertyName = "creation_time")]
    public DateTime CreationTime;

    [JsonProperty(PropertyName = "expiration")]
    public DateTime? Expiration;

    [JsonProperty(PropertyName = "data")]
    public string Data;
}

All the properties are the same as in PersistedGrant. But we map Key to id and change other property names to follow the naming convention.

Let’s create a PersistedGrantMapper class to make conversion between the two easier.

public static class PersistedGrantMapper
{
    public static PersistedGrantItem ToItem(this PersistedGrant persistedGrant)
    {
        var persistedGrantItem = new PersistedGrantItem
        {
            Key = persistedGrant.Key,
            Type = persistedGrant.Type,
            SubjectId = persistedGrant.SubjectId,
            ClientId = persistedGrant.ClientId,
            CreationTime = persistedGrant.CreationTime,
            Expiration = persistedGrant.Expiration,
            Data = persistedGrant.Data
        };
        return persistedGrantItem;
    }

    public static PersistedGrant ToModel(this PersistedGrantItem persistedGrantItem)
    {
        var persistedGrant = new PersistedGrant
        {
            Key = persistedGrantItem.Key,
            Type = persistedGrantItem.Type,
            SubjectId = persistedGrantItem.SubjectId,
            ClientId = persistedGrantItem.ClientId,
            CreationTime = persistedGrantItem.CreationTime,
            Expiration = persistedGrantItem.Expiration,
            Data = persistedGrantItem.Data
        };
        return persistedGrant;
    }
}

CosmosDbPersistedGrantStore

Now we can add CosmosDbPersistedGrantStore.

public class CosmosDbPersistedGrantStore : IPersistedGrantStore
{
    private readonly Container _container;

    public CosmosDbPersistedGrantStore(IConfiguration configuration, CosmosClient cosmosClient)
    {
        _container = cosmosClient.GetContainer(
            configuration["CosmosDB:DatabaseId"],
            configuration["CosmosDB:ContainerId"]);
    }

    public Task StoreAsync(PersistedGrant grant)
    {
        grant.Key = EncodePersistedGrantKey(grant.Key);
        var persistedGrantItem = grant.ToItem();
        return _container.CreateItemAsync(persistedGrantItem);
    }

    public async Task<PersistedGrant> GetAsync(string key)
    {
        key = EncodePersistedGrantKey(key);

        var result = await _container.ReadItemAsync<PersistedGrantItem>(key, new PartitionKey(key));
        if (result.StatusCode == HttpStatusCode.NotFound)
        {
            return null;
        }

        var persistedGrantItem = (PersistedGrantItem)result;
        var persistedGrant = persistedGrantItem.ToModel();
        return persistedGrant;
    }

    public Task<IEnumerable<PersistedGrant>> GetAllAsync(string subjectId)
    {
        throw new NotImplementedException();
    }

    public Task RemoveAsync(string key)
    {
        key = EncodePersistedGrantKey(key);
        return _container.DeleteItemAsync<PersistedGrantItem>(key, new PartitionKey(key));
    }

    public Task RemoveAllAsync(string subjectId, string clientId)
    {
        throw new NotImplementedException();
    }

    public Task RemoveAllAsync(string subjectId, string clientId, string type)
    {
        throw new NotImplementedException();
    }

    private static string EncodePersistedGrantKey(string key)
    {
        key = Base64UrlEncoder.Encode(key);
        return key;
    }
}

We need to encode the key to make it url friendly.

Add a CosmosDB section into appsettings.json.

{
  "Logging": {
    "LogLevel": {
      "Default": "Warning"
    }
  },
  "AllowedHosts": "*",
  "CosmosDB": {
    "AccountEndpoint": "your account endpoint",
    "AccountKey": "your account key",
    "DatabaseId": "your database id",
    "ContainerId": "your container id"
  } 
}

Startup

Finally we need to tell IdentityServer to use CosmosDbPersistedGrantStore and add a singleton binding for CosmosClient. They are both configured through Startup class.

public class Startup
{
    private readonly IConfiguration _configuration;

    public Startup(IConfiguration configuration)
    {
        _configuration = configuration;
    }

    // This method gets called by the runtime. Use this method to add services to the container.
    // For more information on how to configure your application, visit https://go.microsoft.com/fwlink/?LinkID=398940
    public void ConfigureServices(IServiceCollection services)
    {
        var builder = services.AddIdentityServer()
            .AddInMemoryClients(new[]
            {
                new Client
                {
                    ClientId = "client_id1",
                    RequireClientSecret = false,
                    AllowedGrantTypes = { GrantType.ResourceOwnerPassword },
                    AllowedScopes = { "api1" },
                    AllowOfflineAccess = true
                }
            })
            .AddInMemoryApiResources(new[]
            {
                new ApiResource("api1", new [] { "custom_claim" })
            })
            .AddResourceOwnerValidator<ResourceOwnerPasswordValidator>()
            .AddPersistedGrantStore<CosmosDbPersistedGrantStore>()
            .AddDeveloperSigningCredential();

        services.AddSingleton(s =>
        {
            var accountEndpoint = _configuration["CosmosDB:AccountEndpoint"];
            var accountKey = _configuration["CosmosDB:AccountKey"];
            var configurationBuilder = new CosmosClientBuilder(accountEndpoint, accountKey);
            return configurationBuilder.Build();
        });
    }

    // This method gets called by the runtime. Use this method to configure the HTTP request pipeline.
    public void Configure(IApplicationBuilder app, IHostingEnvironment env)
    {
        app.UseIdentityServer();

        if (env.IsDevelopment())
        {
            app.UseDeveloperExceptionPage();
        }

        app.Run(async (context) =>
        {
            await context.Response.WriteAsync("Hello World!");
        });
    }
}

Refresh Token Lifetime

The lifetime of a refresh token is configured via client setting AbsoluteRefreshTokenLifetime. Normally we would need to create a task to delete expired refresh tokens. But Azure Cosmos DB has a nice Time to live feature. When it is configured, expired tokens will be deleted automatically.