Hermes/Hermes/CircuitManagers/OneShotRoundRobinCircuitManager.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Hermes.Objects;
using TorControlLibrary.Responses;

namespace Hermes.CircuitManagers
{
    /// <summary>
    /// One of the most basic circuit managers, using randomization to select the next viable circuit
    /// </summary>
    public class OneShotRoundRobinCircuitManager : EventBasedCircuitManager, ICircuitManager
    {
        #region Event Overrides
        protected override void OnCircuitBuilt(TorControlLibrary.Responses.EventResponse response)
        {
            //Add the completely built circuit
            AddCircuit(ParseCircuitResponse(response));
        }
        protected override void OnCircuitClosed(TorControlLibrary.Responses.EventResponse response)
        {
            Objects.Circuit circ = GetCircuit(response.ID);
            if(circ != null)
                RemoveCircuit(circ);

            if (response.EventInformation.Contains("PURPOSE=USER") || response.EventInformation.Contains("PURPOSE=GENERAL") && CircuitCount < ExpectedCircuits)
                RequestNewCircuit();
        }
        protected override void OnCircuitFailed(TorControlLibrary.Responses.EventResponse response)
        {
            Objects.Circuit circ = GetCircuit(response.ID);
            if(circ != null)
                RemoveCircuit(circ);

            if (response.EventInformation.Contains("PURPOSE=USER") || response.EventInformation.Contains("PURPOSE=GENERAL") && CircuitCount < ExpectedCircuits)
                RequestNewCircuit();
        }
        protected override void OnStreamClosed(TorControlLibrary.Responses.EventResponse response)
        {
            Stream stream = Stream.ParseStream(response);
            RemoveStream(stream);
        }
        protected override void OnStreamDetached(TorControlLibrary.Responses.EventResponse response)
        {
            Stream stream = Stream.ParseStream(response);
            if (stream.CircuitID != 0)
                stream.Circuit = GetCircuit(stream.CircuitID);

            RemoveStream(stream);
            
            if (stream.Purpose != null && (stream.Purpose.Equals("USER") || stream.Purpose.Equals("GENERAL")))
                AssignStream(stream);
        }
        protected override void OnStreamFailed(TorControlLibrary.Responses.EventResponse response)
        {
            Stream stream = Stream.ParseStream(response);
            if (stream.CircuitID != 0)
                stream.Circuit = GetCircuit(stream.CircuitID);

            RemoveStream(stream);
        }
        protected override void OnStreamNewResolved(EventResponse response)
        {
            OnStreamNew(response);
            //base.OnStreamNewResolved(response);
        }
        protected override void OnStreamNew(TorControlLibrary.Responses.EventResponse response)
        {
            Stream stream = Stream.ParseStream(response);
            if (stream.CircuitID != 0)
                stream.Circuit = GetCircuit(stream.CircuitID);

            if (stream.Purpose.Equals("USER") || stream.Purpose.Equals("GENERAL"))
                AssignStream(stream);
        }
        #endregion
        protected virtual void AssignStream(Stream stream)
        {
            Circuit circ = GetCircuitForStream(stream);
            while (circ == null)
            {
                //Wait for more circuits to build
                System.Threading.Thread.Sleep(5000);
                circ = GetCircuitForStream(stream);
            }
            try
            {
                AssignStreamToCircuit(stream, circ);
            }
            catch (TorControlLibrary.Exceptions.TorUnknownCircuitException e)
            { //The circuit has become invalidated, we need to try again
                //cleanup
                RemoveCircuit(circ);
                try
                {
                    _controlConnection.CloseCircuit(circ.CircuitID);
                }
                catch (Exception) { }
                //Call this function again, with hopefully better results next time
                AssignStream(stream);
            }
            catch (TorControlLibrary.Exceptions.TorUnknownStreamException e)
            { //Stream no longer exists, clean up and exit
                RemoveStream(stream);
                try
                {
                    _controlConnection.CloseStream(stream.StreamID);
                }
                catch (Exception) { }
            }
        }
        public override void AssignStreamToCircuit(Stream stream, Circuit circuit)
        {
            //Handles internal objects
            base.AssignStreamToCircuit(stream, circuit);
            _controlConnection.AttachStream(stream.StreamID, circuit.CircuitID);
        }
        /// <summary>
        /// This method chooses the circuit to attach to the specified stream
        /// </summary>
        /// <param name="stream"></param>
        /// <returns></returns>
        protected virtual Circuit GetCircuitForStream(Stream stream)
        {
            Circuit circ;
            //Round Robin the next circuit 
            if (stream.Destination != null)
            {
                circ = Circuits.OrderBy(p => p.AttachedStreams.Count).FirstOrDefault(p => p.Routers.Last().Value.ExitPolicy.IsAcceptableDestination(stream.Destination));
                if (circ != null)
                    return circ;
            }
            //return a circuit that won't necessarily exit successfully to avoid deadlocks and cycle for new circuits
            return Circuits.OrderBy(p => p.AttachedStreams.Count).FirstOrDefault();
        }
    }
}