JD2022-TU1/main/tools/framework/JD.Remoting/Connections/TCPClientConnection.cs

using System;
using System.Net;
using System.Net.Sockets;
using JD.Interfaces.Remoting;
using JD.Remoting.Details;

namespace JD.Remoting.Connections
{
    /// <summary>
    /// Represents the implementation of a TCP connection between a tool and the UAF engine. 
    /// </summary>
    public class TcpClientConnection : TcpConnection
    {
        private DateTime m_lastConnectionAttempt = new DateTime(0);

        public const int TimeBetweenTwoConnectionsAttempt = 1000;
        public const int TimeBetweenTwoPings = 2000;
        public const int TimeToWaitAfterDisconnection = 200;

        /// <summary>
        /// Default constructor. 
        /// </summary>
        /// <param name="plugin">The plugin you want to connect to in the engine. </param>
        public TcpClientConnection(IPlugin plugin)
            : base(plugin)
        {
        }

        /// <summary>
        /// Opens a TCP connection to a distant target. This operation is asynchronous. 
        /// <remarks>Use the Connection.IsConnected property to determine if the connection is established or not. </remarks>
        /// </summary>
        /// <returns>Always zero. </returns>
        public override int Open(IEndPointId endPointId = null)
        {
            var _endPointId = endPointId ?? EndPointId.Default;
            if (IsConnected)
                throw new InvalidOperationException("The connection is already opened. Use Close() function to terminate it before calling Open() again. ");

            m_endPoint = new IPEndPoint(_endPointId.IpAddress, _endPointId.Port);
            m_endPointId = _endPointId;
            return base.Open(_endPointId);
        }

        /// <summary>
        /// Closes the connection. 
        /// </summary>
        /// <returns>Always zero. </returns>
        public override int Close()
        {
            base.Close();
            m_endPoint = null;
            return 0;
        }

        protected override void Dispose(bool disposing)
        {
            if (disposing)
            {
                base.Dispose(true);
                if (IsConnected)
                    throw new InvalidOperationException("Connection is still active. ");
            }
        }

        /// <summary>
        /// Gets a value that indicates whether the connection is established. 
        /// </summary>
        /// <remarks>The connection state may change between the moment you call IsConnected and the moment you use the connection. </remarks>
        public override bool IsConnected
        {
            get
            {
                try
                {
                    return m_socket != null && m_socket.Connected;
                }
                catch
                {
                    return false;
                }
            }
        }

        /// <summary>
        /// Tries to connect the socket. 
        /// </summary>
        /// <returns>0 if operation succeeded, a non-zero value otherwise. 
        /// See Windows Sockets Error Codes page for more information on error codes. </returns>
        /// <exception cref="System.ObjectDisposedException">The System.Net.Socket object has been closed. </exception>
        /// <exception cref="System.Security.SecurityException">A caller higher in the call stack does not have permission for the requested operation. </exception>
        /// <exception cref="System.InvalidOperationException">The Socket is listening.</exception>
        protected override int Connect()
        {
            int errorCode = 0;
            try
            {
                using (new AutoMutex(m_socketMutex))
                {
                    m_socket.Connect(m_endPoint);
                }
                OnConnectedChanged();
            }
            catch (SocketException ex)
            {
                errorCode = ex.ErrorCode;
            }

            return errorCode;
        }

        /// <summary>
        /// Tries to disconnect the socket. 
        /// </summary>
        /// <returns>0 if operation succeeded, a non-zero value otherwise. 
        /// See Windows Sockets Error Codes page for more information on error codes. </returns>
        /// <exception cref="System.ObjectDisposedException">The System.Net.Socket object has been closed. </exception>
        /// <exception cref="System.PlatformNotSupportedException">This method requires Windows 2000 or earlier. </exception>
        protected override int Disconnect()
        {
            int errorCode = 0;
            try
            {
                if (IsConnected)
                {
                    // Update last connection date to prevent an automatic reconnection in case the EndPoint is changed. 
                    m_lastConnectionAttempt = DateTime.Now.Subtract(new TimeSpan(0, 0, 0, 0, TimeBetweenTwoConnectionsAttempt - TimeToWaitAfterDisconnection));

                    SendDisconnect();
                    DisposeSocket(ref m_socket, m_socketMutex);
                    OnConnectedChanged();
                }
                else
                {
                    if (m_socket != null)
                    {
                        DisposeSocket(ref m_socket, m_socketMutex);
                        OnConnectedChanged();
                    }
                }
            }
            catch (SocketException ex)
            {
                errorCode = ex.ErrorCode;
            }

            return errorCode;
        }

        /// <summary>
        /// Receives data from the socket. 
        /// </summary>
        /// <param name="bufferOffset"></param>
        /// <param name="readCount">The number of bytes read from the socket. </param>
        /// <param name="buffer"></param>
        /// <param name="bufferSize"></param>
        /// <returns>0 if operation succeeded, a non-zero value otherwise. 
        /// See Windows Sockets Error Codes page for more information on error codes. </returns>
        /// <exception cref="System.ArgumentOutOfRangeException">Offset or size values are invalid. </exception>
        /// <exception cref="System.ObjectDisposedException">The System.Net.Socket object has been closed. </exception>
        protected override int Receive(byte[] buffer, int bufferSize, int bufferOffset, out int readCount)
        {
            readCount = 0;

            if (!IsConnected)
                throw new InvalidOperationException("Socket is not connected. Call Connect() method before trying to receive a message again. ");

            //using (new AutoMutex(m_socketMutex))
            {
                if (m_socket.Available == 0)
                {
                    // Nothing to receive. 
                    return 0;
                }

                int errorCode = 0;
                try
                {
                    readCount = m_socket.ReceiveFrom(buffer, bufferOffset, bufferSize, SocketFlags.None, ref m_endPoint);
                }
                catch (SocketException ex)
                {
                    errorCode = ex.ErrorCode;
                }

                return errorCode;
            }
        }

        protected override void ConnectionTaskCallback()
        {
            using (new AutoMutex(m_socketMutex))
            {
                if (m_socket == null)
                    m_socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
            }

            if (IsConnected)
            {
                if (DateTime.Now.Subtract(m_lastConnectionAttempt).TotalMilliseconds > TimeBetweenTwoPings)
                {
                    if (SendPing() != 0)
                    {
                        // Connection was lost. 
                        Disconnect();
                    }

                    m_lastConnectionAttempt = DateTime.Now;
                }
            }
            else
            {
                if (DateTime.Now.Subtract(m_lastConnectionAttempt).TotalMilliseconds > TimeBetweenTwoConnectionsAttempt)
                {
                    if (m_socket.LocalEndPoint != null)
                    {
                        // Socket was already connected, but the connection was lost on the other side. 
                        // Dispose this socket and create a new one. 
                        DisposeSocket(ref m_socket);
                        m_socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
                    }

                    // Try to reconnect. 
                    if (m_endPoint != null)
                    {
                        Connect();
                    }

                    m_lastConnectionAttempt = DateTime.Now;
                }
            }

            // I moved the mutex lock a bit earlier, as I got some issues in the following case : 
            //   - IsConnected is true, the connection tries to receive data. 
            //   - Disconnect() occurs, followed by socket destruction. 
            //   - Receive() tries to read in a null socket. 
            using (new AutoMutex(m_socketMutex))
            {
                if (IsConnected)
                {
                    Receive();
                }
            }
        }

        protected override void OnConnectedChanged()
        {
            if (IsConnected)
            {
                SendIdentify();
            }

            base.OnConnectedChanged();
        }
    }
}