Azure/web-farm ready SecurityTokenCache

Question

Our site uses ADFS for auth. To reduce the cookie payload on every request we're turning IsSessionMode on (see Your fedauth cookies on a diet).

The last thing we need to do to get this working in our load balanced environment is to implement a farm ready SecurityTokenCache. The implementation seems pretty straightforward, I'm mainly interested in finding out if there are any gotchas we should consider when dealing with SecurityTokenCacheKey and the TryGetAllEntries and TryRemoveAllEntries methods (SecurityTokenCacheKey has a custom implementation of the Equals and GetHashCode methods).

Does anyone have an example of this? We're planning on using AppFabric as the backing store but an example using any persistent store would be helpful- database table, Azure table-storage, etc.

Here are some places I've searched:

• In Hervey Wilson's PDC09 session he uses a DatabaseSecurityTokenCache. I haven't been able to find the sample code for his session.
• On page 192 of Vittorio Bertocci's excellent book, "Programming Windows Identity Foundation" he mentions uploading a sample implementation of an Azure ready SecurityTokenCache to the book's website. I haven't been able to find this sample either.

Thanks!

jd

3/16/2012 UPDATE Vittorio's blog links to a sample using the new .net 4.5 stuff:

ClaimsAwareWebFarm This sample is an answer to the feedback we got from many of you guys: you wanted a sample showing a farm ready session cache (as opposed to a tokenreplycache) so that you can use sessions by reference instead of exchanging big cookies; and you asked for an easier way of securing cookies in a farm.

Answer 1

To come up with a working implementation we ultimately had to use reflector to analyze the different SessionSecurityToken related classes in Microsoft.IdentityModel. Below is what we came up with. This implementation is deployed on our dev and qa environments, seems to be working fine, it's resiliant to app pool recycles etc.

In global.asax:

protected void Application_Start(object sender, EventArgs e)
{
    FederatedAuthentication.ServiceConfigurationCreated += this.OnServiceConfigurationCreated;
}

private void OnServiceConfigurationCreated(object sender, ServiceConfigurationCreatedEventArgs e)
{
    var sessionTransforms = new List<CookieTransform>(new CookieTransform[]
            {
                new DeflateCookieTransform(),
                new RsaEncryptionCookieTransform(
                    e.ServiceConfiguration.ServiceCertificate),
                new RsaSignatureCookieTransform(
                    e.ServiceConfiguration.ServiceCertificate)
            });

    // following line is pseudo code.  use your own durable cache implementation.
    var durableCache = new AppFabricCacheWrapper();

    var tokenCache = new DurableSecurityTokenCache(durableCache, 5000);
    var sessionHandler = new SessionSecurityTokenHandler(sessionTransforms.AsReadOnly(),
        tokenCache,
        TimeSpan.FromDays(1));

    e.ServiceConfiguration.SecurityTokenHandlers.AddOrReplace(sessionHandler);
}

private void WSFederationAuthenticationModule_SecurityTokenValidated(object sender, SecurityTokenValidatedEventArgs e)
{
    FederatedAuthentication.SessionAuthenticationModule.IsSessionMode = true;
}

DurableSecurityTokenCache.cs:

/// <summary>
/// Two level durable security token cache (level 1: in memory MRU, level 2: out of process cache).
/// </summary>
public class DurableSecurityTokenCache : SecurityTokenCache
{
    private ICache<string, byte[]> durableCache;
    private readonly MruCache<SecurityTokenCacheKey, SecurityToken> mruCache;

    /// <summary>
    /// The constructor.
    /// </summary>
    /// <param name="durableCache">The durable second level cache (should be out of process ie sql server, azure table, app fabric, etc).</param>
    /// <param name="mruCapacity">Capacity of the internal first level cache (in-memory MRU cache).</param>
    public DurableSecurityTokenCache(ICache<string, byte[]> durableCache, int mruCapacity)
    {
        this.durableCache = durableCache;
        this.mruCache = new MruCache<SecurityTokenCacheKey, SecurityToken>(mruCapacity, mruCapacity / 4);
    }

    public override bool TryAddEntry(object key, SecurityToken value)
    {
        var cacheKey = (SecurityTokenCacheKey)key;

        // add the entry to the mru cache.
        this.mruCache.Add(cacheKey, value);

        // add the entry to the durable cache.
        var keyString = GetKeyString(cacheKey);
        var buffer = this.GetSerializer().Serialize((SessionSecurityToken)value);
        this.durableCache.Add(keyString, buffer);

        return true;
    }

    public override bool TryGetEntry(object key, out SecurityToken value)
    {
        var cacheKey = (SecurityTokenCacheKey)key;

        // attempt to retrieve the entry from the mru cache.
        value = this.mruCache.Get(cacheKey);
        if (value != null)
            return true;

        // entry wasn't in the mru cache, retrieve it from the app fabric cache.
        var keyString = GetKeyString(cacheKey);

        var buffer = this.durableCache.Get(keyString);
        var result = buffer != null;
        if (result)
        {
            // we had a cache miss in the mru cache but found the item in the durable cache...

            // deserialize the value retrieved from the durable cache.
            value = this.GetSerializer().Deserialize(buffer);

            // push this item into the mru cache.
            this.mruCache.Add(cacheKey, value);
        }

        return result;
    }

