WingMan – Rev 32

Subversion Repositories:
Rev:
using System;
using System.Threading;
using System.Threading.Tasks;
using Open.Nat;

namespace WingMan.Discovery
{
    public class Discovery : IDisposable
    {
        public delegate void PortMapFailed(object sender, DiscoveryFailedEventArgs args);

        public Discovery()
        {
            NatDiscoverer = new NatDiscoverer();
        }

        public Discovery(CancellationTokenSource cancellationTokenSource, TaskScheduler taskScheduler) : this()
        {
            CancellationTokenSource = cancellationTokenSource;
            TaskScheduler = taskScheduler;
        }

        private static CancellationTokenSource CancellationTokenSource { get; set; }
        private static NatDiscoverer NatDiscoverer { get; set; }
        private static TaskScheduler TaskScheduler { get; set; }

        private static CancellationTokenSource UPnPCancellationTokenSource { get; set; }
        private static CancellationTokenSource PMPCancellationTokenSource { get; set; }

        public void Dispose()
        {
            PMPCancellationTokenSource?.Dispose();
            PMPCancellationTokenSource = null;

            PMPCancellationTokenSource?.Dispose();
            PMPCancellationTokenSource = null;
        }

        public event PortMapFailed OnPortMapFailed;

        public async Task<bool> CreateMapping(DiscoveryType type, int port)
        {
            switch (type)
            {
                case DiscoveryType.UPnP:
                    return await CreateUPnPMapping(port);
                case DiscoveryType.PMP:
                    return await CreatePMPPortMapping(port);
                default:
                    throw new ArgumentException("Unknown disocvery type");
            }
        }

        private async Task<bool> CreateUPnPMapping(int port)
        {
            try
            {
                UPnPCancellationTokenSource = new CancellationTokenSource();
                var linkedCancellationTokenSource =
                    CancellationTokenSource.CreateLinkedTokenSource(UPnPCancellationTokenSource.Token,
                        CancellationTokenSource.Token);
                var device = await NatDiscoverer.DiscoverDeviceAsync(PortMapper.Upnp, linkedCancellationTokenSource);
                await device.CreatePortMapAsync(new Mapping(Protocol.Tcp, port, port, "WingMan Host"));

                return true;
            }
            catch (Exception ex)
            {
                await Task.Delay(0, CancellationTokenSource.Token).ContinueWith(_ =>
                        OnPortMapFailed?.Invoke(this, new DiscoveryFailedEventArgs(DiscoveryType.UPnP, ex)),
                    CancellationTokenSource.Token, TaskContinuationOptions.None, TaskScheduler);

                return false;
            }
        }

        private async Task<bool> CreatePMPPortMapping(int port)
        {
            try
            {
                PMPCancellationTokenSource = new CancellationTokenSource();
                var linkedCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(
                    PMPCancellationTokenSource.Token,
                    CancellationTokenSource.Token);
                var device = await NatDiscoverer.DiscoverDeviceAsync(PortMapper.Pmp, linkedCancellationTokenSource);
                await device.CreatePortMapAsync(new Mapping(Protocol.Tcp, port, port, "WingMan Host"));
                return true;
            }
            catch (Exception ex)
            {
                await Task.Delay(0, CancellationTokenSource.Token).ContinueWith(_ =>
                        OnPortMapFailed?.Invoke(this, new DiscoveryFailedEventArgs(DiscoveryType.PMP, ex)),
                    CancellationTokenSource.Token, TaskContinuationOptions.None, TaskScheduler);

                return false;
            }
        }
    }
}