UdpDiscoveryClient.cs

// 
// Author: John Carruthers ([email protected])
// 
// Copyright (C) 2014 John Carruthers
// 
// Permission is hereby granted, free of charge, to any person obtaining
// a copy of this software and associated documentation files (the
// "Software"), to deal in the Software without restriction, including
// without limitation the rights to use, copy, modify, merge, publish,
// distribute, sublicense, and/or sell copies of the Software, and to
// permit persons to whom the Software is furnished to do so, subject to
// the following conditions:
//  
// The above copyright notice and this permission notice shall be
// included in all copies or substantial portions of the Software.
//  
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
// LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
// WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
// 

using System;
using System.Collections.Generic;
using System.Net;
using System.Net.NetworkInformation;
using System.Net.Sockets;
using System.Threading;
using System.Threading.Tasks;
using AdjutantFramework.Diagnostics.Logging;
using AdjutantFramework.Modules.Domains;

namespace AdjutantFramework.Network.Discovery
{
    /// <summary>
    /// Periodically broadcasts to the network to discover servers.
    /// </summary>
    public class UdpBroadcastClient : DomainShareObject, IDiscoveryClient
    {
        /// <summary>
        /// Gets the port number being pinged when using Locate.
        /// </summary>
        public int BroadcastPort { get; private set; }

        /// <summary>
        /// UDP socket used to broadcast a ping packet.
        /// </summary>
        private List<BroadcastSocket> _broadcastSockets = new List<BroadcastSocket>();

        /// <summary>
        /// UDP socket used to receive a response from any servers found on the network.
        /// </summary>
        private Socket _recvSocket;

        /// <summary>
        /// UDP port the locator can receive replies on.
        /// </summary>
        private int _recvPort;

        /// <summary>
        /// Socket args for receiving.
        /// </summary>
        private SocketAsyncEventArgs _recvArgs;

        private readonly TimeSpan _discoverTime;
        private readonly IConnectionFactory _connectionFactory;

        private CancellationTokenSource _cancellationTokenSource;
        private CancellationToken _cancellationToken;
        private Task _searchTask;

        private Bundle _bundle;

        /// <summary>
        /// 
        /// </summary>
        /// <param name="discoverTime">How often to try and locate a server.</param>
        /// <param name="broadcastPort"></param>
        /// <param name="connectionFactory"></param>
        public UdpBroadcastClient(TimeSpan discoverTime, int broadcastPort, IConnectionFactory connectionFactory)
        {
            if (connectionFactory == null) throw new ArgumentNullException("connectionFactory");
            _discoverTime = discoverTime;
            _connectionFactory = connectionFactory;
            BroadcastPort = broadcastPort;

            SetupRecvSocket();
        }

        private void SendComplete(object sender, SocketAsyncEventArgs socketAsyncEventArgs)
        {

        }

        private void SetupBroadcastSockets()
        {
            var broadcastData = UdpBroadcastPackets.Ping;
            var port = BitConverter.GetBytes(_recvPort);
            Buffer.BlockCopy(port, 0, broadcastData, 4, 4);

            foreach (var nic in NetworkInterface.GetAllNetworkInterfaces())
            {
                var addresses = nic.GetIPProperties().UnicastAddresses;
                if (addresses.Count == 0) continue;
                _bundle.Log.Write(LogLevel.Debug, "Locator binding to {0}", nic.Name);
                foreach (var addrInfo in addresses)
                {
                    try
                    {
                        var socket = new Socket(addrInfo.Address.AddressFamily, SocketType.Dgram, ProtocolType.Udp);
                        socket.EnableBroadcast = true;
                        socket.Bind(new IPEndPoint(addrInfo.Address, 0));

                        var args = new SocketAsyncEventArgs();
                        args.Completed += SendComplete;
                        args.SetBuffer(broadcastData, 0, broadcastData.Length);
                        args.RemoteEndPoint = new IPEndPoint(IPAddress.Broadcast, BroadcastPort);

                        _broadcastSockets.Add(new BroadcastSocket
                        {
                            Socket = socket,
                            AsyncArgs = args
                        });
                        _bundle.Log.Write(LogLevel.Debug, "Bound address {0}", addrInfo.Address);
                    }
                    catch (Exception)
                    {
                    }
                }
            }
        }

        /// <summary>
        /// Setup the receiving socket.
        /// </summary>
        private void SetupRecvSocket()
        {
            _recvSocket = new Socket(AddressFamily.InterNetwork, SocketType.Dgram, ProtocolType.Udp);
            var rand = new Random();
            while (true)
            {
                try
                {
                    var port = rand.Next(1000, 65535);
                    _recvSocket.Bind(new IPEndPoint(IPAddress.Any, port));
                    _recvPort = port;
                    break;
                }
                catch
                {
                }
            }
            _recvArgs = new SocketAsyncEventArgs();
            _recvArgs.Completed += ReceivedResponse;
            _recvArgs.SetBuffer(new byte[UdpBroadcastPackets.Response.Length], 0, UdpBroadcastPackets.Response.Length);
            Receive();
        }

