/*
 * Copyright (c) 2010-2011 Thilo Cestonaro
 * Licensed under the simplified BSD license.
 * See Documentation/Licenses/BSD-simplified.txt for more information.
 */

#include <cppunit/extensions/HelperMacros.h>
#include <cppunit/extensions/TestFactoryRegistry.h>

#include <boost/optional.hpp>
#include <boost/bind.hpp>
#include <boost/smart_ptr/make_shared.hpp>
#include <boost/shared_ptr.hpp>

#include <Swiften/Network/Connection.h>
#include <Swiften/Network/ConnectionFactory.h>
#include <Swiften/Network/SOCKS5ProxiedConnection.h>
#include <Swiften/Network/HostAddressPort.h>
#include <Swiften/EventLoop/DummyEventLoop.h>

using namespace Swift;

class SOCKS5ProxiedConnectionTest : public CppUnit::TestFixture {
		CPPUNIT_TEST_SUITE(SOCKS5ProxiedConnectionTest);
		CPPUNIT_TEST(testConnect_CreatesConnectionToProxy);
		CPPUNIT_TEST(testConnect_ConnectHandshake);
		CPPUNIT_TEST(testConnect_ReceiveMalformedConnectResponse);
		CPPUNIT_TEST(testConnect_ReceiveConnectResponseWithUnsupportedAuthMethod);
		CPPUNIT_TEST_SUITE_END();

	public:
		void setUp() {
			proxyHost = HostAddressPort(HostAddress("1.1.1.1"), 1234);
			host = HostAddressPort(HostAddress("2.2.2.2"), 2345);
			eventLoop = new DummyEventLoop();
			connectionFactory = new MockConnectionFactory(eventLoop);
			connectFinished = false;
			disconnected = false;
		}

		void tearDown() {
			delete connectionFactory;
			delete eventLoop;
		}

		void testConnect_CreatesConnectionToProxy() {
			SOCKS5ProxiedConnection::ref testling(createTestling());

			testling->connect(host);
			eventLoop->processEvents();

			CPPUNIT_ASSERT_EQUAL(1, static_cast<int>(connectionFactory->connections.size()));
			CPPUNIT_ASSERT(connectionFactory->connections[0]->hostAddressPort);
			CPPUNIT_ASSERT(proxyHost == *connectionFactory->connections[0]->hostAddressPort);
			CPPUNIT_ASSERT(!connectFinished);
		}

		void testConnect_ConnectHandshake() {
			ByteArray connectResponse;
			ByteArray connectRequest;
			SOCKS5ProxiedConnection::ref testling(createTestling());
			testling->connect(HostAddressPort(HostAddress("2.2.2.2"), 2345));
			eventLoop->processEvents();

			// response is always first byte 0x05 and second 0x00
			connectResponse += 0x05;
			connectResponse += 0x00;

			// first part of the connection handshake
			connectRequest += 0x05; // Version: SOCKS5 = 0x05
			connectRequest += 0x01; // Number of authentication methods after this byte.
			connectRequest += 0x00; // authmethod = noAuth = 0x00

			CPPUNIT_ASSERT_EQUAL(connectRequest, connectionFactory->connections[0]->dataWritten);

			// response 
			connectionFactory->connections[0]->onDataRead(connectResponse);
			eventLoop->processEvents();


			// second part of the connection handshake
			connectRequest += 0x05; // Version: SOCKS5 = 0x05
			connectRequest += 0x01; // Number of authentication methods after this byte.
			connectRequest += 0x00; // authmethod = noAuth = 0x00
			connectRequest += 0x01; // IPv4
			connectRequest += 0x02; // IP
			connectRequest += 0x02; // IP
			connectRequest += 0x02; // IP
			connectRequest += 0x02; // IP
			connectRequest += 0x09; // Port
			connectRequest += 0x29; // Port
			
			CPPUNIT_ASSERT_EQUAL(connectRequest, connectionFactory->connections[0]->dataWritten);

			// response 
			connectionFactory->connections[0]->onDataRead(connectResponse);
			eventLoop->processEvents();

			CPPUNIT_ASSERT(connectFinished);
			CPPUNIT_ASSERT(!connectFinishedWithError);
			CPPUNIT_ASSERT(dataRead.isEmpty());
		}

