FileTransferService.cs

Go to the documentation of this file.

using System;
using System.Collections.Generic;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
 
using Microsoft.Extensions.Logging;
 
using Tgstation.Server.Api.Models;
using Tgstation.Server.Api.Models.Response;
using Tgstation.Server.Host.IO;
using Tgstation.Server.Host.Security;
using Tgstation.Server.Host.Utils;
 
namespace Tgstation.Server.Host.Transfer
{
    sealed class FileTransferService : IFileTransferTicketProvider, IFileTransferStreamHandler, IAsyncDisposable
    {
        const int TicketValidityMinutes = 5;
 
        readonly ICryptographySuite cryptographySuite;
 
        readonly IIOManager ioManager;
 
        readonly IAsyncDelayer asyncDelayer;
 
        readonly ILogger<FileTransferService> logger;
 
        readonly Dictionary<string, FileUploadProvider> uploadTickets;
 
        readonly Dictionary<string, FileDownloadProvider> downloadTickets;
 
        readonly CancellationTokenSource disposeCts;
 
        readonly object synchronizationLock;
 
        Task expireTask;
 
        bool disposed;
 
        public FileTransferService(
            ICryptographySuite cryptographySuite,
            IIOManager ioManager,
            IAsyncDelayer asyncDelayer,
            ILogger<FileTransferService> logger)
        {
            this.cryptographySuite = cryptographySuite ?? throw new ArgumentNullException(nameof(cryptographySuite));
            this.ioManager = ioManager ?? throw new ArgumentNullException(nameof(ioManager));
            this.asyncDelayer = asyncDelayer ?? throw new ArgumentNullException(nameof(asyncDelayer));
            this.logger = logger ?? throw new ArgumentNullException(nameof(logger));
 
            uploadTickets = new Dictionary<string, FileUploadProvider>();
            downloadTickets = new Dictionary<string, FileDownloadProvider>();
 
            disposeCts = new CancellationTokenSource();
 
            expireTask = Task.CompletedTask;
            synchronizationLock = new object();
        }
 
        public async ValueTask DisposeAsync()
        {
            Task toAwait;
            lock (synchronizationLock)
                if (!disposed)
                {
                    disposeCts.Cancel();
                    disposeCts.Dispose();
                    disposed = true;
                    toAwait = expireTask;
                    expireTask = Task.CompletedTask;
                }
                else
                    toAwait = Task.CompletedTask;
 
            await toAwait;
        }
 
        public FileTicketResponse CreateDownload(FileDownloadProvider downloadProvider)
        {
            ArgumentNullException.ThrowIfNull(downloadProvider);
            ObjectDisposedException.ThrowIf(disposed, this);
 
            logger.LogDebug("Creating download ticket for path {filePath}", downloadProvider.FilePath);
            var ticket = cryptographySuite.GetSecureString();
 
            lock (downloadTickets)
                downloadTickets.Add(ticket, downloadProvider);
 
            QueueExpiry(() =>
            {
                lock (downloadTickets)
                    if (downloadTickets.Remove(ticket))
                        logger.LogTrace("Expired download ticket {ticket}...", ticket);
            });
 
            logger.LogTrace("Created download ticket {ticket}", ticket);
 
            return new FileTicketResponse
            {
                FileTicket = ticket,
            };
        }
 
        public IFileUploadTicket CreateUpload(FileUploadStreamKind streamKind)
        {
            ObjectDisposedException.ThrowIf(disposed, this);
 
            logger.LogDebug("Creating upload ticket...");
            var ticket = cryptographySuite.GetSecureString();
            var uploadTicket = new FileUploadProvider(
                new FileTicketResponse
                {
                    FileTicket = ticket,
                },
                streamKind);
 
            lock (uploadTickets)
                uploadTickets.Add(ticket, uploadTicket);
 
            QueueExpiry(() =>
            {
                lock (uploadTickets)
                    if (uploadTickets.Remove(ticket))
                        logger.LogTrace("Expired upload ticket {ticket}...", ticket);
                    else
                        return;
 
                uploadTicket.Expire();
            });
 
            logger.LogTrace("Created upload ticket {ticket}", ticket);
 
            return uploadTicket;
        }
 
