// DWBS ------------------------------------------------------------------ DWBS
//
// PLATFORM		: WIN32 PS2 UNIX OSX PSP-INFRA PS3
// PRODUCT		: BIT_DEMON
// VISIBILITY	: PUBLIC
//
// DWBS ------------------------------------------------------------------ DWBS

#include <bdSocket/bdNAT/bdNATTypeDiscoveryClient.h>


/// We use a simplified version of the RFC 3489 NAT type discovery.
///
///                     +--------+
///                     |  Test  |
///                     |   I    |
///                     +--------+
///                          |
///                          |
///                          V
///                         /\              /\      .
///                      N /  \ Y          /  \     .
///       UDP     <-------/Resp\--------->/SKIP\    .
///       Blocked         \ ?  /          \TEST/
///                        \  /            \  /               
///                         \/              \/                
///                                          |              
///                                          |                
///                                          V                
///                                      +--------+ 
///                                      |  Test  | 
///                                      |   II   | 
///                                      +--------+ 
///                                          |                
///                                          V                
///               /\                         /\      .         
/// MODERATE  N  /  \       +--------+   N  /  \     .        
///   NAT  <--- / IP \<-----|  Test  |<--- /Resp\    .        
///             \Same/      |   III  |     \ ?  /             
///              \? /       +--------+      \  / 
///               \/                         \/ 
///               |                           |Y
///               |                           |
///               |                           V
///               |                           OPEN
///               |                           NAT
///               V            
///           STRICT NAT
///                            
///                            
///                            


#include <bdSocket/bdNAT/bdNAT.h>
#include <bdSocket/bdNAT/bdPacket/bdNATTypeDiscoveryPacket.h>
#include <bdSocket/bdNAT/bdPacket/bdNATTypeDiscoveryPacketReply.h>

#define BD_NTDCS_SEND_TIMEOUT	(.5f)
#define BD_NTDCS_MAX_RESENDS	(5u)

#define BD_LOG_LEVEL	"bdSocket/nat"

bdNATTypeDiscoveryClient::bdNATTypeDiscoveryClient(const bdInetAddr &clientAddr, const bdAddr &serverAddr)
: m_serverAddr1(serverAddr),
  m_NATtype(BD_NAT_STRICT),
  m_success(false),
  m_state(BD_NTDCS_NOT_RUNNING),
  m_resends(0)
{
	m_socket = new bdDefaultSocket;

	bdBool ok = m_socket && m_socket->create();
	//ok = ok && m_socket.setBlocking(false);

	bdPort port;
	bdAddr addr;
	do
	{
		port = static_cast<bdPort>(bdTrulyRandom::getInstance()->getRandomUInt());

		if(port <= 1024)
		{
			port += 1025;
		}

		addr.set(clientAddr, port);
	}
	while(ok && !m_socket->bind(addr));

	if(!ok)
	{
		BD_WARN(BD_LOG_LEVEL, "Failed to create socket");
	}
}

bdNATTypeDiscoveryClient::bdNATTypeDiscoveryClient()
: m_socket(BD_NULL),
m_NATtype(BD_NAT_STRICT),
m_success(false),
m_state(BD_NTDCS_NOT_RUNNING),
m_resends(0)
{
}

bdNATTypeDiscoveryClient::~bdNATTypeDiscoveryClient()
{
	if (m_socket)
	{
		delete m_socket;
	}
}

bdBool bdNATTypeDiscoveryClient::startDiscovery()
{
	m_state = BD_NTDCS_RUN_TEST_1;
	m_timer.start();
	
	BD_INFO(BD_LOG_LEVEL, "Running test 1");
	bdNATTypeDiscoveryPacket packet(bdNATTypeDiscoveryPacket::BD_NTDP_SAME_ADDR);
	const bdUInt bufferSize = BD_MAX_DATAGRAM_SIZE;
	bdUByte8 buffer[bufferSize];
	bdUInt len = 0;
	BD_ASSERT(packet.serialize(buffer, bufferSize, 0, len), "Failed to serialize a packet.");
	const bdBool sent = m_socket->sendTo(m_serverAddr1, buffer, len) == static_cast<bdInt>(len);

	return sent;
}

