#include "pch.h"
#include "cvar.h"
#include "limits.h"
#include "dedicated.h"
#include "tier0.h"
#include "r2engine.h"
#include "r2client.h"

AUTOHOOK_INIT()

ConVar* ns_exploitfixes_log;
#define SHOULD_LOG (ns_exploitfixes_log->m_Value.m_nValue > 0)
#define BLOCKED_INFO(s)                                                                                                                    \
	(                                                                                                                                      \
		[=]() -> bool                                                                                                                      \
		{                                                                                                                                  \
			if (SHOULD_LOG)                                                                                                                \
			{                                                                                                                              \
				std::stringstream stream;                                                                                                  \
				stream << "ExploitFixes.cpp: " << BLOCK_PREFIX << s;                                                                       \
				spdlog::error(stream.str());                                                                                               \
			}                                                                                                                              \
			return false;                                                                                                                  \
		}())

// Make sure 3 or less floats are valid
bool ValidateFloats(float a, float b = 0, float c = 0)
{
	return !isnan(a) && !isnan(b) && !isnan(c);
}

struct Vector
{
	float x, y, z;

	Vector(float x = 0, float y = 0, float z = 0) : x(x), y(y), z(z) {}

	bool IsValid()
	{
		return ValidateFloats(x, y, z);
	}
};

struct Angle
{
	float pitch, yaw, roll;

	Angle(float pitch = 0, float yaw = 0, float roll = 0) : pitch(pitch), yaw(yaw), roll(roll) {}

	bool IsInvalid()
	{
		return !ValidateFloats(pitch, yaw, roll);

		if (!ValidateFloats(pitch, yaw, roll))
			return false;

		return (pitch > 90 || pitch < -90) || (yaw > 180 || yaw < -180) || (roll > 180 || roll < -180);
	}
};

// block bad netmessages
// Servers can literally request a screenshot from any client, yeah no
AUTOHOOK(CLC_Screenshot_WriteToBuffer, engine.dll + 0x22AF20, 
bool, __fastcall, (void* thisptr, void* buffer)) // 48 89 5C 24 ? 57 48 83 EC 20 8B 42 10
{
	return false;
}

AUTOHOOK(CLC_Screenshot_ReadFromBuffer, engine.dll + 0x221F00, 
bool, __fastcall, (void* thisptr, void* buffer)) // 48 89 5C 24 ? 48 89 6C 24 ? 48 89 74 24 ? 57 48 83 EC 20 48 8B DA 48 8B 52 38
{
	return false;
}

// This is unused ingame and a big client=>server=>client exploit vector
AUTOHOOK(Base_CmdKeyValues_ReadFromBuffer, engine.dll + 0x220040,
bool, __fastcall, (void* thisptr, void* buffer)) // 40 55 48 81 EC ? ? ? ? 48 8D 6C 24 ? 48 89 5D 70
{
	return false;
}

