cs2-sendproxy/src/overrides.cpp

#include "overrides.h"

#include "common.h"
#include "schema.h"
#include "entity2/entitysystem.h"
#include "entity2/entityinstance.h"

#include <cstring>
#include <algorithm>

extern CGameEntitySystem* g_entitySystem;
extern CGameEntitySystem* GameEntitySystem();

OverrideManager g_overrides;

static int EntityIndex(CEntityIdentity* identity)
{
	return identity ? identity->m_EHandle.GetEntryIndex() : -1;
}

static CEntityIdentity* FindIdentityByIndex(int index)
{
	if (!g_entitySystem)
		g_entitySystem = GameEntitySystem();
	if (!g_entitySystem)
		return nullptr;

	int visited = 0;
	for (auto* identity = g_entitySystem->m_EntityList.m_pFirstActiveEntity; identity && visited < 16384; identity = identity->m_pNext, ++visited)
	{
		if (EntityIndex(identity) == index)
			return identity;
	}
	return nullptr;
}

static std::vector<std::string> SplitString(const char* value, char separator)
{
	std::vector<std::string> out;
	if (!value)
		return out;

	const char* start = value;
	for (const char* p = value; ; ++p)
	{
		if (*p != separator && *p != '\0')
			continue;
		out.emplace_back(start, p - start);
		if (*p == '\0')
			break;
		start = p + 1;
	}
	return out;
}

static std::string NormalizeSchemaClassName(std::string typeName)
{
	while (!typeName.empty() && (typeName.back() == '*' || typeName.back() == '&' || typeName.back() == ' '))
		typeName.pop_back();
	while (!typeName.empty() && typeName.front() == ' ')
		typeName.erase(typeName.begin());

	constexpr const char* classPrefix = "class ";
	constexpr const char* structPrefix = "struct ";
	if (typeName.rfind(classPrefix, 0) == 0)
		typeName.erase(0, strlen(classPrefix));
	else if (typeName.rfind(structPrefix, 0) == 0)
		typeName.erase(0, strlen(structPrefix));

	return typeName;
}

static bool IsPointerSchemaType(const std::string& typeName)
{
	for (auto it = typeName.rbegin(); it != typeName.rend(); ++it)
	{
		if (*it == ' ')
			continue;
		return *it == '*';
	}
	return false;
}

void OverrideManager::Clear()
{
	for (const auto& rule : m_rules)
	{
		if (rule.enabled)
			MarkEntityFieldDirtyByOffset(rule.entityIndex, rule.offset);
	}

	m_rules.clear();
	m_rulesByEntity.clear();
}

void OverrideManager::Compact()
{
	m_rules.erase(
		std::remove_if(m_rules.begin(), m_rules.end(), [](const CompiledOverrideRule& rule) {
			return !rule.enabled;
		}),
		m_rules.end());
	RebuildIndex();
}

bool OverrideManager::ResolveFieldPath(const char* className, const char* fieldPath, int32_t& offset, int32_t& size, std::vector<FieldAddressStep>& addressSteps, std::string& error) const
{
	if (!className || !className[0] || !fieldPath || !fieldPath[0])
	{
		error = "class and field path are required";
		return false;
	}

	auto parts = SplitString(fieldPath, '.');
	if (parts.empty())
	{
		error = "field path is empty";
		return false;
	}

	offset = 0;
	size = 0;
	addressSteps.clear();
	std::string currentClass = className;
	for (size_t i = 0; i < parts.size(); ++i)
	{
		auto field = schema::FindField(currentClass.c_str(), parts[i].c_str());
		if (!field.found)
		{
			error = "schema field not found: " + currentClass + "::" + parts[i];
			return false;
		}

		size = field.size;
		FieldAddressStep step {};
		step.offset = field.offset;
		step.dereference = i + 1 < parts.size() && IsPointerSchemaType(field.typeName);
		addressSteps.push_back(step);
		if (i + 1 < parts.size())
		{
			if (field.typeName.empty())
			{
				error = "intermediate field has no schema type name";
				return false;
			}
			currentClass = NormalizeSchemaClassName(field.typeName);
		}
	}

	if (!addressSteps.empty())
		offset = addressSteps.front().offset;

	return true;
}

