#include <stdafx.h>
#include <winsock2.h>
#include <windows.h>
#include <tchar.h>
#include <iphlpapi.h>
#include <vd2/system/binary.h>
#include <vd2/system/error.h>
#include <vd2/system/refcount.h>
#include <vd2/system/thunk.h>
#include <vd2/system/vdstl.h>
#include <vd2/system/w32assist.h>
#include <at/atnetwork/socket.h>
#include <at/atnetworksockets/internal/worker.h>

#pragma comment(lib, "ws2_32.lib")
#pragma comment(lib, "iphlpapi.lib")

ATNetSockBridgeHandler::ATNetSockBridgeHandler(ATNetSockWorker *parent, SOCKET s, IATSocket *s2, uint32 srcIpAddr, uint16 srcPort, uint32 dstIpAddr, uint16 dstPort)
	: mSrcIpAddr(srcIpAddr)
	, mSrcPort(srcPort)
	, mDstIpAddr(dstIpAddr)
	, mDstPort(dstPort)
	, mpParent(parent)
	, mNativeSocket(s)
	, mpLocalSocket(s2)
	, mbLocalClosed(false)
	, mbNativeConnected(false)
	, mbNativeClosed(false)
	, mLocalReadAvail(0)
	, mLocalWriteAvail(0)
	, mRecvBase(0)
	, mRecvLimit(0)
	, mSendBase(0)
	, mSendLimit(0)
{
}

void ATNetSockBridgeHandler::Shutdown() {
	AddRef();

	if (mpLocalSocket) {
		mpLocalSocket->Shutdown();
		mpLocalSocket = NULL;
	}

	if (mNativeSocket != INVALID_SOCKET) {
		mpParent->DeleteConnection(mNativeSocket);
		closesocket(mNativeSocket);
		mNativeSocket = INVALID_SOCKET;
	}

	Release();
}

void ATNetSockBridgeHandler::OnNativeSocketConnect() {
	mbNativeConnected = true;
	TryCopyToNative();
	TryCopyFromNative();
}

void ATNetSockBridgeHandler::OnNativeSocketReadReady() {
	TryCopyFromNative();
}

void ATNetSockBridgeHandler::OnNativeSocketWriteReady() {
	TryCopyToNative();
}

void ATNetSockBridgeHandler::OnNativeSocketClose() {
	mbNativeClosed = true;

	mpLocalSocket->Close();

	if (mbLocalClosed)
		Shutdown();
}

void ATNetSockBridgeHandler::OnNativeSocketError() {
	Shutdown();
}

void ATNetSockBridgeHandler::OnSocketOpen() {
}

void ATNetSockBridgeHandler::OnSocketReadReady(uint32 len) {
	mLocalReadAvail = len;

	TryCopyToNative();
}

void ATNetSockBridgeHandler::OnSocketWriteReady(uint32 len) {
	mLocalWriteAvail = len;

	TryCopyFromNative();
}

void ATNetSockBridgeHandler::OnSocketClose() {
	mbLocalClosed = true;

	shutdown(mNativeSocket, SD_SEND);

	if (mbNativeClosed)
		Shutdown();
}

void ATNetSockBridgeHandler::OnSocketError() {
	Shutdown();
}

void ATNetSockBridgeHandler::TryCopyToNative() {
	if (!mbNativeConnected)
		return;

	for(;;) {
		if (mSendBase == mSendLimit) {
			if (mbLocalClosed)
				break;

			uint32 actual = mpLocalSocket->Read(mSendBuf, sizeof mSendBuf);

			if (!actual)
				break;

			mSendBase = 0;
			mSendLimit = actual;
		}

		int actual2 = ::send(mNativeSocket, mSendBuf + mSendBase, mSendLimit - mSendBase, 0);

		if (actual2 == 0)
			break;

		if (actual2 == SOCKET_ERROR)
			break;

		mSendBase += actual2;
	}
}

