/*
 * Copyright (C) 2009,2010  Barracuda Networks, Inc.
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
 * 02110-1301  USA
 *
 */

#include "icelocaltransport.h"

#include <QHostAddress>
#include <QUdpSocket>
#include <QtCrypto>
#include "objectsession.h"
#include "stunmessage.h"
#include "stuntransaction.h"
#include "stunbinding.h"
#include "stunallocate.h"
#include "turnclient.h"

// don't queue more incoming packets than this per transmit path
#define MAX_PACKET_QUEUE 64

namespace XMPP {

enum
{
	Direct,
	Relayed
};

//----------------------------------------------------------------------------
// SafeUdpSocket
//----------------------------------------------------------------------------
// DOR-safe wrapper for QUdpSocket
class SafeUdpSocket : public QObject
{
	Q_OBJECT

private:
	ObjectSession sess;
	QUdpSocket *sock;
	int writtenCount;

public:
	SafeUdpSocket(QUdpSocket *_sock, QObject *parent = 0) :
		QObject(parent),
		sess(this),
		sock(_sock)
	{
		sock->setParent(this);
		connect(sock, SIGNAL(readyRead()), SLOT(sock_readyRead()));
		connect(sock, SIGNAL(bytesWritten(qint64)), SLOT(sock_bytesWritten(qint64)));

		writtenCount = 0;
	}

	~SafeUdpSocket()
	{
		if(sock)
		{
			QUdpSocket *out = release();
			out->deleteLater();
		}
	}

	QUdpSocket *release()
	{
		sock->disconnect(this);
		sock->setParent(0);
		QUdpSocket *out = sock;
		sock = 0;
		return out;
	}

	QHostAddress localAddress() const
	{
		return sock->localAddress();
	}

	quint16 localPort() const
	{
		return sock->localPort();
	}

	bool hasPendingDatagrams() const
	{
		return sock->hasPendingDatagrams();
	}

	QByteArray readDatagram(QHostAddress *address = 0, quint16 *port = 0)
	{
		if(!sock->hasPendingDatagrams())
			return QByteArray();

		QByteArray buf;
		buf.resize(sock->pendingDatagramSize());
		sock->readDatagram(buf.data(), buf.size(), address, port);
		return buf;
	}

	void writeDatagram(const QByteArray &buf, const QHostAddress &address, quint16 port)
	{
		sock->writeDatagram(buf, address, port);
	}

signals:
	void readyRead();
	void datagramsWritten(int count);

private slots:
	void sock_readyRead()
	{
		emit readyRead();
	}

	void sock_bytesWritten(qint64 bytes)
	{
		Q_UNUSED(bytes);

		++writtenCount;
		sess.deferExclusive(this, "processWritten");
	}

	void processWritten()
	{
		int count = writtenCount;
		writtenCount = 0;

		emit datagramsWritten(count);
	}
};

//----------------------------------------------------------------------------
// IceLocalTransport
//----------------------------------------------------------------------------
class IceLocalTransport::Private : public QObject
{
	Q_OBJECT

public:
	class WriteItem
	{
	public:
		enum Type
		{
			Direct,
			Pool,
			Turn
		};

		Type type;
		QHostAddress addr;
		int port;
	};

	class Written
	{
	public:
		QHostAddress addr;
		int port;
		int count;
	};

	class Datagram
	{
	public:
		QHostAddress addr;
		int port;
		QByteArray buf;
	};

	IceLocalTransport *q;
	ObjectSession sess;
	QUdpSocket *extSock;
	SafeUdpSocket *sock;
	StunTransactionPool *pool;
	StunBinding *stunBinding;
	TurnClient *turn;
	bool turnActivated;
	QHostAddress addr;
	int port;
	QHostAddress refAddr;
	int refPort;
	QHostAddress relAddr;
	int relPort;
	QHostAddress stunBindAddr;
	int stunBindPort;
	QHostAddress stunRelayAddr;
	int stunRelayPort;
	QString stunUser;
	QCA::SecureArray stunPass;
	QString clientSoftware;
	QList<Datagram> in;
	QList<Datagram> inRelayed;
	QList<WriteItem> pendingWrites;
	int retryCount;
	bool stopping;
	int debugLevel;