bdBool bdNATTypeDiscoveryClient::isPending() const
{
	return m_state != BD_NTDCS_FINI;
}

bdBool bdNATTypeDiscoveryClient::wasSuccess() const
{
	return m_success;
}

const bdNATType& bdNATTypeDiscoveryClient::getNATType() const
{
	return m_NATtype;
}

void bdNATTypeDiscoveryClient::pump()
{
	// try to receive first
	bdAddr addr;
	const bdUInt dataSize = BD_MAX_DATAGRAM_SIZE;
	bdUByte8 data[dataSize];
	bdBool receiving = true;

	while(receiving)
	{
		const bdInt received = m_socket->receiveFrom(addr, data, dataSize);
		bdNATTypeDiscoveryPacketReply reply;
		bdUInt tmpUInt;

		if (received > 0 && reply.deserialize(data, static_cast<bdUInt>(received), 0, tmpUInt))
		{
			if(reply.getType() == BD_NTDP_BINDING_RESPONSE)
			{
				handleResponse(addr, reply);
			}
		}
		else if(received < 0)
		{
			switch(received)
			{
				case BD_NET_WOULD_BLOCK:
				case BD_NET_BLOCKING_CALL_CANCELLED:
					receiving = false;
					break;
				case BD_NET_CONNECTION_RESET:
					break;
				case BD_NET_MSG_SIZE:
					break;
				case BD_NET_NOT_BOUND:
				default:
					receiving = false;
					BD_WARN(BD_LOG_LEVEL, "net error.");
					break;
			}
		}
	}

	// continue tests
	switch(m_state)
	{
		case BD_NTDCS_NOT_RUNNING:
			break;
		case BD_NTDCS_RUN_TEST_1:
			sendForTest1();
			break;
		case BD_NTDCS_RUN_TEST_2:
			sendForTest2();
			break;
		case BD_NTDCS_RUN_TEST_3:
			sendForTest3();
			break;
		case BD_NTDCS_FINI:
			break;
	}
}


// protected:

void bdNATTypeDiscoveryClient::sendForTest1()
{
	if(m_timer.getElapsedTimeInSeconds() > BD_NTDCS_SEND_TIMEOUT)
	{
		if(m_resends++ < BD_NTDCS_MAX_RESENDS)
		{
			m_timer.start();
			BD_INFO(BD_LOG_LEVEL, "Running test 1");

			bdNATTypeDiscoveryPacket packet(bdNATTypeDiscoveryPacket::BD_NTDP_SAME_ADDR);
			const bdUInt bufferSize = BD_MAX_DATAGRAM_SIZE;
			bdUByte8 buffer[bufferSize];
			bdUInt len = 0;
			BD_ASSERT(packet.serialize(buffer, bufferSize, 0, len), "Failed to serialize a packet.");
			m_socket->sendTo(m_serverAddr1, buffer, len);
		}
		else
		{
			BD_INFO(BD_LOG_LEVEL, "Test 1 failed. Not online.");
			m_state = BD_NTDCS_FINI;
		}
	}
}

void bdNATTypeDiscoveryClient::sendForTest2()
{
	if(m_timer.getElapsedTimeInSeconds() > BD_NTDCS_SEND_TIMEOUT)
	{
		if(m_resends++ < BD_NTDCS_MAX_RESENDS)
		{
			m_timer.start();
			BD_INFO(BD_LOG_LEVEL, "Running test 2");

			bdNATTypeDiscoveryPacket packet(bdNATTypeDiscoveryPacket::BD_NTDP_DIFF_PORT_IP);
			const bdUInt bufferSize = BD_MAX_DATAGRAM_SIZE;
			bdUByte8 buffer[bufferSize];
			bdUInt len = 0;
			BD_ASSERT(packet.serialize(buffer, bufferSize, 0, len), "Failed to serialize a packet.");

			m_socket->sendTo(m_serverAddr1, buffer, len);
		}
		else
		{
			BD_INFO(BD_LOG_LEVEL, "Test 2 failed.");
			m_resends = 0;
			m_state = BD_NTDCS_RUN_TEST_3;
		}
	}
}