        public async ValueTask<Tuple<Stream?, ErrorMessageResponse?>> RetrieveDownloadStream(FileTicketResponse ticketResponse, CancellationToken cancellationToken)
        {
            ArgumentNullException.ThrowIfNull(ticketResponse);
            ObjectDisposedException.ThrowIf(disposed, this);
 
            var ticket = ticketResponse.FileTicket ?? throw new InvalidOperationException("ticketResponse must have FileTicket!");
            FileDownloadProvider? downloadProvider;
            lock (downloadTickets)
            {
                if (!downloadTickets.TryGetValue(ticket, out downloadProvider))
                {
                    logger.LogTrace("Download ticket {ticket} not found!", ticket);
                    return Tuple.Create<Stream?, ErrorMessageResponse?>(null, null);
                }
 
                downloadTickets.Remove(ticket);
            }
 
            var errorCode = downloadProvider.ActivationCallback();
            if (errorCode.HasValue)
            {
                logger.LogDebug("Download ticket {ticket} failed activation!", ticket);
                return Tuple.Create<Stream?, ErrorMessageResponse?>(null, new ErrorMessageResponse(errorCode.Value));
            }
 
            Stream stream;
            try
            {
                if (downloadProvider.StreamProvider != null)
                    stream = await downloadProvider.StreamProvider(cancellationToken);
                else
                    stream = ioManager.GetFileStream(downloadProvider.FilePath, downloadProvider.ShareWrite);
            }
            catch (IOException ex)
            {
                return Tuple.Create<Stream?, ErrorMessageResponse?>(
                    null,
                    new ErrorMessageResponse(ErrorCode.IOError)
                    {
                        AdditionalData = ex.ToString(),
                    });
            }
 
            try
            {
                logger.LogTrace("Ticket {ticket} downloading...", ticket);
                return Tuple.Create<Stream?, ErrorMessageResponse?>(stream, null);
            }
            catch
            {
                await stream.DisposeAsync();
                throw;
            }
        }
 
        public async ValueTask<ErrorMessageResponse?> SetUploadStream(FileTicketResponse ticketResponse, Stream stream, CancellationToken cancellationToken)
        {
            ArgumentNullException.ThrowIfNull(ticketResponse);
            ObjectDisposedException.ThrowIf(disposed, this);
 
            var ticket = ticketResponse.FileTicket ?? throw new InvalidOperationException("ticketResponse must have FileTicket!");
            FileUploadProvider? uploadProvider;
            lock (uploadTickets)
            {
                if (!uploadTickets.TryGetValue(ticket, out uploadProvider))
                {
                    logger.LogTrace("Upload ticket {ticket} not found!", ticket);
                    return new ErrorMessageResponse(ErrorCode.ResourceNotPresent);
                }
 
                uploadTickets.Remove(ticket);
            }
 
            return await uploadProvider.Completion(stream, cancellationToken);
        }
 
        void QueueExpiry(Action expireAction)
        {
            Task oldExpireTask;
            async Task ExpireAsync()
            {
                var expireAt = DateTimeOffset.UtcNow + TimeSpan.FromMinutes(TicketValidityMinutes);
                try
                {
                    await oldExpireTask.WaitAsync(disposeCts.Token);
 
                    var now = DateTimeOffset.UtcNow;
                    if (now < expireAt)
                        await asyncDelayer.Delay(expireAt - now, disposeCts.Token);
                }
                finally
                {
                    expireAction();
                }
            }
 
            lock (synchronizationLock)
            {
                oldExpireTask = expireTask;
                expireTask = ExpireAsync();
            }
        }
    }
}