Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Reworked BlueprintHookManager to better-prevent hooks from corrupting the code or each other #320

Open
wants to merge 6 commits into
base: dev
Choose a base branch
from
Open

Reworked BlueprintHookManager to better-prevent hooks from corrupting the code or each other #320

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
5de489d
Reworked BlueprintHookManager to better-prevent hooks from corrupting…
Epp-code Nov 7, 2024
2d37282
Addressed first round of code review comments
Epp-code Nov 9, 2024
97f9226
Second round of code review comments.
Epp-code Nov 9, 2024
a88eb30
Fixed scenario where hooking function return by exact index AND by EP…
Epp-code Nov 9, 2024
93412b0
Added assertion that Function is not null to HookBlueprintFunction to…
Epp-code Nov 9, 2024
190a9a3
Changed additional checkf to fgcheckf so it won't be compiled out
Epp-code Nov 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
244 changes: 167 additions & 77 deletions Mods/SML/Source/SML/Private/Patching/BlueprintHookManager.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,18 +12,18 @@ DEFINE_LOG_CATEGORY(LogBlueprintHookManager);

#define WRITE_UNALIGNED(Arr, Type, Value) \
Arr.AddUninitialized(sizeof(Type)); \
FPlatformMemory::WriteUnaligned<Type>(&AppendedCode[Arr.Num() - sizeof(Type)], (Type) Value);
FPlatformMemory::WriteUnaligned<Type>(&Arr[Arr.Num() - sizeof(Type)], (Type) Value);

void UBlueprintHookManager::HandleHookedFunctionCall(FFrame& Stack, int32 HookOffset) {
FFunctionHookInfo& FunctionHookInfo = HookedFunctions.FindChecked(Stack.Node);
FunctionHookInfo.InvokeBlueprintHook(Stack, HookOffset);
}

