siadrive/UnitTests/MockSiad.cpp

#include "stdafx.h"
#include "MockSiad.h"
#include <future>
#include <SiaCommon.h>
using namespace Sia::Api;

class CMockClient
{
private:
	const std::string SUCESS_RESPONSE_JSON = "HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Length: {ContentLength}\r\nContent-Type: application/json\r\n\r\n{Content}";
	const std::string NOT_FOUND_RESPONSE = "HTTP/1.1 404 Not Found\r\nConnection: close\r\nContent-Length: 17\r\nContent-Type: application/json\r\n\r\n{\"message\":\"404\"}";
	const std::string SERVER_VERSION_SUCCESS_CONTENT = "{\"version\": \"1.1.0\"}";

public:
	CMockClient(SOCKET socket, std::function<void(CMockClient* client)> removedCallback) :
		_socket(socket),
		_stopEvent(::CreateEvent(nullptr, FALSE, FALSE, nullptr))
	{
	}

public:
	~CMockClient()
	{
		::SetEvent(_stopEvent);
		_clientThread->join();
		::CloseHandle(_stopEvent);
	}

private:
	SOCKET _socket;
	HANDLE _stopEvent;
	std::function<void(CMockClient* client)> NotifyClientRemoved;
	std::unique_ptr<std::thread> _clientThread;

private:
	void Run()
	{
		bool connected = true;
		while (connected && (::WaitForSingleObject(_stopEvent, 1) == WAIT_TIMEOUT))
		{
			timeval tv = { 1, 0 };
			FD_SET readFds, exceptFds;
			FD_ZERO(&readFds);
			FD_ZERO(&exceptFds);
			FD_SET(_socket, &readFds);
			FD_SET(_socket, &exceptFds);
			int selectResult = ::select(0, &readFds, nullptr, &exceptFds, &tv);
			if (selectResult != SOCKET_ERROR)
			{
				if (FD_ISSET(_socket, &readFds))
				{
					std::vector<char> buffer(1024);
					int recvd;
					std::string response;
					{
						std::string tmp;
						while (!response.length() && ((recvd = ::recv(_socket, &buffer[0], buffer.size(), 0))))
						{
							tmp += std::string(&buffer[0], recvd);
							int position = tmp.find_first_of("\r\n\r\n");
							if (position > -1)
							{
								std::string header = tmp.substr(0, position);
								if (header == "GET /daemon/version HTTP/1.1")
								{
									response = SUCESS_RESPONSE_JSON;
									response = ReplaceStringInPlace(ReplaceStringInPlace(response, "{ContentLength}", std::to_string(SERVER_VERSION_SUCCESS_CONTENT.length())), "{Content}", SERVER_VERSION_SUCCESS_CONTENT);
								}
								else
								{
									response = NOT_FOUND_RESPONSE;
								}
							}
						}

					}
					
					if (response.size())
					{
						::send(_socket, &response[0], response.length(), 0);
					}

				}
				else if (FD_ISSET(_socket, &exceptFds))
				{
					connected = false;
				}
			}
		}

		closesocket(_socket);
		NotifyClientRemoved(this);
	}

public:
	void Start()
	{
		_clientThread.reset(new std::thread([this]()
		{
			Run();
		}));
	}

	void Close()
	{
		::SetEvent(_stopEvent);
		_clientThread->join();
	}
};


CMockSiad::CMockSiad(const SiaHostConfig& hostConfig) :
	_hostConfig(hostConfig),
	_stopEvent(::CreateEvent(nullptr, FALSE, FALSE, nullptr))
{
	const WORD ver = MAKEWORD(2, 2);
	WSADATA d = { 0 };
	::WSAStartup(ver, &d);
}

CMockSiad::~CMockSiad()
{
	Stop();
	CloseHandle(_stopEvent);
	::WSACleanup();
}

void CMockSiad::Start(const SiadTestType& testType)
{
	if (!_serverThread)
	{
		HANDLE startedEvent = ::CreateEvent(nullptr, false, false, nullptr);

		_serverThread.reset(new std::thread([this,&startedEvent]()
		{
			bool started = false;
			bool active;
			std::mutex clientMutex;
			std::vector<CMockClient*> connected;
			std::vector<CMockClient*> disconnected;
			auto ClearDisconnected = [&]()
			{
				std::lock_guard<std::mutex> l(clientMutex);
				while (disconnected.size())
				{
					CMockClient* client = disconnected.back();
					disconnected.pop_back();
					delete client;
				}
			};

			do
			{
				SOCKET listenSocket = ::socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);

				struct sockaddr_in sockAddr = { 0 };
				sockAddr.sin_family = AF_INET;
				sockAddr.sin_addr.s_addr = inet_addr(CW2A(_hostConfig.HostName.c_str()));
				sockAddr.sin_port = htons(_hostConfig.HostPort);

				bool listening = ((listenSocket != INVALID_SOCKET) && (::bind(listenSocket, reinterpret_cast<const struct sockaddr *>(&sockAddr), sizeof(sockAddr)) != SOCKET_ERROR)) && (::listen(listenSocket, 1) != SOCKET_ERROR);
				while (listening && ((active = (::WaitForSingleObject(_stopEvent, 1) == WAIT_TIMEOUT))))
				{
					if (!started)
					{
						::SetEvent(startedEvent);
						started = true;
					}

					ClearDisconnected();

					timeval tv = { 1, 0 };
					FD_SET readFds, exceptFds;
					FD_ZERO(&readFds);
					FD_ZERO(&exceptFds);
					FD_SET(listenSocket, &readFds);
					FD_SET(listenSocket, &exceptFds);
					int selectResult = ::select(0, &readFds, nullptr, &exceptFds, &tv);
					if (selectResult != SOCKET_ERROR)
					{
						if (FD_ISSET(listenSocket, &readFds))
						{
							std::lock_guard<std::mutex> l(clientMutex);
							SOCKET clientSocket = ::accept(listenSocket, nullptr, nullptr);
							CMockClient* client = new CMockClient(clientSocket, [this, &clientMutex, &connected, &disconnected](CMockClient* removedClient)
							{
								std::lock_guard<std::mutex> l(clientMutex);
								connected.erase(std::remove_if(connected.begin(), connected.end(), [&](CMockClient* client) -> bool { return client == removedClient; }), connected.end());
								disconnected.push_back(removedClient);
							});
							client->Start();
						}
						else if (FD_ISSET(listenSocket, &exceptFds))
						{
							listening = false;
						}
					}
				}

				closesocket(listenSocket);
				while (connected.size())
				{
					CMockClient* client = nullptr;
					{
						std::lock_guard<std::mutex> l(clientMutex);
						if (connected.size())
						{
							client = connected.back();
						}
					}

					if (client)
					{
						client->Close();
					}
				}
			} while (active);

			ClearDisconnected();
		}));

		::WaitForSingleObject(startedEvent, INFINITE);
		::CloseHandle(startedEvent);
	}
}

void CMockSiad::Stop()
{
	if (_serverThread)
	{
		::SetEvent(_stopEvent);
		_serverThread->join();
		_serverThread.reset(nullptr);
	}
}