	Private(IceLocalTransport *_q) :
		QObject(_q),
		q(_q),
		sess(this),
		extSock(0),
		sock(0),
		pool(0),
		stunBinding(0),
		turn(0),
		turnActivated(false),
		port(-1),
		refPort(-1),
		relPort(-1),
		retryCount(0),
		stopping(false),
		debugLevel(IceTransport::DL_None)
	{
	}

	~Private()
	{
		reset();
	}

	void reset()
	{
		sess.reset();

		delete stunBinding;
		stunBinding = 0;

		delete turn;
		turn = 0;
		turnActivated = false;

		if(sock)
		{
			if(extSock)
			{
				sock->release();
				extSock = 0;
			}

			delete sock;
			sock = 0;
		}

		addr = QHostAddress();
		port = -1;

		refAddr = QHostAddress();
		refPort = -1;

		relAddr = QHostAddress();
		relPort = -1;

		in.clear();
		inRelayed.clear();
		pendingWrites.clear();

		retryCount = 0;
		stopping = false;
	}

	void start()
	{
		Q_ASSERT(!sock);

		sess.defer(this, "postStart");
	}

	void stop()
	{
		Q_ASSERT(sock);
		Q_ASSERT(!stopping);

		stopping = true;

		if(turn)
			turn->close();
		else
			sess.defer(this, "postStop");
	}

	void stunStart()
	{
		Q_ASSERT(!pool);

		pool = new StunTransactionPool(StunTransaction::Udp, this);
		pool->setDebugLevel((StunTransactionPool::DebugLevel)debugLevel);
		connect(pool, SIGNAL(outgoingMessage(QByteArray,QHostAddress,int)), SLOT(pool_outgoingMessage(QByteArray,QHostAddress,int)));
		connect(pool, SIGNAL(needAuthParams()), SLOT(pool_needAuthParams()));
		connect(pool, SIGNAL(debugLine(QString)), SLOT(pool_debugLine(QString)));

		pool->setLongTermAuthEnabled(true);
		if(!stunUser.isEmpty())
		{
			pool->setUsername(stunUser);
			pool->setPassword(stunPass);
		}

		if(!stunBindAddr.isNull())
		{
			stunBinding = new StunBinding(pool);
			connect(stunBinding, SIGNAL(success()), SLOT(binding_success()));
			connect(stunBinding, SIGNAL(error(XMPP::StunBinding::Error)), SLOT(binding_error(XMPP::StunBinding::Error)));
			stunBinding->start(stunBindAddr, stunBindPort);
		}

		if(!stunRelayAddr.isNull())
		{
			do_turn();
		}
	}

	void do_turn()
	{
		turn = new TurnClient(this);
		turn->setDebugLevel((TurnClient::DebugLevel)debugLevel);
		connect(turn, SIGNAL(connected()), SLOT(turn_connected()));
		connect(turn, SIGNAL(tlsHandshaken()), SLOT(turn_tlsHandshaken()));
		connect(turn, SIGNAL(closed()), SLOT(turn_closed()));
		connect(turn, SIGNAL(activated()), SLOT(turn_activated()));
		connect(turn, SIGNAL(packetsWritten(int,QHostAddress,int)), SLOT(turn_packetsWritten(int,QHostAddress,int)));
		connect(turn, SIGNAL(error(XMPP::TurnClient::Error)), SLOT(turn_error(XMPP::TurnClient::Error)));
		connect(turn, SIGNAL(outgoingDatagram(QByteArray)), SLOT(turn_outgoingDatagram(QByteArray)));
		connect(turn, SIGNAL(debugLine(QString)), SLOT(turn_debugLine(QString)));

		turn->setClientSoftwareNameAndVersion(clientSoftware);

		turn->connectToHost(pool, stunRelayAddr, stunRelayPort);
	}

private:
	// note: emits signal on error
	QUdpSocket *createSocket()
	{
		QUdpSocket *qsock = new QUdpSocket(this);
		if(!qsock->bind(addr, 0))
		{
			delete qsock;
			emit q->error(IceLocalTransport::ErrorBind);
			return 0;
		}

		return qsock;
	}

