/*
 * Copyright 2006-2009, Haiku, Inc. All Rights Reserved.
 * Distributed under the terms of the MIT License.
 *
 * Authors:
 *		Axel Dörfler, axeld@pinc-software.de
 *		Hugo Santos, hugosantos@gmail.com
 */


#include "EndpointManager.h"

#include <new>
#include <unistd.h>

#include <KernelExport.h>

#include <NetUtilities.h>
#include <tracing.h>

#include "TCPEndpoint.h"


//#define TRACE_ENDPOINT_MANAGER
#ifdef TRACE_ENDPOINT_MANAGER
#	define TRACE(x) dprintf x
#else
#	define TRACE(x)
#endif

#if TCP_TRACING
#	define ENDPOINT_TRACING
#endif
#ifdef ENDPOINT_TRACING
namespace EndpointTracing {

class Bind : public AbstractTraceEntry {
public:
	Bind(TCPEndpoint* endpoint, ConstSocketAddress& address, bool ephemeral)
		:
		fEndpoint(endpoint),
		fEphemeral(ephemeral)
	{
		address.AsString(fAddress, sizeof(fAddress), true);
		Initialized();
	}

	Bind(TCPEndpoint* endpoint, SocketAddress& address, bool ephemeral)
		:
		fEndpoint(endpoint),
		fEphemeral(ephemeral)
	{
		address.AsString(fAddress, sizeof(fAddress), true);
		Initialized();
	}

	virtual void AddDump(TraceOutput& out)
	{
		out.Print("tcp:e:%p bind%s address %s", fEndpoint,
			fEphemeral ? " ephemeral" : "", fAddress);
	}

protected:
	TCPEndpoint*	fEndpoint;
	char			fAddress[32];
	bool			fEphemeral;
};

class Connect : public AbstractTraceEntry {
public:
	Connect(TCPEndpoint* endpoint)
		:
		fEndpoint(endpoint)
	{
		endpoint->LocalAddress().AsString(fLocal, sizeof(fLocal), true);
		endpoint->PeerAddress().AsString(fPeer, sizeof(fPeer), true);
		Initialized();
	}

	virtual void AddDump(TraceOutput& out)
	{
		out.Print("tcp:e:%p connect local %s, peer %s", fEndpoint, fLocal,
			fPeer);
	}

protected:
	TCPEndpoint*	fEndpoint;
	char			fLocal[32];
	char			fPeer[32];
};

class Unbind : public AbstractTraceEntry {
public:
	Unbind(TCPEndpoint* endpoint)
		:
		fEndpoint(endpoint)
	{
		//fStackTrace = capture_tracing_stack_trace(10, 0, false);

		endpoint->LocalAddress().AsString(fLocal, sizeof(fLocal), true);
		endpoint->PeerAddress().AsString(fPeer, sizeof(fPeer), true);
		Initialized();
	}

#if 0
	virtual void DumpStackTrace(TraceOutput& out)
	{
		out.PrintStackTrace(fStackTrace);
	}
#endif

	virtual void AddDump(TraceOutput& out)
	{
		out.Print("tcp:e:%p unbind, local %s, peer %s", fEndpoint, fLocal,
			fPeer);
	}

protected:
	TCPEndpoint*	fEndpoint;
	//tracing_stack_trace* fStackTrace;
	char			fLocal[32];
	char			fPeer[32];
};

}	// namespace EndpointTracing

#	define T(x)	new(std::nothrow) EndpointTracing::x
#else
#	define T(x)
#endif	// ENDPOINT_TRACING


static const uint16 kLastReservedPort = 1023;
static const uint16 kFirstEphemeralPort = 40000;


ConnectionHashDefinition::ConnectionHashDefinition(EndpointManager* manager)
	:
	fManager(manager)
{
}


size_t
ConnectionHashDefinition::HashKey(const KeyType& key) const
{
	return ConstSocketAddress(fManager->AddressModule(),
		key.first).HashPair(key.second);
}


size_t
ConnectionHashDefinition::Hash(TCPEndpoint* endpoint) const
{
	return endpoint->LocalAddress().HashPair(*endpoint->PeerAddress());
}


