/*
 * Copyright (C) 2009  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"

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

namespace XMPP {

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

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

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

		writtenCount = 0;
	}

	~SafeUdpSocket()
	{
		sock->disconnect(this);
		sock->setParent(0);
		sock->deleteLater();
	}

	bool bind(const QHostAddress &addr, quint16 port = 0)
	{
		return sock->bind(addr, port);
	}

	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:
	enum WriteType
	{
		InternalWrite,
		DirectWrite,
		RelayedWrite
	};

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

	IceLocalTransport *q;
	ObjectSession sess;
	SafeUdpSocket *sock;
	StunTransactionPool *pool;
	StunBinding *stunBinding;
	StunAllocate *stunAllocate;
	bool alloc_started;
	bool changing_perms;
	QHostAddress addr;
	int port;
	QHostAddress refAddr;
	int refPort;
	QHostAddress relAddr;
	int relPort;
	QHostAddress stunAddr;
	int stunPort;
	IceLocalTransport::StunServiceType stunType;
	QString stunUser;
	QCA::SecureArray stunPass;
	QList<Datagram> in;
	QList<Datagram> inRelayed;
	QList<Datagram> outRelayed;
	QList<WriteType> pendingWrites;
	QList<QHostAddress> pendingPerms;

	Private(IceLocalTransport *_q) :
		QObject(_q),
		q(_q),
		sess(this),
		sock(0),
		stunBinding(0),
		stunAllocate(0),
		alloc_started(false),
		changing_perms(false),
		port(-1),
		refPort(-1),
		relPort(-1)
	{
		pool = new StunTransactionPool(StunTransaction::Udp, this);
		connect(pool, SIGNAL(retransmit(XMPP::StunTransaction*)), SLOT(pool_retransmit(XMPP::StunTransaction*)));
	}

	~Private()
	{
		reset();
	}

	void reset()
	{
		sess.reset();

		delete stunBinding;
		stunBinding = 0;

		delete stunAllocate;
		stunAllocate = 0;
		alloc_started = false;
		changing_perms = false;

		delete sock;
		sock = 0;

		addr = QHostAddress();
		port = -1;

		refAddr = QHostAddress();
		refPort = -1;

		relAddr = QHostAddress();
		relPort = -1;

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

	void start()
	{
		Q_ASSERT(!sock);

		sock = new SafeUdpSocket(this);
		connect(sock, SIGNAL(readyRead()), SLOT(sock_readyRead()));
		connect(sock, SIGNAL(datagramsWritten(int)), SLOT(sock_datagramsWritten(int)));

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

	void stop()
	{
		Q_ASSERT(sock);

		if(stunAllocate)
			stunAllocate->stop();
		else
			sess.defer(this, "postStop");
	}

	void stunStart()
	{
		Q_ASSERT(!stunBinding && !stunAllocate);

		if(stunType == IceLocalTransport::Relay)
		{
			stunAllocate = new StunAllocate(pool);
			connect(stunAllocate, SIGNAL(started()), SLOT(allocate_started()));
			connect(stunAllocate, SIGNAL(stopped()), SLOT(allocate_stopped()));
			connect(stunAllocate, SIGNAL(error(XMPP::StunAllocate::Error)), SLOT(allocate_error(XMPP::StunAllocate::Error)));
			connect(stunAllocate, SIGNAL(permissionsChanged()), SLOT(allocate_permissionsChanged()));
			connect(stunAllocate, SIGNAL(readyRead()), SLOT(allocate_readyRead()));
			connect(stunAllocate, SIGNAL(datagramsWritten(int)), SLOT(allocate_datagramsWritten(int)));
			stunAllocate->start();
		}
		else // Basic
		{
			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();
		}
	}

	void tryWriteRelayed(const QByteArray &buf, const QHostAddress &addr, int port)
	{
		QList<QHostAddress> perms = stunAllocate->permissions();

		// do we have permission to relay to this address yet?
		if(perms.contains(addr))
		{
			writeRelayed(buf, addr, port);
		}
		else
		{
			// no?  then queue while we ask the server to grant
			Datagram dg;
			dg.addr = addr;
			dg.port = port;
			dg.buf = buf;
			outRelayed += dg;

			if(!changing_perms)
			{
				perms += addr;
				stunAllocate->setPermissions(perms);
			}
			else
				pendingPerms += addr;
		}
	}

	void writeRelayed(const QByteArray &buf, const QHostAddress &addr, int port)
	{
		QByteArray enc = stunAllocate->encode(buf, addr, port);
		if(enc.isEmpty())
		{
			printf("Warning: could not encode packet for sending.\n");
			return;
		}

		pendingWrites += RelayedWrite;
		sock->writeDatagram(enc, stunAddr, stunPort);
	}

	// return true if we received a relayed packet
	bool processIncomingStun(const QByteArray &buf, Datagram *dg)
	{
		// this might be a ChannelData message.  check the first
		//   two bits:
		if(stunAllocate && alloc_started && buf.size() >= 1 && (buf[0] & 0xC0) == 0x40)
		{
			QHostAddress fromAddr;
			int fromPort;
			QByteArray buf = stunAllocate->decode(buf, &fromAddr, &fromPort);
			if(fromAddr.isNull())
			{
				printf("Warning: server responded with what appears to be an invalid packet, skipping.\n");
				return false;
			}

			dg->addr = fromAddr;
			dg->port = fromPort;
			dg->buf = buf;
			return true;
		}

		// else, interpret it as a stun message
		StunMessage message = StunMessage::fromBinary(buf);
		if(message.isNull())
		{
			printf("Warning: server responded with what doesn't seem to be a STUN packet, skipping.\n");
			return false;
		}

		// indication?  maybe it's a relayed packet
		if(message.mclass() == StunMessage::Indication)
		{
			QHostAddress fromAddr;
			int fromPort;
			QByteArray buf = stunAllocate->decode(message, &fromAddr, &fromPort);
			if(fromAddr.isNull())
			{
				printf("Warning: server responded with an unknown Indication packet, skipping.\n");
				return false;
			}

			dg->addr = fromAddr;
			dg->port = fromPort;
			dg->buf = buf;
			return true;
		}

		if(!pool->writeIncomingMessage(message))
		{
			printf("Warning: received unexpected message, skipping.\n");
		}

		return false;
	}

public slots:
	void postStart()
	{
		bool ok;
		if(port != -1)
			ok = sock->bind(addr, port);
		else
			ok = sock->bind(addr, 0);

		if(ok)
		{
			port = sock->localPort();
			emit q->started();
		}
		else
		{
			reset();
			emit q->error(IceLocalTransport::ErrorGeneric);
		}
	}

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

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

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

		while(sock->hasPendingDatagrams())
		{
			QHostAddress from;
			quint16 fromPort;
			QByteArray buf = sock->readDatagram(&from, &fromPort);

			Datagram dg;

			if(from == stunAddr && fromPort == stunPort)
			{
				// came from stun server
				if(processIncomingStun(buf, &dg))
					rreads += dg;

				if(!watcher.isValid())
					return;
			}
			else
			{
				dg.addr = from;
				dg.port = fromPort;
				dg.buf = buf;
				dreads += dg;
			}
		}

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

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

	void sock_datagramsWritten(int count)
	{
		Q_ASSERT(count <= pendingWrites.count());

		int dwrites = 0;
		int rwrites = 0;
		for(int n = 0; n < count; ++n)
		{
			WriteType type = pendingWrites.takeFirst();
			if(type == DirectWrite)
				++dwrites;
			else if(type == RelayedWrite)
				++rwrites;
		}

		ObjectSessionWatcher watch(&sess);

		if(dwrites > 0)
		{
			emit q->datagramsWritten(IceLocalTransport::Direct, dwrites);
			if(!watch.isValid())
				return;
		}

		if(rwrites > 0)
			emit q->datagramsWritten(IceLocalTransport::Relayed, rwrites);
	}

	void pool_retransmit(XMPP::StunTransaction *trans)
	{
		// warning: read StunTransactionPool docs before modifying
		//   this function

		pendingWrites += InternalWrite;
		sock->writeDatagram(trans->packet(), stunAddr, stunPort);
	}

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

		delete stunBinding;
		stunBinding = 0;

		emit q->stunFinished();
	}

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

		delete stunBinding;
		stunBinding = 0;

		emit q->stunFinished();
	}

	void allocate_started()
	{
		refAddr = stunAllocate->reflexiveAddress();
		refPort = stunAllocate->reflexivePort();
		relAddr = stunAllocate->relayedAddress();
		relPort = stunAllocate->relayedPort();
		alloc_started = true;

		emit q->stunFinished();
	}

	void allocate_stopped()
	{
		// allocation deleted
		delete stunAllocate;
		stunAllocate = 0;
		alloc_started = false;

		postStop();
	}

	void allocate_error(XMPP::StunAllocate::Error e)
	{
		delete stunAllocate;
		stunAllocate = 0;
		bool wasStarted = alloc_started;
		alloc_started = false;

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

		// if we get an error during initialization, fall back to
		//   binding
		if(e != StunAllocate::ErrorTimeout)
		{
			stunType = IceLocalTransport::Basic;
			stunStart();
		}
		else
		{
			emit q->stunFinished();
		}
	}

	void allocate_permissionsChanged()
	{
		// get updated list
		QList<QHostAddress> perms = stunAllocate->permissions();

		// extract any sendable packets from the out queue
		QList<Datagram> sendable;
		for(int n = 0; n < outRelayed.count(); ++n)
		{
			if(perms.contains(outRelayed[n].addr))
			{
				sendable += outRelayed[n];
				outRelayed.removeAt(n);
				--n; // adjust position
			}
		}

		// and send them
		foreach(const Datagram &dg, sendable)
			writeRelayed(dg.buf, dg.addr, dg.port);

		changing_perms = false;

		if(!pendingPerms.isEmpty())
		{
			perms += pendingPerms;
			pendingPerms.clear();

			changing_perms = true;
			stunAllocate->setPermissions(perms);
		}
	}
};

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

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

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

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

void IceLocalTransport::setStunService(StunServiceType type, const QHostAddress &addr, int port)
{
	d->stunType = type;
	d->stunAddr = addr;
	d->stunPort = port;
}

void IceLocalTransport::setStunUsername(const QString &user)
{
	d->stunUser = user;
}

void IceLocalTransport::setStunPassword(const QCA::SecureArray &pass)
{
	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;
}

bool IceLocalTransport::hasPendingDatagrams(TransmitPath 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(TransmitPath 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
	{
		*addr = QHostAddress();
		*port = -1;
		return QByteArray();
	}
}

void IceLocalTransport::writeDatagram(TransmitPath path, const QByteArray &buf, const QHostAddress &addr, int port)
{
	if(path == Direct)
	{
		d->pendingWrites += Private::DirectWrite;
		d->sock->writeDatagram(buf, addr, port);
	}
	else if(path == Relayed)
	{
		if(d->stunAllocate && d->alloc_started)
			d->tryWriteRelayed(buf, addr, port);
	}
	else
		Q_ASSERT(0);
}

}

#include "icelocaltransport.moc"
