gecko-dev/mozglue/tests/interceptor/TestDllInterceptorCrossProcess.cpp

/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
/* vim: set ts=8 sts=2 et sw=2 tw=80: */
/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

#include "mozilla/Attributes.h"
#include "mozilla/ArrayUtils.h"
#include "mozilla/CmdLineAndEnvUtils.h"
#include "nsWindowsDllInterceptor.h"
#include "nsWindowsHelpers.h"

#include <string>

using std::wstring;

extern "C" __declspec(dllexport) int
ReturnResult()
{
  return 2;
}

static mozilla::CrossProcessDllInterceptor::FuncHookType<decltype(&ReturnResult)>
  gOrigReturnResult;

static int
ReturnResultHook()
{
  if (gOrigReturnResult() != 2) {
    return 3;
  }

  return 0;
}

int ParentMain(int argc, wchar_t* argv[])
{
  mozilla::SetArgv0ToFullBinaryPath(argv);

  // We'll add the child process to a job so that, in the event of a failure in
  // this parent process, the child process will be automatically terminated.
  nsAutoHandle job(::CreateJobObject(nullptr, nullptr));
  if (!job) {
    printf("TEST-UNEXPECTED-FAIL | DllInterceptorCrossProcess | Job creation failed\n");
    return 1;
  }

  JOBOBJECT_EXTENDED_LIMIT_INFORMATION jobInfo = {};
  jobInfo.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;

  if (!::SetInformationJobObject(job.get(), JobObjectExtendedLimitInformation,
                                 &jobInfo, sizeof(jobInfo))) {
    printf("TEST-UNEXPECTED-FAIL | DllInterceptorCrossProcess | Job config failed\n");
    return 1;
  }

  wchar_t childArgv_1[] = L"-child";

  wchar_t* childArgv[] = {
    argv[0],
    childArgv_1
  };

  mozilla::UniquePtr<wchar_t[]>
    cmdLine(mozilla::MakeCommandLine(mozilla::ArrayLength(childArgv), childArgv));

  STARTUPINFOW si = { sizeof(si) };
  PROCESS_INFORMATION pi;
  if (!::CreateProcessW(argv[0], cmdLine.get(), nullptr, nullptr, FALSE,
                        CREATE_SUSPENDED, nullptr, nullptr, &si, &pi)) {
    printf("TEST-UNEXPECTED-FAIL | DllInterceptorCrossProcess | Failed to spawn child process\n");
    return 1;
  }

  nsAutoHandle childProcess(pi.hProcess);
  nsAutoHandle childMainThread(pi.hThread);

  if (!::AssignProcessToJobObject(job.get(), childProcess.get())) {
    printf("TEST-UNEXPECTED-FAIL | DllInterceptorCrossProcess | Failed to assign child process to job\n");
    ::TerminateProcess(childProcess.get(), 1);
    return 1;
  }

  mozilla::CrossProcessDllInterceptor intcpt(childProcess.get());
  intcpt.Init("TestDllInterceptorCrossProcess.exe");

  if (!gOrigReturnResult.Set(childProcess.get(), intcpt, "ReturnResult",
                             &ReturnResultHook)) {
    printf("TEST-UNEXPECTED-FAIL | DllInterceptorCrossProcess | Failed to add hook\n");
    return 1;
  }

  printf("TEST-PASS | DllInterceptorCrossProcess | Hook added\n");

  if (::ResumeThread(childMainThread.get()) == static_cast<DWORD>(-1)) {
    printf("TEST-UNEXPECTED-FAIL | DllInterceptorCrossProcess | Failed to resume child thread\n");
    return 1;
  }

  BOOL remoteDebugging;
  bool debugging = ::IsDebuggerPresent() ||
                   (::CheckRemoteDebuggerPresent(childProcess.get(),
                                                 &remoteDebugging) &&
                    remoteDebugging);

  DWORD waitResult = ::WaitForSingleObject(childProcess.get(),
                                           debugging ? INFINITE : 60000);
  if (waitResult != WAIT_OBJECT_0) {
    printf("TEST-UNEXPECTED-FAIL | DllInterceptorCrossProcess | Child process failed to finish\n");
    return 1;
  }

  DWORD childExitCode;
  if (!::GetExitCodeProcess(childProcess.get(), &childExitCode)) {
    printf("TEST-UNEXPECTED-FAIL | DllInterceptorCrossProcess | Failed to obtain child process exit code\n");
    return 1;
  }

  if (childExitCode) {
    printf("TEST-UNEXPECTED-FAIL | DllInterceptorCrossProcess | Child process exit code is %lu instead of 0\n", childExitCode);
    return 1;
  }

  printf("TEST-PASS | DllInterceptorCrossProcess | Child process exit code is zero\n");
  return 0;
}

extern "C"
int wmain(int argc, wchar_t* argv[])
{
  if (argc > 1) {
    // clang keeps inlining this call despite every attempt to force it to do
    // otherwise. We'll use GetProcAddress and call its function pointer instead.
    auto pReturnResult =
      reinterpret_cast<decltype(&ReturnResult)>(
        ::GetProcAddress(::GetModuleHandleW(nullptr), "ReturnResult"));
    return pReturnResult();
  }

  return ParentMain(argc, argv);
}