* (at your option) any later version.

 *

 * This program is distributed in the hope that it will be useful,

 * but WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

 * GNU General Public License for more details.

 *

 * You should have received a copy of the GNU General Public License

 * along with this program; if not, write to the Free Software

 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02111-1301  USA

 */

#pragma once

#include <Swiften/Network/TLSConnectionFactory.h>

#include <Swiften/Network/HostAddressPort.h>

#include <Swiften/TLS/PlatformTLSFactories.h>

#include <Swiften/Network/DomainNameResolveError.h>

#include <Swiften/Network/DomainNameAddressQuery.h>

#include <Swiften/Network/DomainNameResolver.h>

#include <Swiften/Network/HostAddress.h>

#include <Swiften/Network/Connection.h>

#include <Swiften/Base/SafeByteArray.h>

#include "Swiften/Version.h"

#define HAVE_SWIFTEN_3  (SWIFTEN_VERSION >= 0x030000)

#if HAVE_SWIFTEN_3

#include <Swiften/TLS/TLSOptions.h>

#endif

#include <string>

#include <algorithm>

#include <map>

#include <boost/signal.hpp>

namespace Transport {

class Component;

class WebSocketClient {

	public:

		WebSocketClient(Component *component);

		virtual ~WebSocketClient();

		void connectServer(const std::string &u);

		void write(const std::string &data);

		boost::signal<void (const std::string &payload)> onPayloadReceived;

	private:

		void handleDNSResult(const std::vector<Swift::HostAddress>&, boost::optional<Swift::DomainNameResolveError>);

		void handleDataRead(boost::shared_ptr<Swift::SafeByteArray> data);

		void handleConnected(bool error);

	private:

		Component *m_component;

		boost::shared_ptr<Swift::DomainNameAddressQuery> m_dnsQuery;

		boost::shared_ptr<Swift::Connection> m_conn;

		Swift::TLSConnectionFactory *m_tlsConnectionFactory;

		Swift::PlatformTLSFactories *m_tlsFactory;

		std::string m_host;

		std::string m_path;

		std::string m_buffer;

		bool m_upgraded;

};

}

#include "transport/Logging.h"

#include <boost/foreach.hpp>

#include <boost/make_shared.hpp>

#include <map>

#include <iterator>

namespace Transport {

DEFINE_LOGGER(logger, "WebSocketClient");

WebSocketClient::WebSocketClient(Component *component) {

	m_component = component;

	m_upgraded = false;

#if HAVE_SWIFTEN_3

	Swift::TLSOptions o;

#endif

	m_tlsFactory = new Swift::PlatformTLSFactories();

#if HAVE_SWIFTEN_3

	m_tlsConnectionFactory = new Swift::TLSConnectionFactory(m_tlsFactory->getTLSContextFactory(), component->getNetworkFactories()->getConnectionFactory(), o);

#else

	m_tlsConnectionFactory = new Swift::TLSConnectionFactory(m_tlsFactory->getTLSContextFactory(), component->getNetworkFactories()->getConnectionFactory());

#endif

}

WebSocketClient::~WebSocketClient() {

	if (m_conn) {

		m_conn->onDataRead.disconnect(boost::bind(&WebSocketClient::handleDataRead, this, _1));

		m_conn->disconnect();

	}

	delete m_tlsFactory;

	delete m_tlsConnectionFactory;

}

	LOG4CXX_INFO(logger, "Starting DNS query for " << m_host << " " << m_path);

	m_dnsQuery = m_component->getNetworkFactories()->getDomainNameResolver()->createAddressQuery(m_host);

	m_dnsQuery->onResult.connect(boost::bind(&WebSocketClient::handleDNSResult, this, _1, _2));

	m_dnsQuery->run();

}

void WebSocketClient::write(const std::string &data) {

	if (!m_conn) {

		return;

	}

	uint8_t opcode = 129; // UTF8

	if (data.empty()) {

		LOG4CXX_INFO(logger, "pong");

		opcode = 138; // PONG

	}

	// Mask the payload

	char mask_bits[4] = {0x11, 0x22, 0x33, 0x44};

	std::string payload = data;

	for (size_t i = 0; i < data.size(); i++ ) {

		payload[i] = payload[i] ^ mask_bits[i&3];

	}

	if (data.size() <= 125) {

		uint8_t size7 = data.size() + 128; // Mask bit

		m_conn->write(Swift::createSafeByteArray(std::string((char *) &opcode, 1)

			+ std::string((char *) &size7, 1)

			unsigned int size = (size16 == 0 ? size7 : size16);

			if (m_buffer.size() >= size + header_size) {

				std::string payload = m_buffer.substr(header_size, size);

				LOG4CXX_INFO(logger, "< " << payload);

				onPayloadReceived(payload);

				m_buffer.erase(0, size + header_size);

			}

			else if (size == 0) {

				m_buffer.erase(0, header_size);

			}

			else {

				return;

			}

		}

		else {

			return;

		}

	}

}

void WebSocketClient::handleConnected(bool error) {

	if (error) {

		return;

	}

	LOG4CXX_INFO(logger, "Connected to " << m_host);

	std::string req = "";

	req += "GET " + m_path + " HTTP/1.1\r\n";

	req += "Host: " + m_host + ":443\r\n";

	req += "Upgrade: websocket\r\n";

	req += "Connection: Upgrade\r\n";

	req += "Sec-WebSocket-Key: x3JJHMbDL1EzLkh9GBhXDw==\r\n";

	req += "Sec-WebSocket-Version: 13\r\n";

	req += "\r\n";

	m_conn->write(Swift::createSafeByteArray(req));

}

void WebSocketClient::handleDNSResult(const std::vector<Swift::HostAddress> &addrs, boost::optional<Swift::DomainNameResolveError>) {

	m_conn = m_tlsConnectionFactory->createConnection();

	m_conn->onDataRead.connect(boost::bind(&WebSocketClient::handleDataRead, this, _1));

	m_conn->onConnectFinished.connect(boost::bind(&WebSocketClient::handleConnected, this, _1));

	m_conn->connect(Swift::HostAddressPort(addrs[0], 443));

}