        /// <summary>
        /// Listen for a response on the receive socket.
        /// </summary>
        private void Receive()
        {
            if (_recvArgs == null || _recvSocket == null) return;
            _recvArgs.RemoteEndPoint = new IPEndPoint(IPAddress.Any, 0);
            if (!_recvSocket.ReceiveFromAsync(_recvArgs))
                ReceivedResponse(_recvSocket, _recvArgs);
        }

        private void ReceivedResponse(object sender, SocketAsyncEventArgs socketAsyncEventArgs)
        {
            if (socketAsyncEventArgs.SocketError != SocketError.Success ||
                socketAsyncEventArgs.BytesTransferred != UdpBroadcastPackets.Response.Length) return;
            var receivedFrom = (IPEndPoint)socketAsyncEventArgs.RemoteEndPoint;
            var headerOk = true;
            for (var i = 0; i < 4; i++)
            {
                if (socketAsyncEventArgs.Buffer[i] == UdpBroadcastPackets.Response[i]) continue;
                headerOk = false;
                break;
            }
            if (headerOk)
            {
                var port = BitConverter.ToInt32(socketAsyncEventArgs.Buffer, 4);
                var networkIdBytes = new byte[16];
                Buffer.BlockCopy(socketAsyncEventArgs.Buffer, 8, networkIdBytes, 0, 16);
                var networkId = new Guid(networkIdBytes);

                if (_cancellationTokenSource == null || _cancellationToken.IsCancellationRequested)
                    return;

                StopAsync();

                var connection = _connectionFactory.CreateConnection(networkId, Guid.Empty, new IPEndPoint(receivedFrom.Address, port));

                if (connection == null)
                {
                    StartAsync(_bundle);
                }
                else
                {
                    connection.Disconnected += (o, args) => StartAsync(_bundle);
                    _bundle.Log.Write(LogLevel.Debug, "Connected to {0} via {1}", connection.RemoteAddress, connection.LocalAddress);
                }
            }
            Receive();
        }

        /// <summary>
        /// Attempts to locate a device on the network.
        /// One Locate pulse can cause multiple DeviceLocated events to fire.
        /// </summary>
        private void Locate()
        {
            if (_isDisposed)
                throw new ObjectDisposedException("BroadcastLocator");
            try
            {
                foreach (var socket in _broadcastSockets)
                {
                    _bundle.Log.Write(LogLevel.Debug, "Searching for nodes on {0}", socket.Socket.LocalEndPoint);
                    socket.Socket.SendToAsync(socket.AsyncArgs);
                }
            }
            catch
            {
            }
        }

        private bool _isDisposed;
        public void Dispose()
        {
            if (_isDisposed)
                throw new ObjectDisposedException("BroadcastLocator");
            _isDisposed = true;
            foreach (var socket in _broadcastSockets)
            {
                socket.Socket.Dispose();
                socket.AsyncArgs.Dispose();
            }
            _broadcastSockets.Clear();
            _broadcastSockets = null;
            _recvSocket.Close();
            _recvSocket.Dispose();
            _recvSocket = null;
            _recvArgs.Dispose();
            _recvArgs = null;
        }

        public void StartAsync(Bundle bundle)
        {
            if (bundle == null) throw new ArgumentNullException("bundle");
            _bundle = bundle;

            if (_cancellationTokenSource != null)
            {
                if (!_cancellationToken.IsCancellationRequested)
                    //  todo: replace with proper Exception type
                    throw new Exception("Already started.");
                else
                    _cancellationTokenSource.Dispose();
            }

            if (_broadcastSockets.Count == 0)
                SetupBroadcastSockets();

            _cancellationTokenSource = new CancellationTokenSource();
            _cancellationToken = _cancellationTokenSource.Token;
            _searchTask = Task.Factory.StartNew(() =>
            {
                while (!_cancellationToken.IsCancellationRequested)
                {
                    Locate();
                    Thread.Sleep(_discoverTime);
                }
            }, _cancellationToken, TaskCreationOptions.LongRunning, TaskScheduler.Current);
        }

        public void StopAsync()
        {
            _cancellationTokenSource.Cancel();
        }
    }

    /// <summary>
    /// Holds socket references in one logical place.
    /// </summary>
    internal class BroadcastSocket
    {
        public Socket Socket { get; set; }
        public SocketAsyncEventArgs AsyncArgs { get; set; }
    }
}