AUTOHOOK(CClient_ProcessSetConVar, engine.dll + 0x75CF0,
bool, __fastcall, (void* pMsg)) // 48 8B D1 48 8B 49 18 48 8B 01 48 FF 60 10
{

	constexpr int ENTRY_STR_LEN = 260;
	struct SetConVarEntry
	{
		char name[ENTRY_STR_LEN];
		char val[ENTRY_STR_LEN];
	};

	struct NET_SetConVar
	{
		void* vtable;
		void* unk1;
		void* unk2;
		void* m_pMessageHandler;
		SetConVarEntry* m_ConVars; // convar entry array
		void* unk5; // these 2 unks are just vector capacity or whatever
		void* unk6;
		int m_ConVars_count; // amount of cvar entries in array (this will not be out of bounds)
	};

	auto msg = (NET_SetConVar*)pMsg;
	bool bIsServerFrame = Tier0::ThreadInServerFrameThread();

	std::string BLOCK_PREFIX = std::string {"NET_SetConVar ("} + (bIsServerFrame ? "server" : "client") + "): Blocked dangerous/invalid msg: ";

	if (bIsServerFrame)
	{
		constexpr int SETCONVAR_SANITY_AMOUNT_LIMIT = 69;
		if (msg->m_ConVars_count < 1 || msg->m_ConVars_count > SETCONVAR_SANITY_AMOUNT_LIMIT)
		{
			return BLOCKED_INFO("Invalid m_ConVars_count (" << msg->m_ConVars_count << ")");
		}
	}

	for (int i = 0; i < msg->m_ConVars_count; i++)
	{
		auto entry = msg->m_ConVars + i;

		// Safety check for memory access
		if (MemoryAddress(entry).IsMemoryReadable(sizeof(*entry)))
		{

			// Find null terminators
			bool nameValid = false, valValid = false;
			for (int i = 0; i < ENTRY_STR_LEN; i++)
			{
				if (!entry->name[i])
					nameValid = true;
				if (!entry->val[i])
					valValid = true;
			}

			if (!nameValid || !valValid)
				return BLOCKED_INFO("Missing null terminators");

			auto realVar = R2::g_pCVar->FindVar(entry->name);

			if (realVar)
				memcpy(
					entry->name,
					realVar->m_ConCommandBase.m_pszName,
					strlen(realVar->m_ConCommandBase.m_pszName) + 1); // Force name to match case

			bool isValidFlags = true;
			if (bIsServerFrame)
			{
				if (realVar)
					isValidFlags = realVar->IsFlagSet(FCVAR_USERINFO); // ConVar MUST be userinfo var
			}
			else
			{
				// TODO: Should probably have some sanity checks, but can't find any that are consistent
			}

			if (!isValidFlags)
			{
				if (!realVar)
				{
					return BLOCKED_INFO("Invalid flags on nonexistant cvar (how tho???)");
				}
				else
				{
					return BLOCKED_INFO(
						"Invalid flags (" << std::hex << "0x" << realVar->m_ConCommandBase.m_nFlags << "), var is " << entry->name);
				}
			}
		}
		else
		{
			return BLOCKED_INFO("Unreadable memory at " << (void*)entry); // Not risking that one, they all gotta be readable
		}
	}

	return CClient_ProcessSetConVar(msg);
}

// prevent invalid user CMDs
AUTOHOOK(CClient_ProcessUsercmds, engine.dll + 0x1040F0,
bool, __fastcall, (void* thisptr, void* pMsg)) // 40 55 56 48 83 EC 58
{
	struct CLC_Move
	{
		BYTE gap0[24];
		void* m_pMessageHandler;
		int m_nBackupCommands;
		int m_nNewCommands;
		int m_nLength;
		// bf_read m_DataIn;
		// bf_write m_DataOut;
	};

	auto msg = (CLC_Move*)pMsg;

	const char* BLOCK_PREFIX = "ProcessUserCmds: ";

	if (msg->m_nBackupCommands < 0)
	{
		return BLOCKED_INFO("Invalid m_nBackupCommands (" << msg->m_nBackupCommands << ")");
	}

	if (msg->m_nNewCommands < 0)
	{
		return BLOCKED_INFO("Invalid m_nNewCommands (" << msg->m_nNewCommands << ")");
	}

	// removing, as vanilla already limits num usercmds per frame
	/*constexpr int NUMCMD_SANITY_LIMIT = 16;
	if ((msg->m_nNewCommands + msg->m_nBackupCommands) > NUMCMD_SANITY_LIMIT)
	{
		return BLOCKED_INFO("Command count is too high (new: " << msg->m_nNewCommands << ", backup: " << msg->m_nBackupCommands << ")");

	}*/

	if (msg->m_nLength <= 0)
		return BLOCKED_INFO("Invalid message length (" << msg->m_nLength << ")");

	return CClient_ProcessUsercmds(thisptr, pMsg);
}