static uint8_t* ResolveRuleAddress(const CompiledOverrideRule& rule, void* entityData)
{
	auto* address = static_cast<uint8_t*>(entityData);
	for (const auto& step : rule.addressSteps)
	{
		address += step.offset;
		if (!step.dereference)
			continue;

		address = *reinterpret_cast<uint8_t**>(address);
		if (!address)
			return nullptr;
	}
	return address;
}

static uint8_t* ResolveAddressSteps(const std::vector<FieldAddressStep>& steps, void* entityData)
{
	auto* address = static_cast<uint8_t*>(entityData);
	for (const auto& step : steps)
	{
		address += step.offset;
		if (!step.dereference)
			continue;

		address = *reinterpret_cast<uint8_t**>(address);
		if (!address)
			return nullptr;
	}
	return address;
}

static void SaveAndWrite(std::vector<AppliedOverride>& applied, uint8_t* address, const void* value, size_t size)
{
	AppliedOverride saved {};
	saved.address = address;
	saved.original.resize(size);
	memcpy(saved.original.data(), address, saved.original.size());
	memcpy(address, value, size);
	applied.push_back(saved);
}

static void SaveAndWriteIfChanged(std::vector<AppliedOverride>& applied, uint8_t* address, const void* value, size_t size)
{
	if (memcmp(address, value, size) == 0)
		return;

	SaveAndWrite(applied, address, value, size);
}

static void ApplyVectorFirstFromField(const CompiledOverrideRule& rule, void* entityData, std::vector<AppliedOverride>& applied)
{
	auto* vectorAddress = ResolveAddressSteps(rule.addressSteps, entityData);
	auto* sourceAddress = ResolveAddressSteps(rule.sourceAddressSteps, entityData);
	if (!vectorAddress || !sourceAddress)
		return;

	const uint32_t sourceHandle = *reinterpret_cast<uint32_t*>(sourceAddress);
	constexpr uint32_t invalidHandle = 0xffffffff;
	const int32_t count = *reinterpret_cast<int32_t*>(vectorAddress);
	auto* elements = *reinterpret_cast<uint8_t**>(vectorAddress + 8);
	if (!rule.debugLogged)
	{
		uint32_t h0 = elements && count > 0 ? *reinterpret_cast<uint32_t*>(elements) : invalidHandle;
		uint32_t h1 = elements && count > 1 ? *reinterpret_cast<uint32_t*>(elements + 4) : invalidHandle;
		uint32_t h2 = elements && count > 2 ? *reinterpret_cast<uint32_t*>(elements + 8) : invalidHandle;
		const auto* raw = reinterpret_cast<const uint8_t*>(vectorAddress);
		SPMessage("vector apply rule %d recipient=%d entity=%d vec=%p src=%p count=%d elems=%p active=%08x first=[%08x %08x %08x] raw=%02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x\n",
			rule.id,
			rule.recipientSlot,
			rule.entityIndex,
			vectorAddress,
			sourceAddress,
			count,
			elements,
			sourceHandle,
			h0,
			h1,
			h2,
			raw[0], raw[1], raw[2], raw[3], raw[4], raw[5], raw[6], raw[7],
			raw[8], raw[9], raw[10], raw[11], raw[12], raw[13], raw[14], raw[15],
			raw[16], raw[17], raw[18], raw[19], raw[20], raw[21], raw[22], raw[23]);
		rule.debugLogged = true;
	}
	if (count < 0 || count > 128)
		return;

	auto* countAddress = reinterpret_cast<uint8_t*>(&*reinterpret_cast<int32_t*>(vectorAddress));
	if (sourceHandle == invalidHandle)
	{
		const int32_t emptyCount = 0;
		SaveAndWriteIfChanged(applied, countAddress, &emptyCount, sizeof(emptyCount));
		return;
	}

	if (!elements)
		return;

	const int32_t activeOnlyCount = 1;
	SaveAndWriteIfChanged(applied, countAddress, &activeOnlyCount, sizeof(activeOnlyCount));
	SaveAndWriteIfChanged(applied, elements, &sourceHandle, sizeof(sourceHandle));
}

