/*
 *  IXWebSocketHandshake.h
 *  Author: Benjamin Sergeant
 *  Copyright (c) 2019 Machine Zone, Inc. All rights reserved.
 *  Copyright (c) 2026 Dan R.
 */

#include "IXWebSocketHandshake.h"

#include "IXBase64.h"
#include "IXHttp.h"
#include "IXSocketConnect.h"
#include "IXStrCaseCompare.h"
#include "IXUrlParser.h"
#include "IXUserAgent.h"
#include "IXWebSocketHandshakeKeyGen.h"
#include <algorithm>
#include <iostream>
#include <random>
#include <sstream>

namespace ix
{
	WebSocketHandshake::WebSocketHandshake(
		std::atomic<bool>& requestInitCancellation,
		std::unique_ptr<Socket>& socket,
		WebSocketPerMessageDeflatePtr& perMessageDeflate,
		WebSocketPerMessageDeflateOptions& perMessageDeflateOptions,
		std::atomic<bool>& enablePerMessageDeflate,
		const std::string& serverName)
		: _requestInitCancellation(requestInitCancellation)
		, _socket(socket)
		, _perMessageDeflate(perMessageDeflate)
		, _perMessageDeflateOptions(perMessageDeflateOptions)
		, _enablePerMessageDeflate(enablePerMessageDeflate)
		, _serverName(serverName)
	{
	}

	bool WebSocketHandshake::insensitiveStringCompare(const std::string& a, const std::string& b)
	{
		return CaseInsensitiveLess::cmp(a, b) == 0;
	}

	std::string WebSocketHandshake::genRandomString(const int len)
	{
		std::string alphanum = "0123456789"
							   "ABCDEFGH"
							   "abcdefgh";

		std::random_device r;
		std::default_random_engine e1(r());
		std::uniform_int_distribution<int> dist(0, (int) alphanum.size() - 1);

		std::string s;
		s.resize(len);

		for (int i = 0; i < len; ++i)
		{
			int x = dist(e1);
			s[i] = alphanum[x];
		}

		return s;
	}

	WebSocketInitResult WebSocketHandshake::sendErrorResponse(int code, const std::string& reason)
	{
		std::stringstream ss;
		ss << "HTTP/1.1 ";
		ss << code;
		ss << " ";
		ss << reason;
		ss << "\r\n";
		ss << "Server: " << (_serverName.empty() ? userAgent() : _serverName) << "\r\n";

		// Socket write can only be cancelled through a timeout here, not manually.
		static std::atomic<bool> requestInitCancellation(false);
		auto isCancellationRequested =
			makeCancellationRequestWithTimeout(1, requestInitCancellation);

		if (!_socket->writeBytes(ss.str(), isCancellationRequested))
		{
			return WebSocketInitResult(false, 500, "Timed out while sending error response");
		}

		return WebSocketInitResult(false, code, reason);
	}

