Akari.Prototype/Akari.Prototype.Server/Services/TcpProviderService.cs

using Akari.Prototype.Server.Options;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using System;
using System.Buffers.Text;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net;
using System.Net.Sockets;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace Akari.Prototype.Server.Services
{
    // TODO Process multiple clients simultaneously
    public sealed class TcpProviderService : BackgroundService
    {
        public const int BufferLength = 4096;
        public const int MaxDataLength = MaxTokenLength + 1 + MaxNameLength;
        public const int MaxTokenLength = 24;
        public const int MaxNameLength = 200;
        public const int ReceiveTimeout = 5_000;

        private readonly ILogger<TcpProviderService> _logger;
        private readonly TcpProviderOptions _options;
        private readonly IHostApplicationLifetime _hostApplicationLifetime;
        private readonly IFingerprintManager _fingerprintManager;
        private TcpListener _listener;

        public TcpProviderService(ILogger<TcpProviderService> logger, IOptions<TcpProviderOptions> options,
                                    IHostApplicationLifetime hostApplicationLifetime, IFingerprintManager fingerprintManager)
        {
            _logger = logger;
            _options = options.Value;
            _hostApplicationLifetime = hostApplicationLifetime;
            _fingerprintManager = fingerprintManager;
        }

        public override Task StartAsync(CancellationToken cancellationToken)
        {
            _listener = new TcpListener(IPAddress.Parse(_options.ListeningAddress), _options.Port);
            _listener.Start();

            _logger.LogInformation($"Now listening on: {_listener.LocalEndpoint}");

            return Task.CompletedTask;
        }

        protected override async Task ExecuteAsync(CancellationToken stoppingToken)
        {
            await TcpLoop(stoppingToken);
        }

        private async Task TcpLoop(CancellationToken cancellationToken)
        {
            if (cancellationToken.IsCancellationRequested)
            {
                return;
            }

            try
            {
                while (!cancellationToken.IsCancellationRequested)
                {
                    _logger.LogDebug("Waiting for client...");
                    var task = await Task.WhenAny(_listener.AcceptTcpClientAsync(), Task.Delay(Timeout.Infinite, cancellationToken));

                    if (task is not Task<TcpClient> clientTask)
                    {
                        return;
                    }

                    using var client = clientTask.Result;

                    _logger.LogDebug($"Client connected from: {client.Client.RemoteEndPoint}");

                    try
                    {
                        await ReceiveAuth(client, cancellationToken);
                    }
                    catch (Exception e)
                    {
                        _logger.LogDebug($"Couldn't proccess client request: {e.Message}");
                    }
                }
            }
            catch (Exception e)
            {
                _logger.LogError(e.ToString());
                _logger.LogError("An error occured in TcpProvider, exiting...");

                _hostApplicationLifetime.StopApplication();

                Environment.Exit(1);
            }
        }

        private async Task ReceiveAuth(TcpClient client, CancellationToken cancellationToken)
        {
            using var stream = client.GetStream();

            Memory<byte> data = new byte[MaxDataLength];
            Memory<byte> buffer = new byte[BufferLength];

            // Receive token
            int position = 0;
            int read = 0;

            var timeoutToken = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, new CancellationTokenSource(ReceiveTimeout).Token).Token;

            _logger.LogDebug("Waiting for data...");
            while (position < MaxDataLength && (read = await stream.ReadAsync(buffer, timeoutToken)) != 0)
            {
                if (position + read > data.Length)
                {
                    // Invalid data
                    return;
                }

                buffer[..read].CopyTo(data[position..]);

                position += read;
            }

            var text = Encoding.UTF8.GetString(data.Span);
            var splitIndex = text.IndexOf('$');

            _logger.LogDebug($"Received text: {text}");

            if (cancellationToken.IsCancellationRequested)
            {
                return;
            }

            var handle = GCHandle.Alloc(text, GCHandleType.Pinned);

            _fingerprintManager.VerifyFingerprint(text[..splitIndex], Convert.FromBase64String(text[(splitIndex + 1)..]));

            handle.Free();
        }

        public override void Dispose()
        {
            _listener?.Stop();

            base.Dispose();
        }
    }
}