SuperReturn

SuperReturn.c

/// Return, but across multiple frames.
///
/// This function unwinds the given number of frames, then sets the return value provided, emulating as if this number
/// of functions returned, with the last one returning the value provided in RetVal. Can be used to hook a callee when
/// you don't have a convenient way to hook it directly and actually just want to stub it out with a return value.
///
/// @param FramesToSkip The number of frames to skip, starting from the current frame.
/// @param RetVal The value to return from the last frame.
/// @param Context Context to start from, in case you want to SuperReturn from somewhere deeper.
DECLSPEC_NOINLINE void SuperReturn(
  _In_ ULONG FramesToSkip,
  _In_opt_ ULONG_PTR RetVal,
  _In_opt_ PCONTEXT Context
) {

  CONTEXT LocalContext;
  if (!Context) {
    FramesToSkip += 1; // skip this frame
    LocalContext.ContextFlags = CONTEXT_CONTROL | CONTEXT_INTEGER;
    RtlCaptureContext(&LocalContext);
    Context = &LocalContext;
  }

#if defined(_M_X64)
#define CTX_IP(Ctx) (Ctx->Rip)
#define CTX_SP(Ctx) (Ctx->Rsp)
#define CTX_RV(Ctx) (Ctx->Rax)
#elif defined(_M_ARM64)
#define CTX_IP(Ctx) (Ctx->Pc)
#define CTX_SP(Ctx) (Ctx->Sp)
#define CTX_RV(Ctx) (Ctx->X0)
#elif defined(_M_IX86)
#error Can't possibly work on x86: no way to restore nonvolatile registers.
#else
#error Unsupported architecture!
#endif

  ULONG64 ControlPc = CTX_IP(Context);

  for (ULONG i = 0; i < FramesToSkip; i++) {
    ULONG_PTR ImageBase = 0;
    PRUNTIME_FUNCTION FunctionEntry = RtlLookupFunctionEntry(ControlPc, &ImageBase, NULL);

    if (!FunctionEntry) {
      // leaf
      CTX_IP(Context) = *(ULONG64*)CTX_SP(Context);
      CTX_SP(Context) += sizeof(ULONG64);
    } else {
      PVOID HandlerData;
      ULONG64 EstablisherFrame;
      RtlVirtualUnwind(
        UNW_FLAG_NHANDLER,
        ImageBase,
        ControlPc,
        FunctionEntry,
        Context,
        &HandlerData,
        &EstablisherFrame,
        NULL
      );
    }

    ControlPc = CTX_IP(Context);
  }

  CTX_RV(Context) = RetVal;

#undef CTX_IP
#undef CTX_SP
#undef CTX_RV

  NtContinue(Context, FALSE);
}