	void prepareSocket()
	{
		addr = sock->localAddress();
		port = sock->localPort();

		connect(sock, SIGNAL(readyRead()), SLOT(sock_readyRead()));
		connect(sock, SIGNAL(datagramsWritten(int)), SLOT(sock_datagramsWritten(int)));
	}

	// return true if we are retrying, false if we should error out
	bool handleRetry()
	{
		// don't allow retrying if activated or stopping)
		if(turnActivated || stopping)
			return false;

		++retryCount;
		if(retryCount < 3)
		{
			if(debugLevel >= IceTransport::DL_Info)
				emit q->debugLine("retrying...");

			delete sock;
			sock = 0;

			// to receive this error, it is a Relay, so change
			//   the mode
			//stunType = IceLocalTransport::Relay;

			QUdpSocket *qsock = createSocket();
			if(!qsock)
			{
				// signal emitted in this case.  bail.
				//   (return true so caller takes no action)
				return true;
			}

			sock = new SafeUdpSocket(qsock, this);

			prepareSocket();

			refAddr = QHostAddress();
			refPort = -1;

			relAddr = QHostAddress();
			relPort = -1;

			do_turn();

			// tell the world that our local address probably
			//   changed, and that we lost our reflexive address
			emit q->addressesChanged();
			return true;
		}

		return false;
	}

	// return true if data packet, false if pool or nothing
	bool processIncomingStun(const QByteArray &buf, const QHostAddress &fromAddr, int fromPort, Datagram *dg)
	{
		QByteArray data;
		QHostAddress dataAddr;
		int dataPort;

		bool notStun;
		if(!pool->writeIncomingMessage(buf, &notStun, fromAddr, fromPort) && turn)
		{
			data = turn->processIncomingDatagram(buf, notStun, &dataAddr, &dataPort);
			if(!data.isNull())
			{
				dg->addr = dataAddr;
				dg->port = dataPort;
				dg->buf = data;
				return true;
			}
			else
			{
				if(debugLevel >= IceTransport::DL_Packet)
					emit q->debugLine("Warning: server responded with what doesn't seem to be a STUN or data packet, skipping.");
			}
		}

		return false;
	}

private slots:
	void postStart()
	{
		if(stopping)
			return;

		if(extSock)
		{
			sock = new SafeUdpSocket(extSock, this);
		}
		else
		{
			QUdpSocket *qsock = createSocket();
			if(!qsock)
			{
				// signal emitted in this case.  bail
				return;
			}

			sock = new SafeUdpSocket(qsock, this);
		}

		prepareSocket();

		emit q->started();
	}

	void postStop()
	{
		reset();
		emit q->stopped();
	}

	void sock_readyRead()
	{
		ObjectSessionWatcher watch(&sess);

		QList<Datagram> dreads;
		QList<Datagram> rreads;

		while(sock->hasPendingDatagrams())
		{
			QHostAddress from;
			quint16 fromPort;

			Datagram dg;

			QByteArray buf = sock->readDatagram(&from, &fromPort);
			if((from == stunBindAddr && fromPort == stunBindPort) || (from == stunRelayAddr && fromPort == stunRelayPort))
			{
				bool haveData = processIncomingStun(buf, from, fromPort, &dg);

				// processIncomingStun could cause signals to
				//   emit.  for example, stopped()
				if(!watch.isValid())
					return;

				if(haveData)
					rreads += dg;
			}
			else
			{
				dg.addr = from;
				dg.port = fromPort;
				dg.buf = buf;
				dreads += dg;
			}
		}

		if(dreads.count() > 0)
		{
			in += dreads;
			emit q->readyRead(Direct);
			if(!watch.isValid())
				return;
		}

		if(rreads.count() > 0)
		{
			inRelayed += rreads;
			emit q->readyRead(Relayed);
		}
	}