	WebSocketInitResult WebSocketHandshake::clientHandshake(
		const std::string& url,
		const WebSocketHttpHeaders& extraHeaders,
		const std::string& protocol,
		const std::string& host,
		const std::string& path,
		int port,
		int timeoutSecs)
	{
		_requestInitCancellation = false;

		auto isCancellationRequested =
			makeCancellationRequestWithTimeout(timeoutSecs, _requestInitCancellation);

		std::string errMsg;
		bool success = _socket->connect(host, port, errMsg, isCancellationRequested);
		if (!success)
		{
			std::stringstream ss;
			ss << "Unable to connect to " << host << " on port " << port << ", error: " << errMsg;
			return WebSocketInitResult(false, 0, ss.str());
		}

		// Generate a random 16 bytes string and base64 encode it.
		//
		// See https://stackoverflow.com/questions/18265128/what-is-sec-websocket-key-for
		std::string secWebSocketKey = macaron::Base64::Encode(genRandomString(16));

		// For IPv6 addresses, brackets are required in Host and Origin headers (RFC 7230)
		bool isIPv6Host = host.find(':') != std::string::npos;
		std::string bracketedHost = isIPv6Host ? "[" + host + "]" : host;

		std::stringstream ss;
		ss << "GET " << path << " HTTP/1.1\r\n";
		if (extraHeaders.find("Host") == extraHeaders.end())
		{
			ss << "Host: " << bracketedHost << ":" << port << "\r\n";
		}
		ss << "Upgrade: websocket\r\n";
		ss << "Connection: Upgrade\r\n";
		ss << "Sec-WebSocket-Version: 13\r\n";
		ss << "Sec-WebSocket-Key: " << secWebSocketKey << "\r\n";

		// User-Agent can be customized by users
		if (extraHeaders.find("User-Agent") == extraHeaders.end())
		{
			ss << "User-Agent: " << userAgent() << "\r\n";
		}

		// Set an origin header if missing
		if (extraHeaders.find("Origin") == extraHeaders.end())
		{
			ss << "Origin: " << protocol << "://" << bracketedHost << ":" << port << "\r\n";
		}

		for (auto& it : extraHeaders)
		{
			ss << it.first << ": " << it.second << "\r\n";
		}

		if (_enablePerMessageDeflate)
		{
			ss << _perMessageDeflateOptions.generateHeader();
		}

		ss << "\r\n";

		if (!_socket->writeBytes(ss.str(), isCancellationRequested))
		{
			return WebSocketInitResult(
				false, 0, std::string("Failed sending GET request to ") + url);
		}

		// Read HTTP status line
		auto lineResult = _socket->readLine(isCancellationRequested);
		auto lineValid = lineResult.first;
		auto line = lineResult.second;

		if (!lineValid)
		{
			return WebSocketInitResult(
				false, 0, std::string("Failed reading HTTP status line from ") + url);
		}

		// Validate status
		auto statusLine = Http::parseStatusLine(line);
		std::string httpVersion = statusLine.first;
		int status = statusLine.second;

		// HTTP/1.0 is too old.
		if (httpVersion != "HTTP/1.1")
		{
			std::stringstream ss;
			ss << "Expecting HTTP/1.1, got " << httpVersion << ". "
			   << "Rejecting connection to " << url << ", status: " << status
			   << ", HTTP Status line: " << line;
			return WebSocketInitResult(false, status, ss.str());
		}

		auto result = parseHttpHeaders(_socket, isCancellationRequested);
		auto headersValid = result.first;
		auto headers = result.second;

		if (!headersValid)
		{
			return WebSocketInitResult(false, status, "Error parsing HTTP headers");
		}

		// We want an 101 HTTP status for websocket, otherwise it could be
		// a redirection (like 301)
		if (status != 101)
		{
			std::stringstream ss;
			ss << "Expecting status 101 (Switching Protocol), got " << status
			   << " status connecting to " << url << ", HTTP Status line: " << line;

			return WebSocketInitResult(false, status, ss.str(), headers, path);
		}

		// Check the presence of the connection field
		if (headers.find("connection") == headers.end())
		{
			std::string errorMsg("Missing connection value");
			return WebSocketInitResult(false, status, errorMsg);
		}

		// Check the value of the connection field
		// Some websocket servers (Go/Gorilla?) send lowercase values for the
		// connection header, so do a case insensitive comparison
		//
		// See https://github.com/apache/thrift/commit/7c4bdf9914fcba6c89e0f69ae48b9675578f084a
		//
		if (!insensitiveStringCompare(headers["connection"], "Upgrade"))
		{
			std::stringstream ss;
			ss << "Invalid connection value: " << headers["connection"];
			return WebSocketInitResult(false, status, ss.str());
		}

		char output[29] = {};
		WebSocketHandshakeKeyGen::generate(secWebSocketKey, output);
		if (std::string(output) != headers["sec-websocket-accept"])
		{
			std::string errorMsg("Invalid Sec-WebSocket-Accept value");
			return WebSocketInitResult(false, status, errorMsg);
		}

		if (_enablePerMessageDeflate)
		{
			// Parse the server response. Does it support deflate ?
			std::string header = headers["sec-websocket-extensions"];
			WebSocketPerMessageDeflateOptions webSocketPerMessageDeflateOptions(header);

			// If the server does not support that extension, disable it.
			if (!webSocketPerMessageDeflateOptions.enabled())
			{
				_enablePerMessageDeflate = false;
			}
			// Otherwise try to initialize the deflate engine (zlib)
			else if (!_perMessageDeflate->init(webSocketPerMessageDeflateOptions))
			{
				return WebSocketInitResult(
					false, 0, "Failed to initialize per message deflate engine");
			}
		}

		return WebSocketInitResult(true, status, "", headers, path);
	}