bool
ConnectionHashDefinition::Compare(const KeyType& key,
	TCPEndpoint* endpoint) const
{
	return endpoint->LocalAddress().EqualTo(key.first, true)
		&& endpoint->PeerAddress().EqualTo(key.second, true);
}


TCPEndpoint*&
ConnectionHashDefinition::GetLink(TCPEndpoint* endpoint) const
{
	return endpoint->fConnectionHashLink;
}


//	#pragma mark -


size_t
EndpointHashDefinition::HashKey(uint16 port) const
{
	return port;
}


size_t
EndpointHashDefinition::Hash(TCPEndpoint* endpoint) const
{
	return endpoint->LocalAddress().Port();
}


bool
EndpointHashDefinition::Compare(uint16 port, TCPEndpoint* endpoint) const
{
	return endpoint->LocalAddress().Port() == port;
}


bool
EndpointHashDefinition::CompareValues(TCPEndpoint* first,
	TCPEndpoint* second) const
{
	return first->LocalAddress().Port() == second->LocalAddress().Port();
}


TCPEndpoint*&
EndpointHashDefinition::GetLink(TCPEndpoint* endpoint) const
{
	return endpoint->fEndpointHashLink;
}


//	#pragma mark -


EndpointManager::EndpointManager(net_domain* domain)
	:
	fDomain(domain),
	fConnectionHash(this),
	fLastPort(kFirstEphemeralPort)
{
	rw_lock_init(&fLock, "TCP endpoint manager");
}


EndpointManager::~EndpointManager()
{
	rw_lock_destroy(&fLock);
}


status_t
EndpointManager::Init()
{
	status_t status = fConnectionHash.Init();
	if (status == B_OK)
		status = fEndpointHash.Init();

	return status;
}


//	#pragma mark - connections


/*!	Returns the endpoint matching the connection.
	You must hold the manager's lock when calling this method (either read or
	write).
*/
TCPEndpoint*
EndpointManager::_LookupConnection(const sockaddr* local, const sockaddr* peer)
{
	return fConnectionHash.Lookup(std::make_pair(local, peer));
}


status_t
EndpointManager::SetConnection(TCPEndpoint* endpoint, const sockaddr* _local,
	const sockaddr* peer, const sockaddr* interfaceLocal)
{
	TRACE(("EndpointManager::SetConnection(%p)\n", endpoint));

	WriteLocker _(fLock);

	SocketAddressStorage local(AddressModule());
	local.SetTo(_local);

	if (local.IsEmpty(false)) {
		uint16 port = local.Port();
		local.SetTo(interfaceLocal);
		local.SetPort(port);
	}

	// We want to create a connection for (local, peer), so check to make sure
	// that this pair is not already in use by an existing connection.
	if (_LookupConnection(*local, peer) != NULL)
		return EADDRINUSE;

	endpoint->LocalAddress().SetTo(*local);
	endpoint->PeerAddress().SetTo(peer);
	T(Connect(endpoint));

	// BOpenHashTable doesn't support inserting duplicate objects. Since
	// BOpenHashTable is a chained hash table where the items are required to
	// be intrusive linked list nodes, inserting the same object twice will
	// create a cycle in the linked list, which is not handled currently.
	//
	// We need to makes sure to remove any existing copy of this endpoint
	// object from the table in order to handle calling connect() on a closed
	// socket to connect to a different remote (address, port) than it was
	// originally used for.
	//
	// We use RemoveUnchecked here because we don't want the hash table to
	// resize itself after this removal when we are planning to just add
	// another.
	fConnectionHash.RemoveUnchecked(endpoint);

	fConnectionHash.Insert(endpoint);
	return B_OK;
}


status_t
EndpointManager::SetPassive(TCPEndpoint* endpoint)
{
	WriteLocker _(fLock);

	if (!endpoint->IsBound()) {
		// if the socket is unbound first bind it to ephemeral
		SocketAddressStorage local(AddressModule());
		local.SetToEmpty();

		status_t status = _BindToEphemeral(endpoint, *local);
		if (status < B_OK)
			return status;
	}

	SocketAddressStorage passive(AddressModule());
	passive.SetToEmpty();

	if (_LookupConnection(*endpoint->LocalAddress(), *passive))
		return EADDRINUSE;

	endpoint->PeerAddress().SetTo(*passive);
	fConnectionHash.Insert(endpoint);
	return B_OK;
}