	void sock_datagramsWritten(int count)
	{
		QList<Written> dwrites;
		int twrites = 0;

		while(count > 0)
		{
			Q_ASSERT(!pendingWrites.isEmpty());
			WriteItem wi = pendingWrites.takeFirst();
			--count;

			if(wi.type == WriteItem::Direct)
			{
				int at = -1;
				for(int n = 0; n < dwrites.count(); ++n)
				{
					if(dwrites[n].addr == wi.addr && dwrites[n].port == wi.port)
					{
						at = n;
						break;
					}
				}

				if(at != -1)
				{
					++dwrites[at].count;
				}
				else
				{
					Written wr;
					wr.addr = wi.addr;
					wr.port = wi.port;
					wr.count = 1;
					dwrites += wr;
				}
			}
			else if(wi.type == WriteItem::Turn)
				++twrites;
		}

		if(dwrites.isEmpty() && twrites == 0)
			return;

		ObjectSessionWatcher watch(&sess);

		if(!dwrites.isEmpty())
		{
			foreach(const Written &wr, dwrites)
			{
				emit q->datagramsWritten(Direct, wr.count, wr.addr, wr.port);
				if(!watch.isValid())
					return;
			}
		}

		if(twrites > 0)
		{
			// note: this will invoke turn_packetsWritten()
			turn->outgoingDatagramsWritten(twrites);
		}
	}

	void pool_outgoingMessage(const QByteArray &packet, const QHostAddress &toAddress, int toPort)
	{
		// warning: read StunTransactionPool docs before modifying
		//   this function

		WriteItem wi;
		wi.type = WriteItem::Pool;
		pendingWrites += wi;
		sock->writeDatagram(packet, toAddress, toPort);
	}

	void pool_needAuthParams()
	{
		// we can get this signal if the user did not provide
		//   creds to us.  however, since this class doesn't support
		//   prompting just continue on as if we had a blank
		//   user/pass
		pool->continueAfterParams();
	}

	void pool_debugLine(const QString &line)
	{
		emit q->debugLine(line);
	}

	void binding_success()
	{
		refAddr = stunBinding->reflexiveAddress();
		refPort = stunBinding->reflexivePort();

		delete stunBinding;
		stunBinding = 0;

		emit q->addressesChanged();
	}

	void binding_error(XMPP::StunBinding::Error e)
	{
		Q_UNUSED(e);

		delete stunBinding;
		stunBinding = 0;

		// don't report any error
		//if(stunType == IceLocalTransport::Basic || (stunType == IceLocalTransport::Auto && !turn))
		//	emit q->addressesChanged();
	}

	void turn_connected()
	{
		if(debugLevel >= IceTransport::DL_Info)
			emit q->debugLine("turn_connected");
	}

	void turn_tlsHandshaken()
	{
		if(debugLevel >= IceTransport::DL_Info)
			emit q->debugLine("turn_tlsHandshaken");
	}

	void turn_closed()
	{
		if(debugLevel >= IceTransport::DL_Info)
			emit q->debugLine("turn_closed");

		delete turn;
		turn = 0;
		turnActivated = false;

		postStop();
	}

	void turn_activated()
	{
		StunAllocate *allocate = turn->stunAllocate();

		// take reflexive address from TURN only if we are not using a
		//   separate STUN server
		if(stunBindAddr.isNull() || stunBindAddr == stunRelayAddr)
		{
			refAddr = allocate->reflexiveAddress();
			refPort = allocate->reflexivePort();
		}

		if(debugLevel >= IceTransport::DL_Info)
			emit q->debugLine(QString("Server says we are ") + allocate->reflexiveAddress().toString() + ';' + QString::number(allocate->reflexivePort()));

		relAddr = allocate->relayedAddress();
		relPort = allocate->relayedPort();
		if(debugLevel >= IceTransport::DL_Info)
			emit q->debugLine(QString("Server relays via ") + relAddr.toString() + ';' + QString::number(relPort));

		turnActivated = true;

		emit q->addressesChanged();
	}

	void turn_packetsWritten(int count, const QHostAddress &addr, int port)
	{
		emit q->datagramsWritten(Relayed, count, addr, port);
	}

	void turn_error(XMPP::TurnClient::Error e)
	{
		if(debugLevel >= IceTransport::DL_Info)
			emit q->debugLine(QString("turn_error: ") + turn->errorString());

		delete turn;
		turn = 0;
		bool wasActivated = turnActivated;
		turnActivated = false;

		if(e == TurnClient::ErrorMismatch)
		{
			if(!extSock && handleRetry())
				return;
		}

		// this means our relay died on us.  in the future we might
		//   consider reporting this
		if(wasActivated)
			return;

		// don't report any error
		//if(stunType == IceLocalTransport::Relay || (stunType == IceLocalTransport::Auto && !stunBinding))
		//	emit q->addressesChanged();
	}