void ATNetSockBridgeHandler::TryCopyFromNative() {
	if (!mbNativeConnected)
		return;

	for(;;) {
		if (mRecvBase == mRecvLimit) {
			int actual = ::recv(mNativeSocket, mRecvBuf, sizeof mRecvBuf, 0);

			if (actual == 0 || actual == SOCKET_ERROR)
				break;

			mRecvBase = 0;
			mRecvLimit = actual;
		}

		uint32 actual2 = mpLocalSocket->Write(mRecvBuf + mRecvBase, mRecvLimit - mRecvBase);
		if (!actual2)
			break;

		mRecvBase += actual2;
	}
}

///////////////////////////////////////////////////////////////////////////

ATNetSockWorker::ATNetSockWorker()
	: mpWndThunk(NULL)
	, mWndClass(NULL)
	, mhwnd(NULL)
	, mpUdpStack(NULL)
{
}

ATNetSockWorker::~ATNetSockWorker() {
	Shutdown();
}

bool ATNetSockWorker::Init(IATNetUdpStack *udp, bool externalAccess) {
	mpUdpStack = udp;
	mbAllowExternalAccess = externalAccess;

	// Bind DNS directly on the gateway, as we have to redirect it to the host's gateway.
	udp->Bind(53, this);

	mpWndThunk = VDCreateFunctionThunkFromMethod(this, &ATNetSockWorker::WndProc, true);
	if (!mpWndThunk) {
		Shutdown();
		return false;
	}

	TCHAR className[64];
	_sntprintf(className, vdcountof(className), _T("ATNetSockWorker_%p"), this);

	WNDCLASS wc = {0};
	wc.lpfnWndProc = (WNDPROC)mpWndThunk;
	wc.hInstance = VDGetLocalModuleHandleW32();
	wc.lpszClassName = className;

	mWndClass = RegisterClass(&wc);
	if (!mWndClass) {
		Shutdown();
		return false;
	}

	mhwnd = CreateWindowEx(0, MAKEINTATOM(mWndClass), _T(""), WS_POPUP, 0, 0, 0, 0, NULL, NULL, wc.hInstance, NULL);
	if (!mhwnd) {
		Shutdown();
		return false;
	}

	return true;
}

void ATNetSockWorker::Shutdown() {
	ResetAllConnections();

	if (mhwnd) {
		DestroyWindow(mhwnd);
		mhwnd = NULL;
	}

	if (mWndClass) {
		UnregisterClass(MAKEINTATOM(mWndClass), VDGetLocalModuleHandleW32());
		mWndClass = NULL;
	}

	if (mpWndThunk) {
		VDDestroyFunctionThunk(mpWndThunk);
		mpWndThunk = NULL;
	}

	if (mpUdpStack) {
		mpUdpStack->Unbind(53, this);
		mpUdpStack = NULL;
	}
}

void ATNetSockWorker::ResetAllConnections() {
	while(!mTcpConnections.empty()) {
		ATNetSockBridgeHandler *h = mTcpConnections.begin()->second;

		h->Shutdown();
	}

	for(UdpSocketMap::const_iterator it = mUdpSocketMap.begin(), itEnd = mUdpSocketMap.end();
		it != itEnd;
		++it)
	{
		::closesocket(it->first);
	}

	mUdpSocketMap.clear();
	mUdpSourceMap.clear();
}

bool ATNetSockWorker::GetHostAddressForLocalAddress(bool tcp, uint32 srcIpAddr, uint16 srcPort, uint32 dstIpAddr, uint16 dstPort, uint32& hostIp, uint16& hostPort) {
	SOCKET s = INVALID_SOCKET;

	if (tcp) {
		for(TcpConnections::const_iterator it = mTcpConnections.begin(), itEnd = mTcpConnections.end();
			it != itEnd;
			++it)
		{
			ATNetSockBridgeHandler *h = it->second;

			if (h->GetSrcIpAddr() == srcIpAddr && h->GetSrcPort() == srcPort &&
				h->GetDstIpAddr() == dstIpAddr && h->GetDstPort() == dstPort)
			{
				s = it->first;
				break;
			}
		}
	} else {
		const UdpConnection conn = { srcIpAddr, dstIpAddr, srcPort, dstPort };

		UdpSourceMap::const_iterator it = mUdpSourceMap.find(conn);

		if (it != mUdpSourceMap.end())
			s = it->second;
	}

	if (s != INVALID_SOCKET) {
		sockaddr_in sa = {};
		int sa_len = sizeof sa;

		if (0 == getsockname(s, (sockaddr *)&sa, &sa_len)) {
			hostIp = sa.sin_addr.S_un.S_addr;
			hostPort = ntohs(sa.sin_port);
			return true;
		}
	}

	return false;
}