		void testConnect_ReceiveMalformedConnectResponse() {
			SOCKS5ProxiedConnection::ref testling(createTestling());
			testling->connect(HostAddressPort(HostAddress("2.2.2.2"), 2345));
			eventLoop->processEvents();

			connectionFactory->connections[0]->onDataRead(ByteArray("FLOOP"));
			eventLoop->processEvents();

			CPPUNIT_ASSERT(connectFinished);
			CPPUNIT_ASSERT(connectFinishedWithError);
			CPPUNIT_ASSERT(connectionFactory->connections[0]->disconnected);
		}

		void testConnect_ReceiveConnectResponseWithUnsupportedAuthMethod() {
			SOCKS5ProxiedConnection::ref testling(createTestling());
			testling->connect(HostAddressPort(HostAddress("2.2.2.2"), 2345));
			eventLoop->processEvents();

			ByteArray connectResponse;
			connectResponse += 0x05;
			connectResponse += 0x01;

			connectionFactory->connections[0]->onDataRead(connectResponse);
			eventLoop->processEvents();

			CPPUNIT_ASSERT(connectFinished);
			CPPUNIT_ASSERT(connectFinishedWithError);
			CPPUNIT_ASSERT(connectionFactory->connections[0]->disconnected);
		}

		void testConnect_ReceiveConnectResponseWithFailure() {
			ByteArray connectResponse;
			ByteArray connectRequest;
			SOCKS5ProxiedConnection::ref testling(createTestling());
			testling->connect(HostAddressPort(HostAddress("2.2.2.2"), 2345));
			eventLoop->processEvents();

			// response is always first byte 0x05 and second 0x00
			connectResponse += 0x05;
			connectResponse += 0x00;

			// first part of the connection handshake
			connectRequest += 0x05; // Version: SOCKS5 = 0x05
			connectRequest += 0x01; // Number of authentication methods after this byte.
			connectRequest += 0x00; // authmethod = noAuth = 0x00

			CPPUNIT_ASSERT_EQUAL(connectRequest, connectionFactory->connections[0]->dataWritten);

			// response 
			connectionFactory->connections[0]->onDataRead(connectResponse);
			eventLoop->processEvents();


			// second part of the connection handshake
			connectRequest += 0x05; // Version: SOCKS5 = 0x05
			connectRequest += 0x01; // Number of authentication methods after this byte.
			connectRequest += 0x00; // authmethod = noAuth = 0x00
			connectRequest += 0x01; // IPv4
			connectRequest += 0x02; // IP
			connectRequest += 0x02; // IP
			connectRequest += 0x02; // IP
			connectRequest += 0x02; // IP
			connectRequest += 0x09; // Port
			connectRequest += 0x29; // Port
			
			CPPUNIT_ASSERT_EQUAL(connectRequest, connectionFactory->connections[0]->dataWritten);

			// response 
			connectResponse.clear();
			connectResponse += 0x05;
			connectResponse += 0x01; // general SOCKS failure
			connectionFactory->connections[0]->onDataRead(connectResponse);
			eventLoop->processEvents();

			CPPUNIT_ASSERT(connectFinished);
			CPPUNIT_ASSERT(connectFinishedWithError);
			CPPUNIT_ASSERT(connectionFactory->connections[0]->disconnected);
		}

/*
		void testConnect_ReceiveDataAfterConnect() {
			SOCKS5ProxiedConnection::ref testling(createTestling());
			testling->connect(HostAddressPort(HostAddress("2.2.2.2"), 2345));
			eventLoop->processEvents();
			connectionFactory->connections[0]->onDataRead(ByteArray("HTTP/1.0 200 Connection established\r\n\r\n"));
			eventLoop->processEvents();

			connectionFactory->connections[0]->onDataRead(ByteArray("abcdef"));

			CPPUNIT_ASSERT_EQUAL(ByteArray("abcdef"), dataRead);
		}

		void testWrite_AfterConnect() {
			SOCKS5ProxiedConnection::ref testling(createTestling());
			testling->connect(HostAddressPort(HostAddress("2.2.2.2"), 2345));
			eventLoop->processEvents();
			connectionFactory->connections[0]->onDataRead(ByteArray("HTTP/1.0 200 Connection established\r\n\r\n"));
			eventLoop->processEvents();
			connectionFactory->connections[0]->dataWritten.clear();

			testling->write(ByteArray("abcdef"));

			CPPUNIT_ASSERT_EQUAL(ByteArray("abcdef"), connectionFactory->connections[0]->dataWritten);
		}

		void testDisconnect_AfterConnectRequest() {
			SOCKS5ProxiedConnection::ref testling(createTestling());
			testling->connect(HostAddressPort(HostAddress("2.2.2.2"), 2345));
			eventLoop->processEvents();

			testling->disconnect();

			CPPUNIT_ASSERT(connectionFactory->connections[0]->disconnected);
			CPPUNIT_ASSERT(disconnected);
			CPPUNIT_ASSERT(!disconnectedError);
		}

		void testDisconnect_AfterConnect() {
			SOCKS5ProxiedConnection::ref testling(createTestling());
			testling->connect(HostAddressPort(HostAddress("2.2.2.2"), 2345));
			eventLoop->processEvents();
			connectionFactory->connections[0]->onDataRead(ByteArray("HTTP/1.0 200 Connection established\r\n\r\n"));
			eventLoop->processEvents();

			testling->disconnect();

			CPPUNIT_ASSERT(connectionFactory->connections[0]->disconnected);
			CPPUNIT_ASSERT(disconnected);
			CPPUNIT_ASSERT(!disconnectedError);
		}
*/
	private:
		SOCKS5ProxiedConnection::ref createTestling() {
			boost::shared_ptr<SOCKS5ProxiedConnection> c = SOCKS5ProxiedConnection::create(connectionFactory, proxyHost);
			c->onConnectFinished.connect(boost::bind(&SOCKS5ProxiedConnectionTest::handleConnectFinished, this, _1));
			c->onDisconnected.connect(boost::bind(&SOCKS5ProxiedConnectionTest::handleDisconnected, this, _1));
			c->onDataRead.connect(boost::bind(&SOCKS5ProxiedConnectionTest::handleDataRead, this, _1));
			return c;
		}