	void turn_outgoingDatagram(const QByteArray &buf)
	{
		WriteItem wi;
		wi.type = WriteItem::Turn;
		pendingWrites += wi;
		sock->writeDatagram(buf, stunRelayAddr, stunRelayPort);
	}

	void turn_debugLine(const QString &line)
	{
		emit q->debugLine(line);
	}
};

IceLocalTransport::IceLocalTransport(QObject *parent) :
	IceTransport(parent)
{
	d = new Private(this);
}

IceLocalTransport::~IceLocalTransport()
{
	delete d;
}

void IceLocalTransport::setClientSoftwareNameAndVersion(const QString &str)
{
	d->clientSoftware = str;
}

void IceLocalTransport::start(QUdpSocket *sock)
{
	d->extSock = sock;
	d->start();
}

void IceLocalTransport::start(const QHostAddress &addr)
{
	d->addr = addr;
	d->start();
}

void IceLocalTransport::stop()
{
	d->stop();
}

void IceLocalTransport::setStunBindService(const QHostAddress &addr, int port)
{
	d->stunBindAddr = addr;
	d->stunBindPort = port;
}

void IceLocalTransport::setStunRelayService(const QHostAddress &addr, int port, const QString &user, const QCA::SecureArray &pass)
{
	d->stunRelayAddr = addr;
	d->stunRelayPort = port;
	d->stunUser = user;
	d->stunPass = pass;
}

void IceLocalTransport::stunStart()
{
	d->stunStart();
}

QHostAddress IceLocalTransport::localAddress() const
{
	return d->addr;
}

int IceLocalTransport::localPort() const
{
	return d->port;
}

QHostAddress IceLocalTransport::serverReflexiveAddress() const
{
	return d->refAddr;
}

int IceLocalTransport::serverReflexivePort() const
{
	return d->refPort;
}

QHostAddress IceLocalTransport::relayedAddress() const
{
	return d->relAddr;
}

int IceLocalTransport::relayedPort() const
{
	return d->relPort;
}

void IceLocalTransport::addChannelPeer(const QHostAddress &addr, int port)
{
	if(d->turn)
		d->turn->addChannelPeer(addr, port);
}

bool IceLocalTransport::hasPendingDatagrams(int path) const
{
	if(path == Direct)
		return !d->in.isEmpty();
	else if(path == Relayed)
		return !d->inRelayed.isEmpty();
	else
	{
		Q_ASSERT(0);
		return false;
	}
}

QByteArray IceLocalTransport::readDatagram(int path, QHostAddress *addr, int *port)
{
	QList<Private::Datagram> *in = 0;
	if(path == Direct)
		in = &d->in;
	else if(path == Relayed)
		in = &d->inRelayed;
	else
		Q_ASSERT(0);

	if(!in->isEmpty())
	{
		Private::Datagram datagram = in->takeFirst();
		*addr = datagram.addr;
		*port = datagram.port;
		return datagram.buf;
	}
	else
		return QByteArray();
}

void IceLocalTransport::writeDatagram(int path, const QByteArray &buf, const QHostAddress &addr, int port)
{
	if(path == Direct)
	{
		Private::WriteItem wi;
		wi.type = Private::WriteItem::Direct;
		wi.addr = addr;
		wi.port = port;
		d->pendingWrites += wi;
		d->sock->writeDatagram(buf, addr, port);
	}
	else if(path == Relayed)
	{
		if(d->turn && d->turnActivated)
			d->turn->write(buf, addr, port);
	}
	else
		Q_ASSERT(0);
}

void IceLocalTransport::setDebugLevel(DebugLevel level)
{
	d->debugLevel = level;
	if(d->pool)
		d->pool->setDebugLevel((StunTransactionPool::DebugLevel)level);
	if(d->turn)
		d->turn->setDebugLevel((TurnClient::DebugLevel)level);
}

}

#include "icelocaltransport.moc"