bool ATNetSockWorker::OnSocketIncomingConnection(uint32 srcIpAddr, uint16 srcPort, uint32 dstIpAddr, uint16 dstPort, IATSocket *socket, IATSocketHandler **handler) {
	// check if we are allowed to do this connection
	uint32 redirectedDstIpAddr = dstIpAddr;
	if (mpUdpStack->GetIpStack()->IsLocalOrBroadcastAddress(dstIpAddr)) {
		redirectedDstIpAddr = VDToBE32(0x7F000001);		// 127.0.0.1
	} else if (!mbAllowExternalAccess)
		return false;

	SOCKET s = ::socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);

	if (s == INVALID_SOCKET)
		return false;

	::WSAAsyncSelect(s, mhwnd, MYWM_TCP_SOCKET, FD_CONNECT | FD_READ | FD_WRITE | FD_CLOSE);

	sockaddr_in addr = {};
	addr.sin_family = AF_INET;
	addr.sin_port = htons(dstPort);
	addr.sin_addr.S_un.S_addr = redirectedDstIpAddr;
	if (SOCKET_ERROR == ::connect(s, (const sockaddr *)&addr, sizeof addr)) {
		if (WSAGetLastError() != WSAEWOULDBLOCK)
			return false;
	}

	vdrefptr<ATNetSockBridgeHandler> h(new_nothrow ATNetSockBridgeHandler(this, s, socket, srcIpAddr, srcPort, dstIpAddr, dstPort));
	if (!h) {
		closesocket(s);
		return false;
	}

	mTcpConnections[s] = h;
	h->AddRef();

	*handler = h.release();
	return true;
}

void ATNetSockWorker::OnUdpDatagram(const ATEthernetAddr& srcHwAddr, uint32 srcIpAddr, uint16 srcPort, uint32 dstIpAddr, uint16 dstPort, const void *data, uint32 dataLen) {
	// Check if this is a DNS packet. If so, we intercept these and redirect them to the
	// host gateway.
	uint32 redirectedDstIpAddr = dstIpAddr;
	bool redirected = false;

	if (dstPort == 53) {
		// Yes, it is DNS -- first, check if it is allowed.
		if (!mbAllowExternalAccess)
			return;
		
		// Look up a gateway.
		ULONG outLen = 0;
		if (ERROR_BUFFER_OVERFLOW != ::GetAdaptersInfo(NULL, &outLen))
			return;

		vdblock<char> buf(outLen);
		if (ERROR_SUCCESS != ::GetAdaptersInfo((IP_ADAPTER_INFO *)buf.data(), &outLen))
			return;

		// find the first valid gateway
		for(const IP_ADAPTER_INFO *info = (const IP_ADAPTER_INFO *)buf.data();
			info;
			info = info->Next)
		{
			if (!info->GatewayList.IpAddress.String[0])
				continue;

			uint32 addr = inet_addr(info->GatewayList.IpAddress.String);

			// The address string will normally be 0.0.0.0 if there is no gateway.
			if (addr && addr != INADDR_NONE) {
				redirectedDstIpAddr = addr;
				redirected = true;
				break;
			}
		}
	} else {
		// if this connection is aimed at the gateway, spoof the dest address to localhost
		if (mpUdpStack->GetIpStack()->IsLocalOrBroadcastAddress(dstIpAddr)) {
			redirectedDstIpAddr = VDToBE32(0x7F000001);		// 127.0.0.1
			redirected = true;
		}
	}

	// see if we already have a socket set up for this source address
	UdpConnection conn = {0};
	conn.mSrcIpAddr = srcIpAddr;
	conn.mSrcPort = srcPort;

	if (redirected) {
		// if we're redirecting the connection, we must retain the original destination so
		// we can spoof it on the reply
		conn.mDstIpAddr = dstIpAddr;
		conn.mDstPort = dstPort;
	}

	UdpSourceMap::insert_return_type r = mUdpSourceMap.insert(conn);

	if (r.second) {
		// Nope -- establish a new socket.
		SOCKET s = ::socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
		if (s == INVALID_SOCKET)
			return;

		::WSAAsyncSelect(s, mhwnd, MYWM_UDP_SOCKET, FD_READ);

		r.first->second = s;

		mUdpSocketMap[s] = conn;
	}

	sockaddr_in dstAddr = {0};
	dstAddr.sin_family = AF_INET;
	dstAddr.sin_port = htons(dstPort);
	dstAddr.sin_addr.S_un.S_addr = redirectedDstIpAddr;

	::sendto(r.first->second, (const char *)data, dataLen, 0, (const sockaddr *)&dstAddr, sizeof dstAddr);
}