TCPEndpoint*
EndpointManager::FindConnection(sockaddr* local, sockaddr* peer)
{
	ReadLocker _(fLock);

	TCPEndpoint *endpoint = _LookupConnection(local, peer);
	if (endpoint != NULL) {
		TRACE(("TCP: Received packet corresponds to explicit endpoint %p\n",
			endpoint));
		if (gSocketModule->acquire_socket(endpoint->socket))
			return endpoint;
	}

	// no explicit endpoint exists, check for wildcard endpoints

	SocketAddressStorage wildcard(AddressModule());
	wildcard.SetToEmpty();

	endpoint = _LookupConnection(local, *wildcard);
	if (endpoint != NULL) {
		TRACE(("TCP: Received packet corresponds to wildcard endpoint %p\n",
			endpoint));
		if (gSocketModule->acquire_socket(endpoint->socket))
			return endpoint;
	}

	SocketAddressStorage localWildcard(AddressModule());
	localWildcard.SetToEmpty();
	localWildcard.SetPort(AddressModule()->get_port(local));

	endpoint = _LookupConnection(*localWildcard, *wildcard);
	if (endpoint != NULL) {
		TRACE(("TCP: Received packet corresponds to local wildcard endpoint "
			"%p\n", endpoint));
		if (gSocketModule->acquire_socket(endpoint->socket))
			return endpoint;
	}

	// no matching endpoint exists
	TRACE(("TCP: no matching endpoint!\n"));

	return NULL;
}


//	#pragma mark - endpoints


status_t
EndpointManager::Bind(TCPEndpoint* endpoint, const sockaddr* address)
{
	// check the family
	if (!AddressModule()->is_same_family(address))
		return EAFNOSUPPORT;

	WriteLocker locker(fLock);

	if (AddressModule()->get_port(address) == 0)
		return _BindToEphemeral(endpoint, address);

	return _BindToAddress(locker, endpoint, address);
}


status_t
EndpointManager::BindChild(TCPEndpoint* endpoint, const sockaddr* address)
{
	WriteLocker _(fLock);
	return _Bind(endpoint, address);
}


/*! You must have fLock write locked when calling this method. */
status_t
EndpointManager::_BindToAddress(WriteLocker& locker, TCPEndpoint* endpoint,
	const sockaddr* _address)
{
	ConstSocketAddress address(AddressModule(), _address);
	uint16 port = address.Port();

	TRACE(("EndpointManager::BindToAddress(%p)\n", endpoint));
	T(Bind(endpoint, address, false));

	// TODO: this check follows very typical UNIX semantics
	// and generally should be improved.
	if (ntohs(port) <= kLastReservedPort && geteuid() != 0)
		return B_PERMISSION_DENIED;

	bool retrying = false;
	int32 retry = 0;
	do {
		EndpointTable::ValueIterator portUsers = fEndpointHash.Lookup(port);
		retry = false;

		while (portUsers.HasNext()) {
			TCPEndpoint* user = portUsers.Next();

			if (user->LocalAddress().IsEmpty(false)
				|| address.EqualTo(*user->LocalAddress(), false)) {
				// Check if this belongs to a local connection

				// Note, while we hold our lock, the endpoint cannot go away,
				// it can only change its state - IsLocal() is safe to be used
				// without having the endpoint locked.
				tcp_state userState = user->State();
				if (user->IsLocal()
					&& (userState > ESTABLISHED || userState == CLOSED)) {
					// This is a closing local connection - wait until it's
					// gone away for real
					locker.Unlock();
					snooze(10000);
					locker.Lock();
						// TODO: make this better
					if (!retrying) {
						retrying = true;
						retry = 5;
					}
					break;
				}

				if ((endpoint->socket->options & SO_REUSEADDR) == 0)
					return EADDRINUSE;

				if (userState != TIME_WAIT && userState != CLOSED)
					return EADDRINUSE;
			}
		}
	} while (retry-- > 0);

	return _Bind(endpoint, *address);
}


