You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

578 lines
14 KiB

#include "NativeCode.h"
_NATIVE_CODE_LINK::_NATIVE_CODE_LINK()
{
XedDecodedInstZero(&XedInstruction);
XedDecodedInstSetMode(&XedInstruction, XED_MACHINE_MODE_LONG_64, XED_ADDRESS_WIDTH_64b);
Flags = 0UL;
Next = Prev = NULL;
Block = NULL;
Label = 0UL;
RawData = NULL;
RawDataSize = 0UL;
}
_NATIVE_CODE_LINK::_NATIVE_CODE_LINK(ULONG LabelId, _NATIVE_CODE_BLOCK* B)
: _NATIVE_CODE_LINK()
{
Block = B;
Label = LabelId;
Flags = CODE_FLAG_IS_LABEL;
}
_NATIVE_CODE_LINK::_NATIVE_CODE_LINK(ULONG F, PVOID Rd, ULONG Rds, BOOL Decode)
: _NATIVE_CODE_LINK()
{
Flags = F;
RawDataSize = Rds;
RawData = new UCHAR[Rds];
if (Rd)
{
RtlCopyMemory(RawData, Rd, Rds);
if (Decode)
XedDecode(&XedInstruction, RawData, RawDataSize);
}
}
_NATIVE_CODE_LINK::~_NATIVE_CODE_LINK()
{
if (RawData)
delete RawData;
}
_NATIVE_CODE_BLOCK::_NATIVE_CODE_BLOCK()
{
Start = End = NULL;
LabelIds.clear();
}
VOID NcAppendToBlock(PNATIVE_CODE_BLOCK Block, PNATIVE_CODE_LINK Link)
{
if (!Link)
return;
Link->Block = Block;
Link->Prev = Block->End;
Link->Next = NULL;
if (!Block->End || !Block->Start)
{
Block->Start = Block->End = Link;
}
else
{
Block->End->Next = Link;
Block->End = Link;
}
}
VOID NcPrependToBlock(PNATIVE_CODE_BLOCK Block, PNATIVE_CODE_LINK Link)
{
if (!Link)
return;
Link->Block = Block;
Link->Next = Block->Start;
Link->Prev = NULL;
if (!Block->End || !Block->Start)
{
Block->Start = Block->End = Link;
}
else
{
Block->Start->Prev = Link;
Block->Start = Link;
}
}
VOID NcInsertLinkAfter(PNATIVE_CODE_LINK Link1, PNATIVE_CODE_LINK Link2)
{
if (Link1)
{
Link2->Prev = Link1;
Link2->Next = Link1->Next;
Link1->Next = Link2;
if (Link2->Next)
Link2->Next->Prev = Link2;
}
}
VOID NcInsertLinkBefore(PNATIVE_CODE_LINK Link1, PNATIVE_CODE_LINK Link2)
{
if (Link1)
{
Link2->Next = Link1;
Link2->Prev = Link1->Prev;
Link1->Prev = Link2;
if (Link2->Prev)
Link2->Prev->Next = Link2;
}
}
VOID NcUnlink(PNATIVE_CODE_LINK Link)
{
if (Link)
{
if (Link->Next)
Link->Next->Prev = Link->Prev;
if (Link->Prev)
Link->Prev->Next = Link->Next;
}
}
ULONG NcCalcBlockSize(PNATIVE_CODE_BLOCK Block)
{
ULONG TotalSize = 0;
for (PNATIVE_CODE_LINK T = Block->Start; T && T != Block->End->Next; T = T->Next)
{
if (T->Flags & CODE_FLAG_IS_LABEL)
continue;
TotalSize += T->RawDataSize;
}
return TotalSize;
}
ULONG NcGenUnusedLabelId(PNATIVE_CODE_BLOCK Block)
{
ULONG ReturnLabelId = rand();
while (StdFind(Block->LabelIds.begin(), Block->LabelIds.end(), ReturnLabelId) != Block->LabelIds.end())
ReturnLabelId = rand();
Block->LabelIds.push_back(ReturnLabelId);
return ReturnLabelId;
}
VOID NcChangeLabelId(PNATIVE_CODE_BLOCK Block, ULONG Original, ULONG New)
{
for (PNATIVE_CODE_LINK T = Block->Start; T && T != Block->End->Next; T = T->Next)
{
if (((T->Flags & CODE_FLAG_IS_LABEL) || (T->Flags & CODE_FLAG_IS_REL_JMP)) && T->Label == Original)
T->Label = New;
}
}
VOID NcFixLabelsForBlocks(PNATIVE_CODE_BLOCK Block1, PNATIVE_CODE_BLOCK Block2)
{
for (PNATIVE_CODE_LINK T = Block2->Start; T && T != Block2->End->Next; T = T->Next)
{
if ((T->Flags & CODE_FLAG_IS_LABEL) && StdFind(Block1->LabelIds.begin(), Block1->LabelIds.end(), T->Label) != Block1->LabelIds.end())
{
NcChangeLabelId(Block2, T->Label, NcGenUnusedLabelId(Block1));
}
}
}
BOOL NcInsertBlockAfter(PNATIVE_CODE_LINK Link, PNATIVE_CODE_BLOCK Block, BOOL FixLabels)
{
if (!Link || !Link->Block || !Block || !Block->Start || !Block->End || Link->Block == Block)
return FALSE;
if (FixLabels && Block->LabelIds.size() && Link->Block->LabelIds.size())
NcFixLabelsForBlocks(Link->Block, Block);
if (Link->Next)
Link->Next->Prev = Block->End;
Block->End->Next = Link->Next;
Block->Start->Prev = Link;
Link->Next = Block->Start;
for (PNATIVE_CODE_LINK T = Block->Start; T && T != Block->End->Next; T = T->Next)
T->Block = Link->Block;
return TRUE;
}
BOOL NcInsertBlockBefore(PNATIVE_CODE_LINK Link, PNATIVE_CODE_BLOCK Block, BOOL FixLabels)
{
if (!Link || !Link->Block || !Block || !Block->Start || !Block->End)
return FALSE;
if (FixLabels && Block->LabelIds.size() && Link->Block->LabelIds.size())
NcFixLabelsForBlocks(Link->Block, Block);
if (Link->Prev)
Link->Prev->Next = Block->Start;
Block->Start->Prev = Link->Prev;
Block->End->Next = Link;
Link->Prev = Block->End;
for (PNATIVE_CODE_LINK T = Block->Start; T && T != Block->End->Next; T = T->Next)
T->Block = Link->Block;
return TRUE;
}
BOOL NcCreateLabels(PNATIVE_CODE_BLOCK Block)
{
ULONG CurrentLabelId = 0;
for (PNATIVE_CODE_LINK T = Block->Start; T; T = T->Next)
{
if (!(T->Flags & CODE_FLAG_IS_INST))
continue;
XED_CATEGORY_ENUM Category = XedDecodedInstGetCategory(&T->XedInstruction);
if (Category != XED_CATEGORY_COND_BR && Category != XED_CATEGORY_UNCOND_BR)
continue;
ULONG OperandCount = XedDecodedInstNumOperands(&T->XedInstruction);
if (OperandCount < 1)
continue;
CONST XED_INST* Inst = XedDecodedInstInst(&T->XedInstruction);
if (!Inst)
continue;
CONST XED_OPERAND* Operand = XedInstOperand(Inst, 0);
if (!Operand)
continue;
XED_OPERAND_TYPE_ENUM OperandType = XedOperandType(Operand);
if (OperandType != XED_OPERAND_TYPE_IMM && OperandType != XED_OPERAND_TYPE_IMM_CONST)
continue;
INT32 BranchDisplacement = XedDecodedInstGetBranchDisplacement(&T->XedInstruction);
PNATIVE_CODE_LINK JmpPos = NcValidateJmp(T, BranchDisplacement);
if (!JmpPos)
{
printf("Failed to validate jump. Type: %s, Displacement: %d\n", XedCategoryEnumToString(Category), BranchDisplacement);
return FALSE;
}
if (JmpPos->Prev && (JmpPos->Prev->Flags & CODE_FLAG_IS_LABEL))
{
T->Label = JmpPos->Prev->Label;
}
else
{
NcInsertLinkBefore(JmpPos, new NATIVE_CODE_LINK(CurrentLabelId, Block));
Block->LabelIds.push_back(CurrentLabelId);
T->Label = CurrentLabelId;
++CurrentLabelId;
}
T->Flags |= CODE_FLAG_IS_REL_JMP;
}
return TRUE;
}
PNATIVE_CODE_LINK NcValidateJmp(PNATIVE_CODE_LINK Jmp, INT32 Delta)
{
PNATIVE_CODE_LINK T;
if (Delta > 0)
{
T = Jmp->Next;
while (Delta > 0 && T)
{
if (T->Flags & CODE_FLAG_IS_LABEL)
continue;
Delta -= XedDecodedInstGetLength(&T->XedInstruction);
T = T->Next;
}
if (Delta != 0 || !T)
return NULL;
while (T && (T->Flags & CODE_FLAG_IS_LABEL))
T = T->Next;
return T;
}
else if (Delta < 0)
{
T = Jmp;
while (T)
{
if (T->Flags & CODE_FLAG_IS_LABEL)
continue;
Delta += XedDecodedInstGetLength(&T->XedInstruction);
if (Delta >= 0)
break;
T = T->Prev;
}
if (Delta != 0 || !T)
return NULL;
while (T && (T->Flags & CODE_FLAG_IS_LABEL))
T = T->Next;
return T;
}
return Jmp;
}
PNATIVE_CODE_LINK NcDeepCopyLink(PNATIVE_CODE_LINK Link)
{
if (Link->Flags & CODE_FLAG_IS_LABEL)
{
return new NATIVE_CODE_LINK(Link->Label, NULL);
}
else
{
PNATIVE_CODE_LINK NewLink = new NATIVE_CODE_LINK(Link->Flags, Link->RawData, Link->RawDataSize);
NewLink->Label = Link->Label;
XED_ERROR_ENUM DecodeError = XedDecode(&NewLink->XedInstruction, Link->RawData, Link->RawDataSize);
if (DecodeError != XED_ERROR_NONE)
{
printf("XedDecode failed in NcDeepCopyLink: %s %u\n", XedErrorEnumToString(DecodeError), Link->RawDataSize);
delete NewLink;
return NULL;
}
return NewLink;
}
}
BOOL NcDeepCopyPartialBlock(PNATIVE_CODE_LINK Start, PNATIVE_CODE_LINK End, PNATIVE_CODE_BLOCK Block)
{
if (!Start || !End || !Start->Block || Start->Block != End->Block || !Block)
return FALSE;
Block->LabelIds.clear();
Block->Start = Block->End = NULL;
for (ULONG L : Start->Block->LabelIds)
Block->LabelIds.push_back(L);
for (PNATIVE_CODE_LINK CurLink = Start; CurLink && CurLink != End->Next; CurLink = CurLink->Next)
{
PNATIVE_CODE_LINK Temp = NcDeepCopyLink(CurLink);
if (!Temp)
{
NcDeleteBlock(Block);
return FALSE;
}
NcAppendToBlock(Block, Temp);
}
return TRUE;
}
BOOL NcDeepCopyBlock(PNATIVE_CODE_BLOCK Block, PNATIVE_CODE_BLOCK BlockCopy)
{
return NcDeepCopyPartialBlock(Block->Start, Block->End, BlockCopy);
}
BOOL NcGetDeltaToLabel(PNATIVE_CODE_LINK Link, PINT32 DeltaOut)
{
INT32 Delta = 0;
//First checking backwards because I feel like thats the direction most jmps are in
for (PNATIVE_CODE_LINK T = Link; T; T = T->Prev)
{
if (T->Flags & CODE_FLAG_IS_LABEL)
{
if (T->Label == Link->Label)
{
*DeltaOut = Delta;
return TRUE;
}
continue;
}
Delta -= T->RawDataSize;
}
//Now check forwards
Delta = 0;
for (PNATIVE_CODE_LINK T = Link->Next; T; T = T->Next)
{
if (T->Flags & CODE_FLAG_IS_LABEL)
{
if (T->Label == Link->Label)
{
*DeltaOut = Delta;
return TRUE;
}
continue;
}
Delta += T->RawDataSize;
}
return FALSE;
}
BOOL NcFixRelJmps(PNATIVE_CODE_BLOCK Block)
{
for (PNATIVE_CODE_LINK T = Block->Start; T && T != Block->End->Next;)
{
if (T->Flags & CODE_FLAG_IS_REL_JMP)
{
INT32 BranchDisp = 0;
if (!NcGetDeltaToLabel(T, &BranchDisp))
return FALSE;
ULONG DispWidth = XedDecodedInstGetBranchDisplacementWidthBits(&T->XedInstruction);
if (log2(abs(BranchDisp)) + 1 > DispWidth)
{
//duh oh
if (DispWidth == 32)
return FALSE;
//Grow displacement width to required size
DispWidth *= 2;
//Check again
if (log2(abs(BranchDisp)) + 1 > DispWidth)
{
if (DispWidth == 32)
return FALSE;
//Grow once more if not already at 32
DispWidth *= 2;
}
//Encode new instruction
XED_STATE MachineState;
MachineState.mmode = XED_MACHINE_MODE_LONG_64;
MachineState.stack_addr_width = XED_ADDRESS_WIDTH_64b;
XED_ENCODER_INSTRUCTION EncoderInstruction;
XED_ENCODER_REQUEST EncoderRequest;
UCHAR EncodeBuffer[15];
UINT ReturnedSize;
XED_ICLASS_ENUM IClass = XedDecodedInstGetIClass(&T->XedInstruction);
//Do the encoding
XedInst1(&EncoderInstruction, MachineState, IClass, DispWidth, XedRelBr(0, DispWidth));
XedEncoderRequestZeroSetMode(&EncoderRequest, &MachineState);
if (!XedConvertToEncoderRequest(&EncoderRequest, &EncoderInstruction))
return FALSE;
if (XED_ERROR_NONE != XedEncode(&EncoderRequest, EncodeBuffer, 15, &ReturnedSize))
return FALSE;
//fixup T->RawData
delete[] T->RawData;
T->RawDataSize = ReturnedSize;
T->RawData = new UCHAR[ReturnedSize];
RtlCopyMemory(T->RawData, EncodeBuffer, ReturnedSize);
//Decode instruction so its proper and all that
XedDecodedInstZeroSetMode(&T->XedInstruction, &MachineState);
if (XED_ERROR_NONE != XedDecode(&T->XedInstruction, T->RawData, T->RawDataSize))
return FALSE;
//Go back to the start and loop through all labels again because now this instruction is larger :))))
T = Block->Start;
continue;
}
else
{
DispWidth = XedDecodedInstGetBranchDisplacementWidth(&T->XedInstruction);
switch (DispWidth)
{
case 1: *(PINT8)&T->RawData[T->RawDataSize - DispWidth] = (INT8)BranchDisp; break;
case 2: *(PINT16)&T->RawData[T->RawDataSize - DispWidth] = (INT16)BranchDisp; break;
case 4: *(PINT32)&T->RawData[T->RawDataSize - DispWidth] = (INT32)BranchDisp; break;
}
}
}
T = T->Next;
}
return TRUE;
}
BOOL NcDisassemble(PNATIVE_CODE_BLOCK Block, PVOID Buffer, ULONG BufferSize)
{
PUCHAR Buf = (PUCHAR)Buffer;
ULONG Offset = 0;
while (Offset < BufferSize)
{
PNATIVE_CODE_LINK Link = new NATIVE_CODE_LINK;
Link->Flags = CODE_FLAG_IS_INST;
ULONG PossibleSize = min(15, BufferSize - Offset);
XED_ERROR_ENUM DecodeError = XedDecode(&Link->XedInstruction, (Buf + Offset), PossibleSize);
if (DecodeError != XED_ERROR_NONE)
{
printf("XedDecode failed with error %s\n", XedErrorEnumToString(DecodeError));
NcDeleteBlock(Block);
delete Link;
return FALSE;
}
Link->RawDataSize = XedDecodedInstGetLength(&Link->XedInstruction);
Link->RawData = new UCHAR[Link->RawDataSize];
RtlCopyMemory(Link->RawData, (Buf + Offset), Link->RawDataSize);
NcAppendToBlock(Block, Link);
Offset += Link->RawDataSize;
}
NcCreateLabels(Block);
return TRUE;
}
PVOID NcAssemble(PNATIVE_CODE_BLOCK Block, PULONG OutSize)
{
if (!NcFixRelJmps(Block))
return NULL;
*OutSize = NcCalcBlockSize(Block);
PUCHAR Buffer = (PUCHAR)malloc(*OutSize);
if (!Buffer)
return NULL;
PUCHAR BufferOffset = Buffer;
for (PNATIVE_CODE_LINK T = Block->Start; T && T != Block->End->Next; T = T->Next)
{
if (T->Flags & CODE_FLAG_IS_LABEL)
continue;
RtlCopyMemory(BufferOffset, T->RawData, T->RawDataSize);
BufferOffset += T->RawDataSize;
}
return Buffer;
}
VOID NcDeleteBlock(PNATIVE_CODE_BLOCK Block)
{
if (!Block->Start || !Block->End)
return;
PNATIVE_CODE_LINK BlockEnding = Block->End->Next;
for (PNATIVE_CODE_LINK T = Block->Start; T && T != BlockEnding;)
{
PNATIVE_CODE_LINK Next = T->Next;
delete T;
T = Next;
}
}
VOID NcDebugPrint(PNATIVE_CODE_BLOCK Block)
{
HANDLE ConsoleHandle = GetStdHandle(STD_OUTPUT_HANDLE);
if (!ConsoleHandle)
return;
for (PNATIVE_CODE_LINK T = Block->Start; T && T != Block->End->Next; T = T->Next)
{
if (T->Flags & CODE_FLAG_IS_LABEL)
{
SetConsoleTextAttribute(ConsoleHandle, FOREGROUND_GREEN | FOREGROUND_RED);
printf("Label: %u\n", T->Label);
}
else
{
XED_ICLASS_ENUM IClass = XedDecodedInstGetIClass(&T->XedInstruction);
if (T->Flags & CODE_FLAG_IS_REL_JMP)
{
SetConsoleTextAttribute(ConsoleHandle, FOREGROUND_GREEN | FOREGROUND_RED);
printf("%s: %u\n", XedIClassEnumToString(IClass), T->Label);
}
else
{
SetConsoleTextAttribute(ConsoleHandle, FOREGROUND_GREEN | FOREGROUND_BLUE);
printf("%s\n", XedIClassEnumToString(IClass));
}
}
}
}
VOID NcPrintBlockCode(PNATIVE_CODE_BLOCK Block)
{
for (PNATIVE_CODE_LINK T = Block->Start; T && T != Block->End->Next; T = T->Next)
{
if (!(T->Flags & CODE_FLAG_IS_LABEL))
{
for (uint32_t i = 0; i < T->RawDataSize; i++)
{
std::cout << std::hex << std::setw(2) << std::setfill('0') << (int)T->RawData[i] << ' ';
}
}
}
}