#if DEBUG_BLUEPRINT_HOOKING
void DebugDumpFunctionScriptCode(UFunction* Function, int32 HookOffset, const FString& Postfix) {
const FString FileLocation = FPaths::RootDir() + FString::Printf(TEXT("BlueprintHookingDebug_%s_%s_at_%d_%s.json"), *Function->GetOuter()->GetName(), *Function->GetName(), HookOffset, *Postfix);
void DebugDumpFunctionScriptCode(UFunction* Function, int32 HookOffset, int32 ResolvedHookOffset, const FString& Postfix) {
const FString FileLocation = FPaths::RootDir() + FString::Printf(TEXT("BlueprintHookingDebug_%s_%s_at_%d_resolved_%d_%s.json"), *Function->GetOuter()->GetName(), *Function->GetName(), HookOffset, ResolvedHookOffset, *Postfix);

FKismetBytecodeDisassemblerJson Disassembler;
FSMLKismetBytecodeDisassembler Disassembler;
Epp-code marked this conversation as resolved.
Show resolved Hide resolved
const TArray<TSharedPtr<FJsonValue>> Statements = Disassembler.SerializeFunction(Function);

FString OutJsonString;
Expand All @@ -35,102 +35,163 @@ void DebugDumpFunctionScriptCode(UFunction* Function, int32 HookOffset, const FS
}
#endif

void UBlueprintHookManager::InstallBlueprintHook(UFunction* Function, int32 HookOffset) {
void UBlueprintHookManager::InstallBlueprintHook(UFunction* Function, const int32 OriginalHookOffset, const int32 ResolvedHookOffset) {
TArray<uint8>& OriginalCode = Function->Script;
checkf(OriginalCode.Num() > HookOffset, TEXT("Invalid hook: HookOffset > Script.Num()"));
fgcheckf(OriginalCode.Num() > ResolvedHookOffset, TEXT("Invalid hook: Resolved HookOffset > Script.Num()"));

#if DEBUG_BLUEPRINT_HOOKING
DebugDumpFunctionScriptCode(Function, HookOffset, TEXT("BeforeHook"));
DebugDumpFunctionScriptCode(Function, OriginalHookOffset, ResolvedHookOffset, TEXT("BeforeHook"));
#endif

//Minimum amount of bytes required to insert unconditional jump with code offset
const int32 MinBytesRequired = 1 + sizeof(CodeSkipSizeType);

// Will throw an error if the resolved hook offset is not at properly-aligned, parseable opcode
FSMLKismetBytecodeDisassembler Disassembler;
int32 BytesAvailable = 0;

//Walk over statements until we collect enough bytes for a replacement
//(or until we consumed all statements in the function's code)
while (BytesAvailable < MinBytesRequired && (HookOffset + BytesAvailable) < OriginalCode.Num()) {
const int32 CurrentStatementIndex = HookOffset + BytesAvailable;
int32 OutStatementLength;

const bool bValid = Disassembler.GetStatementLength(Function, CurrentStatementIndex, OutStatementLength);
checkf(bValid, TEXT("Provided hook offset is not a valid statement index: %d"), HookOffset);
BytesAvailable += OutStatementLength;
}
Disassembler.SerializeStatement(Function, ResolvedHookOffset);
Epp-code marked this conversation as resolved.
Show resolved Hide resolved

//Check that we collected enough bytes
if (BytesAvailable < MinBytesRequired) {
//If we are here, it means we consumed all the statements in the function's code
//And still don't have enough space for inserting a jump. In that case, we append additional
//EX_EndOfScript instructions until we have enough place
const int32 BytesToAppend = MinBytesRequired - BytesAvailable;
OriginalCode.AddUninitialized(BytesToAppend);
FPlatformMemory::Memset(&OriginalCode[OriginalCode.Num() - BytesToAppend], EX_EndOfScript, BytesToAppend);
BytesAvailable = MinBytesRequired;
// First we go over the existing code and add UBlueprintHookManager::JumpBytesRequired to all
// the offsets. Afterwards we will move the relevant code and add the jump. Doing it in this
// order just means we don't have to worry about errantly changing the jump offset we will add.
TArray<TSharedPtr<FJsonValue>> DisassembledFunction = Disassembler.SerializeFunction(Function);
for (TSharedPtr<FJsonValue> JsonValue : DisassembledFunction) {
ModifyOffsetsForNewHookOffset(OriginalCode, JsonValue->AsObject(), ResolvedHookOffset);
}

//Add enough room to add the jump. This will shift all the post-hookoffset bytes down so adding
//the jump at the resolved offset doesn't overwrite valid instructions and data
OriginalCode.InsertUninitialized(ResolvedHookOffset, UBlueprintHookManager::JumpBytesRequired);

//Generate code required for calling our hook
//We use EX_CallMath for speed since our inserted function doesn't need context, and is fine with being called on CDO
TArray<uint8> AppendedCode;

//Make sure hook function is not NULL, otherwise we may experience weird crashes later
UFunction* HookCallFunction = UBlueprintHookManager::StaticClass()->FindFunctionByName(TEXT("ExecuteBPHook"));
check(HookCallFunction);
fgcheck(HookCallFunction);

//We use EX_CallMath for speed since our inserted function doesn't need context, and is fine with being called on CDO
//EX_CallMath requires just UFunction object pointer and argument list
AppendedCode.Add(EX_CallMath);
WRITE_UNALIGNED(AppendedCode, ScriptPointerType, HookCallFunction);

//Begin writing function parameters - we have just hook offset constant
//Begin writing function parameters - we have just the original hook offset constant.
//The hook function needs this because it stores the user hooks according to their original offset, not the resolved offset.
AppendedCode.Add(EX_IntConst);
WRITE_UNALIGNED(AppendedCode, int32, HookOffset);
WRITE_UNALIGNED(AppendedCode, int32, OriginalHookOffset);
AppendedCode.Add(EX_EndFunctionParms);


//Append original code that we replaced earlier with unconditional jump
AppendedCode.AddUninitialized(BytesAvailable);
FPlatformMemory::Memcpy(&AppendedCode[AppendedCode.Num() - BytesAvailable], &OriginalCode[HookOffset], BytesAvailable);

//Insert jump to original location for running code after hook
//Insert jump to after the resolved hook offset location for running code after hook
AppendedCode.Add(EX_Jump);
const int32 JumpDestination = HookOffset + BytesAvailable;
const int32 JumpDestination = ResolvedHookOffset + UBlueprintHookManager::JumpBytesRequired;
WRITE_UNALIGNED(AppendedCode, CodeSkipSizeType, JumpDestination);

//Finish generated code with EX_EndOfScript to avoid any surprises
AppendedCode.Add(EX_EndOfScript);


//Append generated code to the end of the function's original code now
const int32 StartOfAppendedCode = OriginalCode.Num();
OriginalCode.Append(AppendedCode);

//Fill space with EX_EndOfScript before replacement for safety
FPlatformMemory::Memset(&OriginalCode[HookOffset], EX_EndOfScript, BytesAvailable);

//Actually insert jump to the start of appended code to original hook location
OriginalCode[HookOffset] = EX_Jump;
FPlatformMemory::WriteUnaligned<CodeSkipSizeType>(&OriginalCode[HookOffset + 1], StartOfAppendedCode);
OriginalCode[ResolvedHookOffset] = EX_Jump;
FPlatformMemory::WriteUnaligned<CodeSkipSizeType>(&OriginalCode[ResolvedHookOffset + 1], StartOfAppendedCode);

#if DEBUG_BLUEPRINT_HOOKING
DebugDumpFunctionScriptCode(Function, HookOffset, TEXT("AfterHook"));
DebugDumpFunctionScriptCode(Function, OriginalHookOffset, ResolvedHookOffset, TEXT("AfterHook"));
#endif
}

int32 UBlueprintHookManager::PreProcessHookOffset(UFunction* Function, int32 HookOffset) {
if (HookOffset == EPredefinedHookOffset::Return) {
//For now Kismet Compiler will always generate only one Return node, so all
//execution paths will end up either with executing it directly or jumping to it
//So we need to hook only in one place to handle all possible execution paths
FSMLKismetBytecodeDisassembler Disassembler;
int32 ReturnOffset;
const bool bIsValid = Disassembler.FindFirstStatementOfType(Function, 0, EX_Return, ReturnOffset);
checkf(bIsValid, TEXT("EX_Return not found for function %s"), *Function->GetPathName());
return ReturnOffset;
void UBlueprintHookManager::ModifyOffsetsForNewHookOffset(TArray<uint8>& Script, TSharedPtr<FJsonObject> Expression, int32 HookOffset)
{
int32 Opcode;
if (!Expression->TryGetNumberField(TEXT("Opcode"), Opcode)) {
// No opcode means it's not an instruction, so we can just return;
return;
}

//A computed jump could do anything at runtime - it could jump to before or after the hook, so we have no way of knowing how
//it needs to be modified to work correctly. For now, the only predictable solution is to forbid hooking functions with them.
fgcheckf(Opcode != EX_ComputedJump, TEXT("Cannot hook a blueprint function that contains an EX_ComputedJump instruction. There's no way to guarantee it would not crash at some point."));

int32 IndexAfterOpcode = 1 + Expression->GetIntegerField(TEXT("OpcodeIndex"));

// This switch was created by comparing to the serialization done in FSMLKismetBytecodeDisassembler::SerializeExpression to
// identify which opcodes can jump to offsets and where exactly those jump targets reside in the script.
switch (Opcode)
{
case EX_Jump:
case EX_JumpIfNot:
case EX_Skip:
case EX_PushExecutionFlow:
{
int32 IndexOfCurrentJumpOffset = IndexAfterOpcode;
int32 CurrentJumpOffset = Expression->GetIntegerField(TEXT("Offset"));
if (CurrentJumpOffset > HookOffset) {
FPlatformMemory::WriteUnaligned(&Script[IndexOfCurrentJumpOffset], (CodeSkipSizeType)(CurrentJumpOffset + UBlueprintHookManager::JumpBytesRequired));
}
break;
}
case EX_ClassContext:
case EX_Context:
case EX_Context_FailSilent:
{
int32 CurrentJumpOffset = Expression->GetIntegerField(TEXT("SkipOffsetForNull"));
if (CurrentJumpOffset > HookOffset) {
// The offset is just past the Context object, so we have to get its size so we can write to the correct place
TSharedPtr<FJsonObject> Context = Expression->GetObjectField(TEXT("Context"));
int32 SizeOfContext = (int32)Context->GetNumberField(TEXT("OpSizeInBytes"));
int32 IndexOfCurrentJumpOffset = IndexAfterOpcode + SizeOfContext;
FPlatformMemory::WriteUnaligned(&Script[IndexOfCurrentJumpOffset], (CodeSkipSizeType)(CurrentJumpOffset + UBlueprintHookManager::JumpBytesRequired));
}
break;
}
case EX_SwitchValue:
{
int32 CurrentJumpOffset = Expression->GetIntegerField(TEXT("OffsetToSwitchEnd"));
if (CurrentJumpOffset > HookOffset) {
// The switch end offset is just past the a single Word field holding the number of cases in the switch
int32 IndexOfCurrentJumpOffset = IndexAfterOpcode + sizeof(uint16);
FPlatformMemory::WriteUnaligned(&Script[IndexOfCurrentJumpOffset], (CodeSkipSizeType)(CurrentJumpOffset + UBlueprintHookManager::JumpBytesRequired));
IndexOfCurrentJumpOffset += sizeof(CodeSkipSizeType); // Move past the offset we just wrote
int32 SwitchOpSizeInBytes = (int32)Expression->GetObjectField(TEXT("Expression"))->GetNumberField(TEXT("OpSizeInBytes"));
IndexOfCurrentJumpOffset += SwitchOpSizeInBytes; // Move past the switch expression
// Each case in the switch has an absolute offset reference to the next case, so we have to adjust each of them, as well
TArray<TSharedPtr<FJsonValue>> Cases = Expression->GetArrayField(TEXT("Cases"));
for (TSharedPtr<FJsonValue>& Case : Cases) {
TSharedPtr<FJsonObject> NextCase = Case->AsObject();
int32 CaseValueSizeInBytes = (int32)NextCase->GetObjectField(TEXT("CaseValue"))->GetNumberField(TEXT("OpSizeInBytes"));
IndexOfCurrentJumpOffset += CaseValueSizeInBytes;
int32 OffsetToNextCase = (int32)NextCase->GetNumberField(TEXT("OffsetToNextCase"));
FPlatformMemory::WriteUnaligned(&Script[IndexOfCurrentJumpOffset], (CodeSkipSizeType)(OffsetToNextCase + UBlueprintHookManager::JumpBytesRequired));
IndexOfCurrentJumpOffset += sizeof(CodeSkipSizeType);
int32 CaseResultSizeInBytes = (int32)NextCase->GetObjectField(TEXT("CaseResult"))->GetNumberField(TEXT("OpSizeInBytes"));
IndexOfCurrentJumpOffset += CaseResultSizeInBytes;
}
}
break;
}
case EX_AutoRtfmTransact:
{
int32 CurrentJumpOffset = Expression->GetIntegerField(TEXT("Offset"));
if (CurrentJumpOffset > HookOffset) {
// The offset is beyond a single TransactionId int
int32 IndexOfCurrentJumpOffset = IndexAfterOpcode + sizeof(int32);
FPlatformMemory::WriteUnaligned(&Script[IndexOfCurrentJumpOffset], (CodeSkipSizeType)(CurrentJumpOffset + UBlueprintHookManager::JumpBytesRequired));
}
break;
}
}

// Now we search all the children of this node for jump instructions that need to be updated
for (auto& Pair : Expression->Values) {
TSharedPtr<FJsonObject>* ExpressionValue;
if (Pair.Value->TryGetObject(ExpressionValue)) {
ModifyOffsetsForNewHookOffset(Script, *ExpressionValue, HookOffset);
continue;
}
TArray<TSharedPtr<FJsonValue>>* ArrayValue;
if (Pair.Value->TryGetArray(ArrayValue)) {
for (TSharedPtr<FJsonValue> Value : *ArrayValue) {
if (Value->TryGetObject(ExpressionValue)) {
ModifyOffsetsForNewHookOffset(Script, *ExpressionValue, HookOffset);
}
}
}
}
return HookOffset;
}

void FFunctionHookInfo::InvokeBlueprintHook(FFrame& Frame, int32 HookOffset) {
Expand All @@ -143,23 +204,20 @@ void FFunctionHookInfo::InvokeBlueprintHook(FFrame& Frame, int32 HookOffset) {

void FFunctionHookInfo::RecalculateReturnStatementOffset(UFunction* Function) {
FSMLKismetBytecodeDisassembler Disassembler;
int32 ReturnInstructionOffset;
Disassembler.FindFirstStatementOfType(Function, 0, EX_Return, ReturnInstructionOffset);
this->ReturnStatementOffset = ReturnInstructionOffset;
ReturnStatementOffset = Disassembler.GetReturnStatementOffset(Function);
}

void UBlueprintHookManager::HookBlueprintFunction(UFunction* Function, const TFunction<HookFunctionSignature>& Hook, int32 HookOffset) {
void UBlueprintHookManager::HookBlueprintFunction(UFunction* Function, const TFunction<HookFunctionSignature>& Hook, const int32 HookOffset) {
#if !WITH_EDITOR
checkf(Function->Script.Num(), TEXT("HookBPFunction: Function provided is not implemented in BP"));

fgcheckf(Function, TEXT("HookBPFunction: Function provided is null"));
fgcheckf(Function->Script.Num(), TEXT("HookBPFunction: Function provided is not implemented in BP"));

//Make sure to add outer UClass to root set to avoid it being Garbage Collected
//Because otherwise after GC script byte code will be reloaded, without our hooks applied
UClass* OuterUClass = Function->GetTypedOuter<UClass>();
check(OuterUClass);
fgcheck(OuterUClass);
HookedClasses.AddUnique(OuterUClass);

HookOffset = PreProcessHookOffset(Function, HookOffset);


#if UE_BLUEPRINT_EVENTGRAPH_FASTCALLS
if (Function->EventGraphFunction != nullptr) {
UE_LOG(LogBlueprintHookManager, Warning, TEXT("Attempt to hook event graph call stub function with fast-call enabled, disabling fast call for that function"));
Expand All @@ -169,16 +227,48 @@ void UBlueprintHookManager::HookBlueprintFunction(UFunction* Function, const TFu
}
#endif

FSMLKismetBytecodeDisassembler Disassembler;
bool IsFirstTimeFunctionEverHooked = !HookedFunctions.Contains(Function);
FFunctionHookInfo& FunctionHookInfo = HookedFunctions.FindOrAdd(Function);
TArray<TFunction<HookFunctionSignature>>& InstalledHooks = FunctionHookInfo.CodeOffsetByHookList.FindOrAdd(HookOffset);
if (IsFirstTimeFunctionEverHooked)
{
FunctionHookInfo.OriginalReturnStatementOffset = Disassembler.GetReturnStatementOffset(Function);
}

//Each new offset modifies the code but we need to keep track by original offsets because callers cannot know how the code has been modified by other hooks.
int32 StoredHookOffset = HookOffset == EPredefinedHookOffset::Return ? FunctionHookInfo.OriginalReturnStatementOffset : HookOffset;
TArray<TFunction<HookFunctionSignature>>& InstalledHooks = FunctionHookInfo.CodeOffsetByHookList.FindOrAdd(StoredHookOffset);

if (InstalledHooks.Num() == 0) {
//First time function is hooked at this offset, call InstallBlueprintHook
InstallBlueprintHook(Function, HookOffset);
//Update cached return instruction offset
//First time function is hooked at this requested offset. We need to resolve what the offset actually is.

int32 ResolvedHookOffset = HookOffset;
if (ResolvedHookOffset == EPredefinedHookOffset::Return) {
//Special case for hooking the return, which has an absolute location that we can directly find.
//For now Kismet Compiler will always generate only one Return node, so all
//execution paths will end up either with executing it directly or jumping to it
//So we need to hook only in one place to handle all possible execution paths
FSMLKismetBytecodeDisassembler Disassembler;
ResolvedHookOffset = Disassembler.GetReturnStatementOffset(Function);
} else {
//Each new offset moves the subsequent code by UBlueprintHookManager::JumpBytesRequired to make room for the hook jump.
//So the resolved offset must be increased by JumpBytesRequired for every hook installed earlier than it in the instruction list.
TArray<int32> HookOffsetKeys;
FunctionHookInfo.CodeOffsetByHookList.GetKeys(HookOffsetKeys);
for (int32 ExistingHookedOffset : HookOffsetKeys) {
if (ExistingHookedOffset < HookOffset) {
ResolvedHookOffset += UBlueprintHookManager::JumpBytesRequired;
}
}
}

InstallBlueprintHook(Function, StoredHookOffset, ResolvedHookOffset);

//Update cached return instruction offset because we've edited the function and moved instructions around
FunctionHookInfo.RecalculateReturnStatementOffset(Function);
}
//Add provided hook into the array
InstalledHooks.Add(Hook);

#endif
}
Loading