		void handleConnectFinished(bool error) {
			connectFinished = true;
			connectFinishedWithError = error;
		}

		void handleDisconnected(const boost::optional<Connection::Error>& e) {
			disconnected = true;
			disconnectedError = e;
		}

		void handleDataRead(const ByteArray& d) {
			dataRead += d;
		}

		struct MockConnection : public Connection {
			public:
				MockConnection(const std::vector<HostAddressPort>& failingPorts, EventLoop* eventLoop) : eventLoop(eventLoop), failingPorts(failingPorts), disconnected(false) {
				}

				void listen() { assert(false); }

				void connect(const HostAddressPort& address) {
					hostAddressPort = address;
					bool fail = std::find(failingPorts.begin(), failingPorts.end(), address) != failingPorts.end();
					eventLoop->postEvent(boost::bind(boost::ref(onConnectFinished), fail));
				}

				HostAddressPort getLocalAddress() const { return HostAddressPort(); }

				void disconnect() { 
					disconnected = true;
					onDisconnected(boost::optional<Connection::Error>());
				}
				
				void write(const ByteArray& d) { 
					dataWritten += d;
				}

				EventLoop* eventLoop;
				boost::optional<HostAddressPort> hostAddressPort;
				std::vector<HostAddressPort> failingPorts;
				ByteArray dataWritten;
				bool disconnected;
		};

		struct MockConnectionFactory : public ConnectionFactory {
			MockConnectionFactory(EventLoop* eventLoop) : eventLoop(eventLoop) {
			}

			boost::shared_ptr<Connection> createConnection() {
				boost::shared_ptr<MockConnection> connection = boost::make_shared<MockConnection>(failingPorts, eventLoop);
				connections.push_back(connection);
				return connection;
			}

			EventLoop* eventLoop;
			std::vector< boost::shared_ptr<MockConnection> > connections;
			std::vector<HostAddressPort> failingPorts;
		};

	private:
		HostAddressPort proxyHost;
		HostAddressPort host;
		DummyEventLoop* eventLoop;
		MockConnectionFactory* connectionFactory;
		std::vector< boost::shared_ptr<MockConnection> > connections;
		bool connectFinished;
		bool connectFinishedWithError;
		bool disconnected;
		boost::optional<Connection::Error> disconnectedError;
		ByteArray dataRead;
};

CPPUNIT_TEST_SUITE_REGISTRATION(SOCKS5ProxiedConnectionTest);