ssdpsearch changes

SsdpSearch.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Net.Sockets;
using System.Text;
using System.Net;
using System.Threading;
using System.IO;
using PelcoAPI.Integration.Common;
using PelcoAPI.Integration.Threading;

namespace PelcoAPI.Integration.Net.Ssdp
{
    /// <summary>
    /// Class representing an SSDP Search
    /// </summary>
    public class SsdpSearch : IDisposable
    {

        #region Constructors

        /// <summary>
        /// Initializes a new instance of the <see cref="SsdpSearch"/> class.
        /// </summary>
        public SsdpSearch() : this(new SsdpSocket[] {})
        {
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="SsdpSearch"/> class.
        /// </summary>
        public SsdpSearch(params SsdpSocket[] servers)
        {
            if (servers.Length == 0)
            {
                servers = new SsdpSocket[] {new SsdpSocket()};
                this.OwnsServer = true;
            }

            this.Servers = new List<SsdpSocket>(servers);
            this.HostEndpoint = Protocol.DiscoveryEndpoints.IPv4;
            this.SearchType = Protocol.SsdpAll;
            this.Mx = Protocol.DefaultMx;   
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="SsdpSearch"/> class.
        /// </summary>
        /// <param name="localEndPoints">The local end points this search will use.</param>
        public SsdpSearch(params IPEndPoint[] localEndPoints)
        {
            if (localEndPoints.Length == 0)
                localEndPoints = new IPEndPoint[]{new IPEndPoint(Utilities.GetLocalIpAddresses().ToArray()[0], 0)};

            this.Servers = new List<SsdpSocket>();
            foreach (var endPoint in localEndPoints)
                this.Servers.Add(new SsdpSocket(endPoint));

            this.OwnsServer = true;
            this.HostEndpoint = Protocol.DiscoveryEndpoints.IPv4;
            this.SearchType = Protocol.SsdpAll;
            this.Mx = Protocol.DefaultMx;
        }


        /// <summary>
        /// Creates the search.
        /// </summary>
        /// <param name="requireUniqueLocation">if set to <c>true</c> [require unique location].</param>
        /// <returns></returns>
        public static SsdpSearch CreateSearch(bool requireUniqueLocation)
        {
            return CreateSearch(requireUniqueLocation, null);
        }

        /// <summary>
        /// Creates the search.
        /// </summary>
        /// <param name="requireUniqueLocation">if set to <c>true</c> [require unique location].</param>
        /// <param name="filter">additional filtering for this search.</param>
        /// <param name="localEndPoints">The local end points.  if none are specified, this method will create endpoints for all local adapters.</param>
        /// <returns></returns>
        public static SsdpSearch CreateSearch(bool requireUniqueLocation, Func<SsdpMessage, bool> filter, params IPEndPoint[] localEndPoints)
        {
            if (localEndPoints.Length == 0)
            {
                var localIps = Utilities.GetLocalIpAddresses();
                localEndPoints = localIps.Select(ip => new IPEndPoint(ip, 0)).ToArray();
            }

            var search = new SsdpSearch(localEndPoints);

            Dictionary<string, SsdpMessage> dict = new Dictionary<string, SsdpMessage>();
            search.Filter = (msg) =>
            {
                lock (dict)
                {
                    // Restrict duplicate search responses based on location or UDN/USN
                    // The reason for this is that there is potential for devices to share the same UDN
                    // However, each unique location is definitely a separate result
                    // And there's no potential for two devices to share the same location
                    string key = (requireUniqueLocation ? msg.Location : msg.USN);
                    if (dict.ContainsKey(key))
                        return false;

                    dict.Add(key, msg);

                    if (filter != null)
                        return filter(msg);

                    return true;
                }
            };

            return search;
        }
        #endregion

        #region Public Methods

        /// <summary>
        /// Finds the first result.
        /// </summary>
        /// <param name="destinations">The destinations.</param>
        /// <returns></returns>
        public SsdpMessage FindFirst(params IPEndPoint[] destinations)
        {
            return FindFirst(msg => true, destinations);
        }

        /// <summary>
        /// Finds the first message that matches the specified filter.
        /// </summary>
        /// <param name="destinations">The destinations.</param>
        /// <param name="filter">The filter.</param>
        /// <returns></returns>
        public SsdpMessage FindFirst(Func<SsdpMessage, bool> filter, params IPEndPoint[] destinations)
        {
            object syncRoot = new object();
            SsdpMessage result = null;
            EventHandler<EventArgs<SsdpMessage>> resultHandler = null;

            // Create our handler to make all the magic happen
            resultHandler = (sender, e) =>
            {
                lock (syncRoot)
                {
                    // If we already got our first result then ignore this
                    if (result != null)
                        return;

                    // This is our first matching result so set our value, remove the handler, and cancel the search
                    if (filter(e.Value))
                    {
                        result = e.Value;
                        this.ResultFound -= resultHandler;
                        this.CancelSearch();
                    }
                }
            };

            try
            {
                lock (this.SearchLock)
                {
                    // Add our handler and start the async search
                    this.ResultFound += resultHandler;
                    this.SearchAsync(destinations);
                }

                // Wait until our search is complete
                this.WaitForSearch();
            }
            finally
            {
                // Make sure we remove our handler when we're done
                this.ResultFound -= resultHandler;
            }

            return result;
        }

        /// <summary>
        /// Searches the specified destinations.
        /// </summary>
        /// <param name="destinations">The destinations.</param>
        /// <returns></returns>
        public List<SsdpMessage> Search(params IPEndPoint[] destinations)
        {
            List<SsdpMessage> results = new List<SsdpMessage>();
            EventHandler<EventArgs<SsdpMessage>> resultHandler = (sender, e) =>
            {
                lock (results)
                {
                    results.Add(e.Value);
                }
            };
            EventHandler completeHandler = (sender, e) =>
            {
                lock (results)
                {
                    Monitor.PulseAll(results);
                }
            };

            try
            {
                lock (this.SearchLock)
                {
                    // Add our handlers and start the async search
                    this.ResultFound += resultHandler;
                    this.SearchComplete += completeHandler;
                    this.SearchAsync(destinations);
                }

                // Wait until our search is complete
                lock (results)
                {
                    Monitor.Wait(results);
                }

                // Return the results
                return results;
            }
            finally
            {
                // Make sure we remove our handlers when we're done
                this.ResultFound -= resultHandler;
                this.SearchComplete -= completeHandler;
            }
        }

        /// <summary>
        /// Searches asynchronously.
        /// </summary>
        /// <param name="destinations">The destinations.</param>
        public void SearchAsync(params IPEndPoint[] destinations)
        {
            lock (this.SearchLock)
            {
                // If we're already searching then this is not allowed so throw an error
                if (this.IsSearching)
                    throw new InvalidOperationException("Search is already in progress.");

                this.IsSearching = true;
                this.Servers.ForEach(server => server.DataReceived += OnServerDataReceived);

                // TODO: Come up with a good calculation for this
                // Just double our mx value for the timeout for now
                var timeout = this.SearchTimeout != default(TimeSpan) ? this.SearchTimeout : TimeSpan.FromSeconds(this.Mx*2);
                this.CreateSearchTimeout(timeout);
            }

            // If no destinations were specified then default to the IPv4 discovery
            if (destinations == null || destinations.Length == 0)
                destinations = new IPEndPoint[] { Protocol.DiscoveryEndpoints.IPv4 };

            foreach (var server in this.Servers)
            {
                // Start the server
                server.StartListening();

                // Do we really need to join the multicast group to send out multicast messages? Seems that way...
                foreach (IPEndPoint dest in destinations.Where(ep => Utilities.IsMulticast(ep.Address)))
                {
                    // TODO: why join on all interfaces if the search is only going out of one interface?
                    //this.Server.JoinMulticastGroupAllInterfaces(dest);
                    try
                    {
                        //server.JoinMulticastGroup(dest.Address, Utilities.GetLocalIpAddressUsedFor(dest.Address));
                        server.GetSocket().SetSocketOption(SocketOptionLevel.IP, SocketOptionName.AddMembership, new MulticastOption(dest.Address, server.LocalEndpoint.Address));
                    }
                    catch (SocketException)
                    {
                    }
                    try
                    {
                        server.GetSocket().SetSocketOption(SocketOptionLevel.IP, SocketOptionName.MulticastTimeToLive, Protocol.SocketTtl);
                    }
                    catch (SocketException)
                    {
                    }
                    try
                    {
                        server.GetSocket().SetSocketOption(SocketOptionLevel.IP, SocketOptionName.MulticastInterface, BitConverter.ToInt32(server.LocalEndpoint.Address.GetAddressBytes(), 0));
                    }
                    catch (SocketException)
                    {
                    }
                    try
                    {
                        server.GetSocket().SetSocketOption(SocketOptionLevel.IP, SocketOptionName.MulticastLoopback, 1);
                    }
                    catch (SocketException)
                    {
                        // If we're already joined to this group we'll throw an error so just ignore it
                    }

                }

                // If we're sending out any searches to the broadcast address then be sure to enable broadcasts
                if (!server.EnableBroadcast)
                    server.EnableBroadcast = destinations.Any(ep => ep.Address.Equals(IPAddress.Broadcast));

                // Now send out our search data
                foreach (IPEndPoint dest in destinations)
                {
                    // Make sure we respect our option as to whether we use the destination as the host value
                    IPEndPoint host = (this.UseRemoteEndpointAsHost ? dest : (this.UseLocalEndpointAsHost ? server.LocalEndpoint : this.HostEndpoint));
                    if (host.Address.Equals(IPAddress.Any))
                        host.Address = Utilities.GetLocalIpAddressUsedFor(dest.Address);

                    string req = Protocol.CreateDiscoveryRequest(host, this.SearchType, this.Mx);
                    byte[] bytes = Encoding.ASCII.GetBytes(req);

                    // TODO: Should we make this configurable?
                    // NOTE: It's recommended to send two searches
                    server.Send(bytes, bytes.Length, dest);
                    server.Send(bytes, bytes.Length, dest);
                }
            }
        }

        /// <summary>
        /// Cancels the search.
        /// </summary>
        public void CancelSearch()
        {
            lock (this.SearchLock)
            {
                // If we're not searching then nothing to do here
                if (!this.IsSearching)
                    return;

                // If we were called from the timeout then this will be null
                if (this.CurrentSearchTimeout != null)
                {
                    this.CurrentSearchTimeout.Dispose();
                    this.CurrentSearchTimeout = null;
                }

                // Cleanup our handler and notify everyone that we're done
                this.Servers.ForEach(server => server.DataReceived -= OnServerDataReceived);

                // If this is our server then go ahead and stop listening
                if (this.OwnsServer)
                    this.Servers.ForEach(server => server.StopListening());

                this.IsSearching = false;
                this.OnSearchComplete();
            }
        }

        /// <summary>
        /// Waits for the current search to complete.
        /// </summary>
        public void WaitForSearch()
        {
            // Create an object for signaling and an event handler to signal it
            object signal = new object();
            EventHandler handler = (sender, args)  =>
            {
                lock (signal)
                {
                    Monitor.Pulse(signal);
                }
            };

            try
            {
                lock (signal)
                {
                    lock (this.SearchLock)
                    {
                        // If we're not searching then nothing to do here
                        if (!this.IsSearching)
                            return;

                        // Attach our handler
                        this.SearchComplete += handler;
                    }

                    // Wait for our event handler to trigger our signal
                    Monitor.Wait(signal);
                }
            }
            finally
            {
                // Make sure to remove our handler
                this.SearchComplete -= handler;
            }
        }

        #endregion

        #region Protected Methods

        /// <summary>
        /// Creates the search timeout.
        /// </summary>
        /// <param name="timeout">The timeout.</param>
        protected void CreateSearchTimeout(TimeSpan timeout)
        {
            lock (this.SearchLock)
            {
                this.CurrentSearchTimeout = Dispatcher.Add(() =>
                {
                    lock (this.SearchLock)
                    {
                        // Search may have already been cancelled
                        if (!this.IsSearching)
                            return;

                        // Make sure we remove ourself before calling CancelSearch
                        this.CurrentSearchTimeout = null;

                        // Cancel search will clean up everything
                        // Seems kind of wrong but it does exactly what we want
                        // If we add a special cancel event or something then we'll want to change this
                        this.CancelSearch();
                    }
                }, timeout);
            }
        }

        #endregion

        #region Events

        /// <summary>
        /// Called when [server data received].
        /// </summary>
        /// <param name="sender">The sender.</param>
        /// <param name="e">The instance containing the event data.</param>
        protected virtual void OnServerDataReceived(object sender, EventArgs<NetworkData> e)
        {
            // Queue this response to be processed
            ThreadPool.QueueUserWorkItem(data =>
            {
                try
                {
                    // Parse our message and fire our event
                    using (var stream = new MemoryStream(e.Value.Buffer, 0, e.Value.Length))
                    {
                        this.OnResultFound(new SsdpMessage(HttpMessage.Parse(stream), e.Value.RemoteIPEndpoint));
                    }
                }
                catch (ArgumentException ex)
                {
                    System.Diagnostics.Trace.TraceError("Failed to parse SSDP response: {0}", ex.ToString());
                }
            });
        }

        /// <summary>
        /// Occurs when [search complete].
        /// </summary>
        public event EventHandler SearchComplete;

        /// <summary>
        /// Called when [search complete].
        /// </summary>
        protected virtual void OnSearchComplete()
        {
            var handler = this.SearchComplete;
            if (handler != null)
                handler(this, EventArgs.Empty);
        }

        /// <summary>
        /// Occurs when [result found].
        /// </summary>
        public event EventHandler<EventArgs<SsdpMessage>> ResultFound;

        /// <summary>
        /// Called when [result found].
        /// </summary>
        /// <param name="result">The result.</param>
        protected virtual void OnResultFound(SsdpMessage result)
        {
            // This is a search so ignore any advertisements we get
            if (!this.AllowAdvertisements)
                if (result.IsAdvertisement)
                    return;

            // If this is not a notify message then ignore it
            if (result.Message.IsRequest)
                return;

            // Check to make sure this message matches our filter
            var filter = this.Filter;
            if (!filter(result))
                return;
            
            var handler = this.ResultFound;
            if (handler != null)
                handler(this, new EventArgs<SsdpMessage>(result));
        }

        #endregion

        #region Properties

        /// <summary>
        /// 
        /// </summary>
        protected static readonly TimeoutDispatcher Dispatcher = new TimeoutDispatcher();
        /// <summary>
        /// 
        /// </summary>
        protected readonly object SearchLock = new object();

        /// <summary>
        /// Gets or sets the current search timeout object.
        /// </summary>
        /// <value>
        /// The current search timeout.
        /// </value>
        protected IDisposable CurrentSearchTimeout
        {
            get;
            set;
        }

        /// <summary>
        /// Gets or sets a value indicating whether [allow advertisements].
        /// </summary>
        /// <value>
        ///   <c>true</c> if [allow advertisements]; otherwise, <c>false</c>.
        /// </value>
        public bool AllowAdvertisements { get; set; }

        /// <summary>
        /// Gets or sets the search timeout duration.
        /// </summary>
        /// <value>
        /// The search timeout.
        /// </value>
        public TimeSpan SearchTimeout { get; set; }

        /// <summary>
        /// Gets or sets a value indicating whether [owns server].
        /// </summary>
        /// <value>
        ///   <c>true</c> if [owns server]; otherwise, <c>false</c>.
        /// </value>
        protected bool OwnsServer
        {
            get;
            set;
        }

        /// <summary>
        /// Gets or sets the server.
        /// </summary>
        /// <value>
        /// The server.
        /// </value>
        protected List<SsdpSocket> Servers
        {
            get;
            set;
        }

        /// <summary>
        /// Gets or sets a value indicating whether this instance is searching.
        /// </summary>
        /// <value>
        /// 	<c>true</c> if this instance is searching; otherwise, <c>false</c>.
        /// </value>
        public bool IsSearching
        {
            get;
            protected set;
        }

        /// <summary>
        /// Gets or sets the filter.
        /// </summary>
        /// <value>
        /// The filter.
        /// </value>
        public Func<SsdpMessage, bool> Filter
        {
            get;
            set;
        }

        private bool _useRemoteEndpointAsHost;
        /// <summary>
        /// Gets or sets a value indicating whether [use remote endpoint as host].
        /// </summary>
        /// <value>
        /// 	<c>true</c> if [use remote endpoint as host]; otherwise, <c>false</c>.
        /// </value>
        public bool UseRemoteEndpointAsHost
        {
            get { return _useRemoteEndpointAsHost; }
            set
            {
                if (value)
                    _useLocalEndpointAsHost = false;

                _useRemoteEndpointAsHost = value;
            }
        }

        private bool _useLocalEndpointAsHost;
        /// <summary>
        /// Gets or sets a value indicating whether [use local endpoint as host].
        /// </summary>
        /// <value>
        /// 	<c>true</c> if [use local endpoint as host]; otherwise, <c>false</c>.
        /// </value>
        public bool UseLocalEndpointAsHost
        {
            get { return _useLocalEndpointAsHost; }
            set
            {
                if (value)
                    _useRemoteEndpointAsHost = false;

                _useLocalEndpointAsHost = value;
            }
        }


        /// <summary>
        /// Gets or sets the host endpoint.
        /// </summary>
        /// <value>
        /// The host endpoint.
        /// </value>
        public IPEndPoint HostEndpoint
        {
            get;
            set;
        }

        /// <summary>
        /// Gets or sets the type of the search.
        /// </summary>
        /// <value>
        /// The type of the search.
        /// </value>
        public string SearchType
        {
            get;
            set;
        }

        /// <summary>
        /// Gets or sets the mx.
        /// </summary>
        /// <value>
        /// The mx.
        /// </value>
        public ushort Mx
        {
            get;
            set;
        }

        #endregion

        #region IDisposable Interface

        /// <summary>
        /// Performs application-defined tasks associated with freeing, releasing, or resetting unmanaged resources.
        /// </summary>
        public virtual void Dispose()
        {
            // Only close the server if we created it
            if (this.OwnsServer)
                this.Servers.ForEach(server => server.Close());
        }

        #endregion

    }
}