/*! You must have fLock write locked when calling this method. */
status_t
EndpointManager::_BindToEphemeral(TCPEndpoint* endpoint,
	const sockaddr* address)
{
	TRACE(("EndpointManager::BindToEphemeral(%p)\n", endpoint));

	uint32 max = fLastPort + 65536;

	for (int32 i = 1; i < 5; i++) {
		// try to retrieve a more or less random port
		uint32 step = i == 4 ? 1 : (system_time() & 0x1f) + 1;
		uint32 counter = fLastPort + step;

		while (counter < max) {
			uint16 port = counter & 0xffff;
			if (port <= kLastReservedPort)
				port += kLastReservedPort;

			fLastPort = port;
			port = htons(port);

			if (!fEndpointHash.Lookup(port).HasNext()) {
				// found a port
				SocketAddressStorage newAddress(AddressModule());
				newAddress.SetTo(address);
				newAddress.SetPort(port);

				TRACE(("   EndpointManager::BindToEphemeral(%p) -> %s\n",
					endpoint, AddressString(Domain(), *newAddress,
					true).Data()));
				T(Bind(endpoint, newAddress, true));

				return _Bind(endpoint, *newAddress);
			}

			counter += step;
		}
	}

	// could not find a port!
	return EADDRINUSE;
}


status_t
EndpointManager::_Bind(TCPEndpoint* endpoint, const sockaddr* address)
{
	// Thus far we have checked if the Bind() is allowed

	status_t status = endpoint->next->module->bind(endpoint->next, address);
	if (status < B_OK)
		return status;

	fEndpointHash.Insert(endpoint);

	return B_OK;
}


status_t
EndpointManager::Unbind(TCPEndpoint* endpoint)
{
	TRACE(("EndpointManager::Unbind(%p)\n", endpoint));
	T(Unbind(endpoint));

	if (endpoint == NULL || !endpoint->IsBound()) {
		TRACE(("  endpoint is unbound.\n"));
		return B_BAD_VALUE;
	}

	WriteLocker _(fLock);

	if (!fEndpointHash.Remove(endpoint))
		panic("bound endpoint %p not in hash!", endpoint);

	fConnectionHash.Remove(endpoint);

	(*endpoint->LocalAddress())->sa_len = 0;

	return B_OK;
}


status_t
EndpointManager::ReplyWithReset(tcp_segment_header& segment, net_buffer* buffer)
{
	TRACE(("TCP: Sending RST...\n"));

	net_buffer* reply = gBufferModule->create(512);
	if (reply == NULL)
		return B_NO_MEMORY;

	AddressModule()->set_to(reply->source, buffer->destination);
	AddressModule()->set_to(reply->destination, buffer->source);

	tcp_segment_header outSegment(TCP_FLAG_RESET);
	outSegment.sequence = 0;
	outSegment.acknowledge = 0;
	outSegment.advertised_window = 0;
	outSegment.urgent_offset = 0;

	if ((segment.flags & TCP_FLAG_ACKNOWLEDGE) == 0) {
		outSegment.flags |= TCP_FLAG_ACKNOWLEDGE;
		outSegment.acknowledge = segment.sequence + buffer->size;
		// TODO: Confirm:
		if ((segment.flags & (TCP_FLAG_SYNCHRONIZE | TCP_FLAG_FINISH)) != 0)
			outSegment.acknowledge++;
	} else
		outSegment.sequence = segment.acknowledge;

	status_t status = add_tcp_header(AddressModule(), outSegment, reply);
	if (status == B_OK)
		status = Domain()->module->send_data(NULL, reply);

	if (status != B_OK)
		gBufferModule->free(reply);

	return status;
}


void
EndpointManager::Dump() const
{
	kprintf("-------- TCP Domain %p ---------\n", this);
	kprintf("%10s %21s %21s %8s %8s %12s\n", "address", "local", "peer",
		"recv-q", "send-q", "state");

	ConnectionTable::Iterator iterator = fConnectionHash.GetIterator();

	while (iterator.HasNext()) {
		TCPEndpoint *endpoint = iterator.Next();

		char localBuf[64], peerBuf[64];
		endpoint->LocalAddress().AsString(localBuf, sizeof(localBuf), true);
		endpoint->PeerAddress().AsString(peerBuf, sizeof(peerBuf), true);

		kprintf("%p %21s %21s %8lu %8lu %12s\n", endpoint, localBuf, peerBuf,
			endpoint->fReceiveQueue.Available(), endpoint->fSendQueue.Used(),
			name_for_state(endpoint->State()));
	}
}