void OverrideManager::RebuildIndex()
{
	m_rulesByEntity.clear();
	for (size_t i = 0; i < m_rules.size(); ++i)
	{
		if (m_rules[i].enabled)
			m_rulesByEntity[m_rules[i].entityIndex].push_back(i);
	}
}

bool OverrideManager::MarkEntityFieldDirtyByOffset(int entityIndex, int32_t offset)
{
	auto* identity = FindIdentityByIndex(entityIndex);
	if (!identity || !identity->m_pInstance)
		return false;

	NetworkStateChangedData data(static_cast<uint32>(offset));
	identity->m_pInstance->NetworkStateChanged(data);
	return true;
}

bool OverrideManager::MarkEntityFullDirty(int entityIndex)
{
	auto* identity = FindIdentityByIndex(entityIndex);
	if (!identity || !identity->m_pInstance)
		return false;

	NetworkStateChangedData data(true);
	identity->m_pInstance->NetworkStateChanged(data);
	return true;
}

int OverrideManager::AddRuleBytes(int recipientSlot, int entityIndex, const char* className, const char* fieldPath, const void* value, int valueSize, std::string& error)
{
	if (recipientSlot < -1 || recipientSlot >= 64)
	{
		error = "recipient slot must be all/-1 or 0..63";
		return 0;
	}
	if (entityIndex < 0)
	{
		error = "entity index must be >= 0";
		return 0;
	}
	if (!value || valueSize <= 0)
	{
		error = "value bytes are required";
		return 0;
	}

	CompiledOverrideRule rule;
	rule.id = m_nextId++;
	rule.recipientSlot = recipientSlot;
	rule.entityIndex = entityIndex;
	rule.className = className ? className : "";
	rule.fieldPath = fieldPath ? fieldPath : "";
	if (!ResolveFieldPath(rule.className.c_str(), rule.fieldPath.c_str(), rule.offset, rule.size, rule.addressSteps, error))
		return 0;
	if (rule.size > 0 && valueSize > rule.size)
	{
		error = "value is larger than schema field";
		return 0;
	}

	auto* bytes = static_cast<const uint8_t*>(value);
	rule.value.assign(bytes, bytes + valueSize);
	m_rules.push_back(rule);
	if (rule.fieldPath.find("m_hMyWeapons") != std::string::npos)
	{
		SPMessage("rule %d recipient=%d entity=%d %s::%s size=%d valueSize=%d topOffset=%d steps=%zu\n",
			rule.id,
			rule.recipientSlot,
			rule.entityIndex,
			rule.className.c_str(),
			rule.fieldPath.c_str(),
			rule.size,
			valueSize,
			rule.offset,
			rule.addressSteps.size());
	}
	RebuildIndex();
	MarkEntityFullDirty(entityIndex);
	return rule.id;
}

int OverrideManager::AddRuleVectorFirstFromField(int recipientSlot, int entityIndex, const char* className, const char* vectorFieldPath, const char* sourceFieldPath, std::string& error)
{
	if (recipientSlot < -1 || recipientSlot >= 64)
	{
		error = "recipient slot must be all/-1 or 0..63";
		return 0;
	}
	if (entityIndex < 0)
	{
		error = "entity index must be >= 0";
		return 0;
	}

	CompiledOverrideRule rule;
	rule.id = m_nextId++;
	rule.kind = OverrideRuleKind::VectorFirstFromField;
	rule.recipientSlot = recipientSlot;
	rule.entityIndex = entityIndex;
	rule.className = className ? className : "";
	rule.fieldPath = vectorFieldPath ? vectorFieldPath : "";
	rule.sourceFieldPath = sourceFieldPath ? sourceFieldPath : "";
	if (!ResolveFieldPath(rule.className.c_str(), rule.fieldPath.c_str(), rule.offset, rule.size, rule.addressSteps, error))
		return 0;

	int32_t sourceOffset = 0;
	if (!ResolveFieldPath(rule.className.c_str(), rule.sourceFieldPath.c_str(), sourceOffset, rule.sourceSize, rule.sourceAddressSteps, error))
		return 0;
	if (rule.sourceSize > 0 && rule.sourceSize < 4)
	{
		error = "source field is smaller than a handle";
		return 0;
	}

	m_rules.push_back(rule);
	SPMessage("vector-first rule %d recipient=%d entity=%d %s::%s <- %s vectorSize=%d sourceSize=%d topOffset=%d steps=%zu sourceSteps=%zu\n",
		rule.id,
		rule.recipientSlot,
		rule.entityIndex,
		rule.className.c_str(),
		rule.fieldPath.c_str(),
		rule.sourceFieldPath.c_str(),
		rule.size,
		rule.sourceSize,
		rule.offset,
		rule.addressSteps.size(),
		rule.sourceAddressSteps.size());
	RebuildIndex();
	MarkEntityFullDirty(entityIndex);
	return rule.id;
}