	WebSocketInitResult WebSocketHandshake::serverHandshake(
		int timeoutSecs,
		bool enablePerMessageDeflate,
		const ProtocolSelectionCallback& protocolSelectionCallback,
		HttpRequestPtr request)
	{
		_requestInitCancellation = false;

		auto isCancellationRequested =
			makeCancellationRequestWithTimeout(timeoutSecs, _requestInitCancellation);

		std::string method;
		std::string uri;
		std::string httpVersion;

		if (request)
		{
			method = request->method;
			uri = request->uri;
			httpVersion = request->version;
		}
		else
		{
			// Read first line
			auto lineResult = _socket->readLine(isCancellationRequested);
			auto lineValid = lineResult.first;
			auto line = lineResult.second;

			if (!lineValid)
			{
				return sendErrorResponse(400, "Error reading HTTP request line");
			}

			// Validate request line (GET /foo HTTP/1.1\r\n)
			auto requestLine = Http::parseRequestLine(line);
			method = std::get<0>(requestLine);
			uri = std::get<1>(requestLine);
			httpVersion = std::get<2>(requestLine);
		}

		if (method != "GET")
		{
			return sendErrorResponse(400, "Invalid HTTP method, need GET, got " + method);
		}

		if (httpVersion != "HTTP/1.1")
		{
			return sendErrorResponse(400,
									 "Invalid HTTP version, need HTTP/1.1, got: " + httpVersion);
		}

		WebSocketHttpHeaders headers;
		if (request)
		{
			headers = request->headers;
		}
		else
		{
			// Retrieve and validate HTTP headers
			auto result = parseHttpHeaders(_socket, isCancellationRequested);
			auto headersValid = result.first;
			headers = result.second;

			if (!headersValid)
			{
				return sendErrorResponse(400, "Error parsing HTTP headers");
			}
		}

		if (headers.find("sec-websocket-key") == headers.end())
		{
			return sendErrorResponse(400, "Missing Sec-WebSocket-Key value");
		}

		if (headers.find("upgrade") == headers.end())
		{
			return sendErrorResponse(400, "Missing Upgrade header");
		}

		if (!insensitiveStringCompare(headers["upgrade"], "WebSocket") &&
			headers["Upgrade"] != "keep-alive, Upgrade") // special case for firefox
		{
			return sendErrorResponse(400,
									 "Invalid Upgrade header, "
									 "need WebSocket, got " +
										 headers["upgrade"]);
		}

		if (headers.find("sec-websocket-version") == headers.end())
		{
			return sendErrorResponse(400, "Missing Sec-WebSocket-Version value");
		}

		{
			std::stringstream ss;
			ss << headers["sec-websocket-version"];
			int version;
			ss >> version;

			if (version != 13)
			{
				return sendErrorResponse(400,
										 "Invalid Sec-WebSocket-Version, "
										 "need 13, got " +
											 ss.str());
			}
		}

		char output[29] = {};
		WebSocketHandshakeKeyGen::generate(headers["sec-websocket-key"], output);

		// Parse sub-protocols requested by the client
		std::vector<std::string> requestedProtocols;
		if (headers.find("sec-websocket-protocol") != headers.end())
		{
			std::istringstream protocolStream(headers["sec-websocket-protocol"]);
			std::string token;
			while (std::getline(protocolStream, token, ','))
			{
				size_t start = token.find_first_not_of(" \t");
				size_t end   = token.find_last_not_of(" \t");
				if (start != std::string::npos)
				{
					requestedProtocols.push_back(token.substr(start, end - start + 1));
				}
			}
		}

		// Invoke the server's protocol selection callback (if any)
		std::string selectedProtocol;
		if (protocolSelectionCallback)
		{
			if (!protocolSelectionCallback(requestedProtocols, selectedProtocol))
			{
				return sendErrorResponse(400, "Unsupported protocol");
			}
		}

		std::stringstream ss;
		ss << "HTTP/1.1 101 Switching Protocols\r\n";
		ss << "Sec-WebSocket-Accept: " << std::string(output) << "\r\n";
		ss << "Upgrade: websocket\r\n";
		ss << "Connection: Upgrade\r\n";
		ss << "Server: " << (_serverName.empty() ? userAgent() : _serverName) << "\r\n";

		if (!selectedProtocol.empty())
		{
			ss << "Sec-WebSocket-Protocol: " << selectedProtocol << "\r\n";
		}

		// Parse the client headers. Does it support deflate ?
		std::string header = headers["sec-websocket-extensions"];
		WebSocketPerMessageDeflateOptions webSocketPerMessageDeflateOptions(header);

		// If the client has requested that extension,
		if (webSocketPerMessageDeflateOptions.enabled() && enablePerMessageDeflate)
		{
			_enablePerMessageDeflate = true;

			if (!_perMessageDeflate->init(webSocketPerMessageDeflateOptions))
			{
				return WebSocketInitResult(
					false, 0, "Failed to initialize per message deflate engine");
			}
			ss << webSocketPerMessageDeflateOptions.generateHeader();
		}

		ss << "\r\n";

		if (!_socket->writeBytes(ss.str(), isCancellationRequested))
		{
			return WebSocketInitResult(
				false, 0, std::string("Failed sending response to remote end"));
		}

		auto result = WebSocketInitResult(true, 200, "", headers, uri);
		result.protocol = selectedProtocol;
		return result;
	}
}