    public override bool TryRemoveEntry(object key)
    {
        var cacheKey = (SecurityTokenCacheKey)key;

        // remove the entry from the mru cache.
        this.mruCache.Remove(cacheKey);

        // remove the entry from the durable cache.
        var keyString = GetKeyString(cacheKey);
        this.durableCache.Remove(keyString);

        return true;
    }

    public override bool TryReplaceEntry(object key, SecurityToken newValue)
    {
        var cacheKey = (SecurityTokenCacheKey)key;

        // remove the entry in the mru cache.
        this.mruCache.Remove(cacheKey);

        // remove the entry in the durable cache.
        var keyString = GetKeyString(cacheKey);

        // add the new value.
        return this.TryAddEntry(key, newValue);
    }

    public override bool TryGetAllEntries(object key, out IList<SecurityToken> tokens)
    {
        // not implemented... haven't been able to find how/when this method is used.
        tokens = new List<SecurityToken>();
        return true;
        //throw new NotImplementedException();
    }

    public override bool TryRemoveAllEntries(object key)
    {
        // not implemented... haven't been able to find how/when this method is used.
        return true;
        //throw new NotImplementedException();
    }

    public override void ClearEntries()
    {
        // not implemented... haven't been able to find how/when this method is used.
        //throw new NotImplementedException();
    }

    /// <summary>
    /// Gets the string representation of the specified SecurityTokenCacheKey.
    /// </summary>
    private string GetKeyString(SecurityTokenCacheKey key)
    {
        return string.Format("{0}; {1}; {2}", key.ContextId, key.KeyGeneration, key.EndpointId);
    }

    /// <summary>
    /// Gets a new instance of the token serializer.
    /// </summary>
    private SessionSecurityTokenCookieSerializer GetSerializer()
    {
        return new SessionSecurityTokenCookieSerializer();  // may need to do something about handling bootstrap tokens.
    }
}

MruCache.cs:

/// <summary>
/// Most recently used (MRU) cache.
/// </summary>
/// <typeparam name="TKey">The key type.</typeparam>
/// <typeparam name="TValue">The value type.</typeparam>
public class MruCache<TKey, TValue> : ICache<TKey, TValue>
{
    private Dictionary<TKey, TValue> mruCache;
    private LinkedList<TKey> mruList;
    private object syncRoot;
    private int capacity;
    private int sizeAfterPurge;

    /// <summary>
    /// The constructor.
    /// </summary>
    /// <param name="capacity">The capacity.</param>
    /// <param name="sizeAfterPurge">Size to make the cache after purging because it's reached capacity.</param>
    public MruCache(int capacity, int sizeAfterPurge)
    {
        this.mruList = new LinkedList<TKey>();
        this.mruCache = new Dictionary<TKey, TValue>(capacity);
        this.capacity = capacity;
        this.sizeAfterPurge = sizeAfterPurge;
        this.syncRoot = new object();
    }

    /// <summary>
    /// Adds an item if it doesn't already exist.
    /// </summary>
    public void Add(TKey key, TValue value)
    {
        lock (this.syncRoot)
        {
            if (mruCache.ContainsKey(key))
                return;

            if (mruCache.Count + 1 >= this.capacity)
            {
                while (mruCache.Count > this.sizeAfterPurge)
                {
                    var lru = mruList.Last.Value;
                    mruCache.Remove(lru);
                    mruList.RemoveLast();
                }
            }
            mruCache.Add(key, value);
            mruList.AddFirst(key);
        }
    }

    /// <summary>
    /// Removes an item if it exists.
    /// </summary>
    public void Remove(TKey key)
    {
        lock (this.syncRoot)
        {
            if (!mruCache.ContainsKey(key))
                return;

            mruCache.Remove(key);
            mruList.Remove(key);
        }
    }

    /// <summary>
    /// Gets an item.  If a matching item doesn't exist null is returned.
    /// </summary>
    public TValue Get(TKey key)
    {
        lock (this.syncRoot)
        {
            if (!mruCache.ContainsKey(key))
                return default(TValue);

            mruList.Remove(key);
            mruList.AddFirst(key);
            return mruCache[key];
        }
    }

    /// <summary>
    /// Gets whether a key is contained in the cache.
    /// </summary>
    public bool ContainsKey(TKey key)
    {
        lock (this.syncRoot)
            return mruCache.ContainsKey(key);
    }
}

ICache.cs:

/// <summary>
/// A cache.
/// </summary>
/// <typeparam name="TKey">The key type.</typeparam>
/// <typeparam name="TValue">The value type.</typeparam>
public interface ICache<TKey, TValue>
{
    void Add(TKey key, TValue value);
    void Remove(TKey key);
    TValue Get(TKey key);
}