AUTOHOOK(ReadUsercmd, server.dll + 0x2603F0,
void, __fastcall, (void* buf, void* pCmd_move, void* pCmd_from)) // 4C 89 44 24 ? 53 55 56 57
{
	// Let normal usercmd read happen first, it's safe
	ReadUsercmd(buf, pCmd_move, pCmd_from);

	// Now let's make sure the CMD we read isnt messed up to prevent numerous exploits (including server crashing)
	struct __declspec(align(4)) SV_CUserCmd
	{
		DWORD command_number;
		DWORD tick_count;
		float command_time;
		Angle worldViewAngles;
		BYTE gap18[4];
		Angle localViewAngles;
		Angle attackangles;
		Vector move;
		DWORD buttons;
		BYTE impulse;
		short weaponselect;
		DWORD meleetarget;
		BYTE gap4C[24];
		char headoffset;
		BYTE gap65[11];
		Vector cameraPos;
		Angle cameraAngles;
		BYTE gap88[4];
		int tickSomething;
		DWORD dword90;
		DWORD predictedServerEventAck;
		DWORD dword98;
		float frameTime;
	};

	auto cmd = (SV_CUserCmd*)pCmd_move;
	auto fromCmd = (SV_CUserCmd*)pCmd_from;

	std::string BLOCK_PREFIX =
		"ReadUsercmd (command_number delta: " + std::to_string(cmd->command_number - fromCmd->command_number) + "): ";

	if (cmd->worldViewAngles.IsInvalid())
	{
		BLOCKED_INFO("CMD has invalid worldViewAngles");
		goto INVALID_CMD;
	}

	if (cmd->attackangles.IsInvalid())
	{
		BLOCKED_INFO("CMD has invalid attackangles");
		goto INVALID_CMD;
	}

	if (cmd->localViewAngles.IsInvalid())
	{
		BLOCKED_INFO("CMD has invalid localViewAngles");
		goto INVALID_CMD;
	}

	if (cmd->cameraAngles.IsInvalid())
	{
		BLOCKED_INFO("CMD has invalid cameraAngles");
		goto INVALID_CMD;
	}

	if (cmd->frameTime <= 0 || cmd->tick_count == 0 || cmd->command_time <= 0)
	{
		BLOCKED_INFO(
			"Bogus cmd timing (tick_count: " << cmd->tick_count << ", frameTime: " << cmd->frameTime
											 << ", commandTime : " << cmd->command_time << ")");
		goto INVALID_CMD; // No simulation of bogus-timed cmds
	}

	if (!cmd->move.IsValid())
	{
		BLOCKED_INFO("Invalid move vector");
		goto INVALID_CMD;
	}

	if (!cmd->cameraPos.IsValid())
	{
		BLOCKED_INFO("Invalid cameraPos"); // IIRC this can crash spectating clients or anyone watching replays
		goto INVALID_CMD;
	}

	return;

INVALID_CMD:

	// Fix any gameplay-affecting cmd properties
	// NOTE: Currently tickcount/frametime is set to 0, this ~shouldn't~ cause any problems
	cmd->worldViewAngles = cmd->localViewAngles = cmd->attackangles = cmd->cameraAngles = Angle(0, 0, 0);
	cmd->tick_count = cmd->frameTime = 0;
	cmd->move = cmd->cameraPos = Vector(0, 0, 0);
	cmd->buttons = 0;
	cmd->meleetarget = 0;
}

// ensure that GetLocalBaseClient().m_bRestrictServerCommands is set correctly, which the return value of this function controls
// this is IsValveMod in source, but we're making it IsRespawnMod now since valve didn't make this one
AUTOHOOK(IsRespawnMod, engine.dll + 0x1C6360,
bool, __fastcall, (const char* pModName)) // 48 83 EC 28 48 8B 0D ? ? ? ? 48 8D 15 ? ? ? ? E8 ? ? ? ? 85 C0 74 63
{
	return (!strcmp("r2", pModName) || !strcmp("r1", pModName)) 
		&& !Tier0::CommandLine()->CheckParm("-norestrictservercommands");
}