bool OverrideManager::RemoveRule(int id)
{
	for (auto& rule : m_rules)
	{
		if (rule.id != id)
			continue;
		rule.enabled = false;
		MarkEntityFullDirty(rule.entityIndex);
		RebuildIndex();
		return true;
	}
	return false;
}

bool OverrideManager::MarkEntityFieldDirty(int entityIndex, const char* className, const char* fieldPath, std::string& error)
{
	int32_t offset = 0;
	int32_t size = 0;
	std::vector<FieldAddressStep> addressSteps;
	if (!ResolveFieldPath(className, fieldPath, offset, size, addressSteps, error))
		return false;
	return MarkEntityFieldDirtyByOffset(entityIndex, offset);
}

std::vector<AppliedOverride> OverrideManager::ApplyForPackedEntity(CPlayerSlot recipient, int entityIndex, void* entityData)
{
	std::vector<AppliedOverride> applied;
	if (!entityData || m_rulesByEntity.empty())
		return applied;

	auto found = m_rulesByEntity.find(entityIndex);
	if (found == m_rulesByEntity.end())
		return applied;

	int effectiveRecipient = recipient.Get();
	int uniqueRuleRecipient = -2;
	for (size_t ruleIndex : found->second)
	{
		const auto& rule = m_rules[ruleIndex];
		if (!rule.enabled || rule.recipientSlot < 0)
			continue;
		if (uniqueRuleRecipient == -2)
			uniqueRuleRecipient = rule.recipientSlot;
		else if (uniqueRuleRecipient != rule.recipientSlot)
		{
			uniqueRuleRecipient = -2;
			break;
		}
	}

	if ((effectiveRecipient < 0 || effectiveRecipient == 0 || effectiveRecipient == 1) && uniqueRuleRecipient >= 0)
	{
		static int fallbackLogs = 0;
		if (fallbackLogs < 64)
		{
			SPMessage("entity=%d unresolved recipient current=%d using unique rule recipient=%d\n",
				entityIndex,
				effectiveRecipient,
				uniqueRuleRecipient);
			++fallbackLogs;
		}
		effectiveRecipient = uniqueRuleRecipient;
	}

	for (size_t ruleIndex : found->second)
	{
		const auto& rule = m_rules[ruleIndex];
		if (!rule.enabled)
			continue;

		if (rule.recipientSlot != -1 && rule.recipientSlot != effectiveRecipient)
		{
			if (!rule.recipientMismatchLogged)
			{
				SPMessage("rule %d entity=%d matched packed entity but skipped recipient rule=%d current=%d field=%s::%s\n",
					rule.id,
					entityIndex,
					rule.recipientSlot,
					effectiveRecipient,
					rule.className.c_str(),
					rule.fieldPath.c_str());
				rule.recipientMismatchLogged = true;
			}
			continue;
		}

		if (rule.kind == OverrideRuleKind::VectorFirstFromField)
		{
			ApplyVectorFirstFromField(rule, entityData, applied);
			continue;
		}

		auto* address = ResolveRuleAddress(rule, entityData);
		if (!address)
			continue;

		SaveAndWrite(applied, address, rule.value.data(), rule.value.size());
	}

	return applied;
}

void OverrideManager::Restore(const std::vector<AppliedOverride>& applied)
{
	for (auto it = applied.rbegin(); it != applied.rend(); ++it)
		memcpy(it->address, it->original.data(), it->original.size());
}