SessionInvalidationTracker.cs

Go to the documentation of this file.

using System;
using System.Collections.Concurrent;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
 
using HotChocolate.Subscriptions;
 
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
 
using Tgstation.Server.Host.GraphQL;
using Tgstation.Server.Host.GraphQL.Subscriptions;
using Tgstation.Server.Host.Models;
using Tgstation.Server.Host.Utils;
 
namespace Tgstation.Server.Host.Security
{
    sealed class SessionInvalidationTracker : ISessionInvalidationTracker
    {
        readonly ITopicEventSender eventSender;
 
        readonly IAsyncDelayer asyncDelayer;
 
        readonly IHostApplicationLifetime applicationLifetime;
 
        readonly ILogger<SessionInvalidationTracker> logger;
 
        readonly ConcurrentDictionary<(string SessionId, long UserId), TaskCompletionSource<SessionInvalidationReason>> trackedSessions;
 
        public SessionInvalidationTracker(
            ITopicEventSender eventSender,
            IAsyncDelayer asyncDelayer,
            IHostApplicationLifetime applicationLifetime,
            ILogger<SessionInvalidationTracker> logger)
        {
            this.eventSender = eventSender ?? throw new ArgumentNullException(nameof(eventSender));
            this.asyncDelayer = asyncDelayer ?? throw new ArgumentNullException(nameof(asyncDelayer));
            this.applicationLifetime = applicationLifetime ?? throw new ArgumentNullException(nameof(applicationLifetime));
            this.logger = logger ?? throw new ArgumentNullException(nameof(logger));
 
            trackedSessions = new ConcurrentDictionary<(string, long), TaskCompletionSource<SessionInvalidationReason>>();
        }
 
        public void TrackSession(IAuthenticationContext authenticationContext)
        {
            trackedSessions.GetOrAdd(
                (authenticationContext.SessionId, authenticationContext.User.Require(x => x.Id)),
                tuple =>
                {
                    var (localSessionId, localUserId) = tuple;
                    logger.LogTrace("Tracking session ID for user {userId}: {sessionId}", localUserId, localSessionId);
                    var tcs = new TaskCompletionSource<SessionInvalidationReason>();
                    async void SendInvalidationTopic()
                    {
                        try
                        {
                            SessionInvalidationReason invalidationReason;
                            try
                            {
                                var otherCancellationReason = tcs.Task;
                                var timeTillSessionExpiry = authenticationContext.SessionExpiry - DateTimeOffset.UtcNow;
                                if (timeTillSessionExpiry > TimeSpan.Zero)
                                {
                                    var delayTask = asyncDelayer.Delay(timeTillSessionExpiry, applicationLifetime.ApplicationStopping).AsTask();
 
                                    await Task.WhenAny(delayTask, otherCancellationReason);
 
                                    if (delayTask.IsCompleted)
                                        await delayTask;
                                }
 
                                invalidationReason = otherCancellationReason.IsCompleted
                                    ? await otherCancellationReason
                                    : SessionInvalidationReason.TokenExpired;
 
                                logger.LogTrace("Invalidating session ID {sessionID}: {reason}", localSessionId, invalidationReason);
                            }
                            catch (OperationCanceledException ex)
                            {
                                logger.LogTrace(ex, "Invalidating session ID {sessionID} due to server shutdown", localSessionId);
                                invalidationReason = SessionInvalidationReason.ServerShutdown;
                            }
 
                            var topicName = Subscription.SessionInvalidatedTopic(authenticationContext);
                            await eventSender.SendAsync(topicName, invalidationReason, CancellationToken.None); // DCT: Session close messages should always be sent
                            await eventSender.CompleteAsync(topicName);
                        }
                        catch (Exception ex)
                        {
                            logger.LogError(ex, "Error tracking session {sessionId}!", localSessionId);
                        }
                    }
 
                    SendInvalidationTopic();
                    return tcs;
                });
        }
 
        public void UserModifiedInvalidateSessions(User user)
        {
            ArgumentNullException.ThrowIfNull(user);
 
            var userId = user.Require(x => x.Id);
            user.LastPasswordUpdate = DateTimeOffset.UtcNow;
 
            foreach (var key in trackedSessions
                .Keys
                .Where(key => key.UserId == userId)
                .ToList())
                if (trackedSessions.TryRemove(key, out var tcs))
                    tcs.TrySetResult(SessionInvalidationReason.UserUpdated);
        }
    }
}