// ratelimit stringcmds, and prevent remote clients from calling non-FCVAR_GAMEDLL_FOR_REMOTE_CLIENTS commands
bool (*CCommand__Tokenize)(CCommand& self, const char* pCommandString, R2::cmd_source_t commandSource);
AUTOHOOK(CGameClient__ExecuteStringCommand, 
engine.dll + 0x1022E0, bool, , (R2::CBaseClient* self, uint32_t unknown, const char* pCommandString))
{
	if (!g_pServerLimits->CheckStringCommandLimits(self))
	{
		R2::CBaseClient__Disconnect(self, 1, "Sent too many stringcmd commands");
		return false;
	}

	// verify the command we're trying to execute is FCVAR_GAMEDLL_FOR_REMOTE_CLIENTS, if it's a concommand
	char* commandBuf[1040]; // assumedly this is the size of CCommand since we don't have an actual constructor
	memset(commandBuf, 0, sizeof(commandBuf));
	CCommand tempCommand = *(CCommand*)&commandBuf;

	if (!CCommand__Tokenize(tempCommand, pCommandString, R2::cmd_source_t::kCommandSrcCode) || !tempCommand.ArgC())
		return false;

	ConCommand* command = R2::g_pCVar->FindCommand(tempCommand.Arg(0));

	// if the command doesn't exist pass it on to ExecuteStringCommand for script clientcommands and stuff
	if (command && !command->IsFlagSet(FCVAR_GAMEDLL_FOR_REMOTE_CLIENTS))
	{
		// ensure FCVAR_GAMEDLL concommands without FCVAR_GAMEDLL_FOR_REMOTE_CLIENTS can't be executed by remote clients
		if (IsDedicatedServer())
			return false;

		if (strcmp(self->m_UID, R2::g_pLocalPlayerUserID))
			return false;
	}

	return CGameClient__ExecuteStringCommand(self, unknown, pCommandString);
}

// prevent clients from crashing servers through overflowing CNetworkStringTableContainer::WriteBaselines
bool bWasWritingStringTableSuccessful;
AUTOHOOK(CBaseClient__SendServerInfo, 
engine.dll + 0x104FB0, void, , (void* self))
{
	bWasWritingStringTableSuccessful = true;
	CBaseClient__SendServerInfo(self);
	if (!bWasWritingStringTableSuccessful)
		R2::CBaseClient__Disconnect(
			self, 1, "Overflowed CNetworkStringTableContainer::WriteBaselines, try restarting your client and reconnecting");
}

ON_DLL_LOAD("engine.dll", EngineExploitFixes, (CModule module))
{
	AUTOHOOK_DISPATCH_MODULE(engine.dll)

	CCommand__Tokenize = module.Offset(0x418380).As<bool(*)(CCommand&, const char*, R2::cmd_source_t)>();

	// allow client/ui to run clientcommands despite restricting servercommands
	module.Offset(0x4FB65).Patch("EB 11");
	module.Offset(0x4FBAC).Patch("EB 16");

	// patch to set bWasWritingStringTableSuccessful in CNetworkStringTableContainer::WriteBaselines if it fails
	{
		MemoryAddress writeAddress(&bWasWritingStringTableSuccessful - module.Offset(0x234EDC).m_nAddress);

		MemoryAddress addr = module.Offset(0x234ED2);
		addr.Patch("C7 05");
		addr.Offset(2).Patch((BYTE*)&writeAddress, sizeof(writeAddress));

		addr.Offset(6).Patch("00 00 00 00");

		addr.Offset(10).NOP(5);
	}
}

ON_DLL_LOAD_RELIESON("server.dll", ServerExploitFixes, ConVar, (CModule module))
{
	AUTOHOOK_DISPATCH_MODULE(server.dll)

	// ret at the start of CServerGameClients::ClientCommandKeyValues as it has no benefit and is forwarded to client (i.e. security issue)
	// this prevents the attack vector of client=>server=>client, however server=>client also has clientside patches
	module.Offset(0x153920).Patch("C3");

	// Dumb ANTITAMPER patches (they negatively impact performance and security)
	constexpr const char* ANTITAMPER_EXPORTS[] = {
		"ANTITAMPER_SPOTCHECK_CODEMARKER",
		"ANTITAMPER_TESTVALUE_CODEMARKER",
		"ANTITAMPER_TRIGGER_CODEMARKER",
	};

	// Prevent these from actually doing anything
	for (auto exportName : ANTITAMPER_EXPORTS)
	{
		MemoryAddress exportAddr = module.GetExport(exportName);
		if (exportAddr)
		{
			// Just return, none of them have any args or are userpurge
			exportAddr.Patch("C3");
			spdlog::info("Patched AntiTamper function export \"{}\"", exportName);
		}
	}

	ns_exploitfixes_log =
		new ConVar("ns_exploitfixes_log", "1", FCVAR_GAMEDLL, "Whether to log whenever ExploitFixes.cpp blocks/corrects something");
}