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.Extensions;
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;
 
        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 (expireTask != null)
                {
                    disposeCts.Cancel();
                    disposeCts.Dispose();
                    toAwait = expireTask;
                    expireTask = null;
                }
                else
                    toAwait = Task.CompletedTask;
 
            await toAwait;
        }
 
        public FileTicketResponse CreateDownload(FileDownloadProvider downloadProvider)
        {
            ArgumentNullException.ThrowIfNull(downloadProvider);
 
            logger.LogDebug("Creating download ticket for path {filePath}", downloadProvider.FilePath);
            var ticketResult = CreateTicket();
 
            lock (downloadTickets)
                downloadTickets.Add(ticketResult.FileTicket, downloadProvider);
 
            QueueExpiry(() =>
            {
                lock (downloadTickets)
                    if (downloadTickets.Remove(ticketResult.FileTicket))
                        logger.LogTrace("Expired download ticket {ticket}...", ticketResult.FileTicket);
            });
 
            logger.LogTrace("Created download ticket {ticket}", ticketResult.FileTicket);
 
            return ticketResult;
        }
 
        public IFileUploadTicket CreateUpload(FileUploadStreamKind streamKind)
        {
            logger.LogDebug("Creating upload ticket...");
            var uploadTicket = new FileUploadProvider(CreateTicket(), streamKind);
 
            lock (uploadTickets)
                uploadTickets.Add(uploadTicket.Ticket.FileTicket, uploadTicket);
 
            QueueExpiry(() =>
            {
                lock (uploadTickets)
                    if (uploadTickets.Remove(uploadTicket.Ticket.FileTicket))
                        logger.LogTrace("Expired upload ticket {ticket}...", uploadTicket.Ticket.FileTicket);
                    else
                        return;
 
                uploadTicket.Expire();
            });
 
            logger.LogTrace("Created upload ticket {ticket}", uploadTicket.Ticket.FileTicket);
 
            return uploadTicket;
        }
 
        public async Task<Tuple<Stream, ErrorMessageResponse>> RetrieveDownloadStream(FileTicketResponse ticket, CancellationToken cancellationToken)
        {
            ArgumentNullException.ThrowIfNull(ticket);
 
            FileDownloadProvider downloadProvider;
            lock (downloadTickets)
            {
                if (!downloadTickets.TryGetValue(ticket.FileTicket, out downloadProvider))
                {
                    logger.LogTrace("Download ticket {ticket} not found!", ticket.FileTicket);
                    return Tuple.Create<Stream, ErrorMessageResponse>(null, null);
                }
 
                downloadTickets.Remove(ticket.FileTicket);
            }
 
            var errorCode = downloadProvider.ActivationCallback();
            if (errorCode.HasValue)
            {
                logger.LogDebug("Download ticket {ticket} failed activation!", ticket.FileTicket);
                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.FileTicket);
                return Tuple.Create<Stream, ErrorMessageResponse>(stream, null);
            }
            catch
            {
                await stream.DisposeAsync();
                throw;
            }
        }
 
        public async Task<ErrorMessageResponse> SetUploadStream(FileTicketResponse ticket, Stream stream, CancellationToken cancellationToken)
        {
            ArgumentNullException.ThrowIfNull(ticket);
 
            FileUploadProvider uploadProvider;
            lock (uploadTickets)
            {
                if (!uploadTickets.TryGetValue(ticket.FileTicket, out uploadProvider))
                {
                    logger.LogTrace("Upload ticket {ticket} not found!", ticket.FileTicket);
                    return new ErrorMessageResponse(ErrorCode.ResourceNotPresent);
                }
 
                uploadTickets.Remove(ticket.FileTicket);
            }
 
            return await uploadProvider.Completion(stream, cancellationToken);
        }
 
        FileTicketResponse CreateTicket() => new ()
        {
            FileTicket = cryptographySuite.GetSecureString(),
        };
 
        void QueueExpiry(Action expireAction)
        {
            Task oldExpireTask;
            async Task ExpireAsync()
            {
                var expireAt = DateTimeOffset.UtcNow + TimeSpan.FromMinutes(TicketValidityMinutes);
                try
                {
                    await oldExpireTask.WithToken(disposeCts.Token);
 
                    var now = DateTimeOffset.UtcNow;
                    if (now < expireAt)
                        await asyncDelayer.Delay(expireAt - now, disposeCts.Token);
                }
                finally
                {
                    expireAction();
                }
            }
 
            lock (synchronizationLock)
            {
                oldExpireTask = expireTask;
                expireTask = ExpireAsync();
            }
        }
    }
}