void ATNetSockWorker::DeleteConnection(SOCKET s) {
	TcpConnections::iterator it = mTcpConnections.find(s);

	if (it != mTcpConnections.end()) {
		ATNetSockBridgeHandler *h = it->second;

		mTcpConnections.erase(it);

		h->Release();
	}
}

LRESULT ATNetSockWorker::WndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) {
	if (msg == MYWM_TCP_SOCKET) {
		const SOCKET sock = (SOCKET)wParam;
		const UINT event = LOWORD(lParam);
		const UINT error = HIWORD(lParam);

		TcpConnections::iterator it = mTcpConnections.find(sock);

		if (it == mTcpConnections.end()) {
			VDASSERT(!"Received message for invalid connection.");
		} else {
			ATNetSockBridgeHandler *h = it->second;

			switch(event) {
				case FD_ACCEPT:
					break;

				case FD_CONNECT:
					h->OnNativeSocketConnect();
					break;

				case FD_READ:
					h->OnNativeSocketReadReady();
					break;
					
				case FD_WRITE:
					h->OnNativeSocketWriteReady();
					break;

				case FD_CLOSE:
					if (error)
						h->OnNativeSocketError();
					else
						h->OnNativeSocketClose();
					break;
			}
		}

		return 0;
	} else if (msg == MYWM_UDP_SOCKET) {
		const SOCKET sock = (SOCKET)wParam;
		const UINT event = LOWORD(lParam);
		const UINT error = HIWORD(lParam);

		UdpSocketMap::iterator it = mUdpSocketMap.find(sock);

		if (it == mUdpSocketMap.end()) {
			VDASSERT(!"Received message for invalid connection.");
		} else {
			const UdpConnection& conn = it->second;

			// The connection info is for the original outbound datagram that bound the socket,
			// so we must reverse it for the incoming datagram we're handling here.
			ProcessUdpDatagram(it->first, conn.mDstIpAddr, conn.mDstPort, conn.mSrcIpAddr, conn.mSrcPort);
		}

		return 0;
	}

	return DefWindowProc(hwnd, msg, wParam, lParam);
}

void ATNetSockWorker::ProcessUdpDatagram(SOCKET s, uint32 srcIpAddr, uint16 srcPort, uint32 dstIpAddr, uint16 dstPort) {
	char buf[4096];
	sockaddr_in from;
	int fromlen = sizeof from;

	int len = ::recvfrom(s, buf, sizeof buf, 0, (sockaddr *)&from, &fromlen);

	if (len && len != SOCKET_ERROR)
		mpUdpStack->SendDatagram(srcIpAddr ? srcIpAddr : from.sin_addr.S_un.S_addr, srcPort ? srcPort : ntohs(from.sin_port), dstIpAddr, dstPort, buf, len);
}

///////////////////////////////////////////////////////////////////////////

void ATCreateNetSockWorker(IATNetUdpStack *udp, bool externalAccess, IATNetSockWorker **pp) {
	ATNetSockWorker *p = new ATNetSockWorker;

	if (!p->Init(udp, externalAccess)) {
		delete p;
		throw MyMemoryError();
	}

	p->AddRef();
	*pp = p;
}