void bdNATTypeDiscoveryClient::sendForTest3()
{
	if(m_timer.getElapsedTimeInSeconds() > BD_NTDCS_SEND_TIMEOUT)
	{
		if(m_resends++ < BD_NTDCS_MAX_RESENDS)
		{
			m_timer.start();
			BD_INFO(BD_LOG_LEVEL, "Running test 3");

			bdNATTypeDiscoveryPacket packet(bdNATTypeDiscoveryPacket::BD_NTDP_DIFF_IP);
			const bdUInt bufferSize = BD_MAX_DATAGRAM_SIZE;
			bdUByte8 buffer[bufferSize];
			bdUInt len = 0;
			BD_ASSERT(packet.serialize(buffer, bufferSize, 0, len), "Failed to serialize a packet.");

			m_socket->sendTo(m_serverAddr2, buffer, len);
		}
		else
		{
			BD_INFO(BD_LOG_LEVEL, "Test 3 failed. Strict NAT.");
			m_success = true;
			m_NATtype = BD_NAT_STRICT;
			m_state = BD_NTDCS_FINI;
		}
	}
}

void bdNATTypeDiscoveryClient::handleResponse(const bdAddr &addr, 
											  const bdNATTypeDiscoveryPacketReply &reply)
{
	switch(m_state)
	{
		case BD_NTDCS_NOT_RUNNING:
			break;
		case BD_NTDCS_RUN_TEST_1:
			{
				m_serverAddr2 = reply.getSecAddr();
				m_mappedAddr = reply.getMappedAddr();

				m_state = BD_NTDCS_RUN_TEST_2;
				m_resends = 0;

				BD_INFO(BD_LOG_LEVEL, "Reply for test 1. Start test 2.");

				// send for test 2:
				bdNATTypeDiscoveryPacket packet(bdNATTypeDiscoveryPacket::BD_NTDP_DIFF_PORT_IP);
				const bdUInt bufferSize = BD_MAX_DATAGRAM_SIZE;
				bdUByte8 buffer[bufferSize];
				bdUInt len = 0;
				BD_ASSERT(packet.serialize(buffer, bufferSize, 0, len), "Failed to serialize a packet.");

				m_socket->sendTo(m_serverAddr1, buffer, len);
			}
			break;
		case BD_NTDCS_RUN_TEST_2:
			if(addr.getAddress() == m_serverAddr2.getAddress() 
				&& addr.getPort() != m_serverAddr2.getPort())
			{
				// response to test 2 means we are behind an OPEN nat (or no NAT)

				BD_INFO(BD_LOG_LEVEL, "Reply for test 2. Open NAT (or no NAT).");

				m_success = true;
				m_state = BD_NTDCS_FINI;
				m_NATtype = BD_NAT_OPEN;
			}
			break;
		case BD_NTDCS_RUN_TEST_3:
			// if the address the STUN server sees me as from test 1 is the same
			// as now:
			if(m_mappedAddr == reply.getMappedAddr())
			{
				BD_INFO(BD_LOG_LEVEL, "Reply for test 3. Moderate NAT.");
				m_success = true;
				m_state = BD_NTDCS_FINI;
				m_NATtype = BD_NAT_MODERATE;
			}
			else
			{
				BD_INFO(BD_LOG_LEVEL, "Reply for test 3. Strict NAT.");
				m_success = true;
				m_state = BD_NTDCS_FINI;
				m_NATtype = BD_NAT_STRICT;
			}
			break;
		case BD_NTDCS_FINI:
			break;
	}
}

