/* 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 http://mozilla.org/MPL/2.0/. */

#include <windows.h>
#include <bits.h>
#include <utility>

// Avoid conversions, we will only build Unicode anyway.
#if !(defined(UNICODE) && defined(_UNICODE))
#  error "Unicode required"
#endif

static HINSTANCE gHInst;

// ***** Section: ScopeExit
// Derived from mfbt mozilla::ScopeExit, I have removed the use of
// GuardObjectNotifier and the MOZ_* annotations.
template <typename ExitFunction>
class ScopeExit {
  ExitFunction mExitFunction;
  bool mExecuteOnDestruction;

 public:
  explicit ScopeExit(ExitFunction &&cleanup)
      : mExitFunction(cleanup), mExecuteOnDestruction(true) {}

  ScopeExit(ScopeExit &&rhs)
      : mExitFunction(std::move(rhs.mExitFunction)),
        mExecuteOnDestruction(rhs.mExecuteOnDestruction) {
    rhs.release();
  }

  ~ScopeExit() {
    if (mExecuteOnDestruction) {
      mExitFunction();
    }
  }

  void release() { mExecuteOnDestruction = false; }

 private:
  explicit ScopeExit(const ScopeExit &) = delete;
  ScopeExit &operator=(const ScopeExit &) = delete;
  ScopeExit &operator=(ScopeExit &&) = delete;
};

template <typename ExitFunction>
ScopeExit<ExitFunction> MakeScopeExit(ExitFunction &&exitFunction) {
  return ScopeExit<ExitFunction>(std::move(exitFunction));
}

// ***** Section: NSIS stack
typedef struct _stack_t {
  struct _stack_t *next;
  WCHAR text[1];  // this should be the length of g_stringsize when allocating
} stack_t;

static unsigned int g_stringsize;
static stack_t **g_stacktop;

static int popstringn(LPWSTR str, int maxlen) {
  stack_t *th;
  if (!g_stacktop || !*g_stacktop) return 1;
  th = (*g_stacktop);
  if (str) lstrcpynW(str, th->text, maxlen ? maxlen : g_stringsize);
  *g_stacktop = th->next;
  GlobalFree((HGLOBAL)th);
  return 0;
}

static void pushstring(LPCWSTR str) {
  stack_t *th;
  if (!g_stacktop) return;
  th = (stack_t *)GlobalAlloc(
      GPTR, (sizeof(stack_t) + (g_stringsize) * sizeof(*str)));
  lstrcpynW(th->text, str, g_stringsize);
  th->next = *g_stacktop;
  *g_stacktop = th;
}

// ***** Section: NSIS Plug-In API (from NSIS api.h)
// NSIS Plug-In Callback Messages
enum NSPIM {
  NSPIM_UNLOAD,     // This is the last message a plugin gets, do final cleanup
  NSPIM_GUIUNLOAD,  // Called after .onGUIEnd
};

// Prototype for callbacks registered with
// extra_parameters->RegisterPluginCallback() Return NULL for unknown messages
// Should always be __cdecl for future expansion possibilities
typedef UINT_PTR (*NSISPLUGINCALLBACK)(enum NSPIM);

#define NSISCALL __stdcall

typedef struct {
  LPVOID exec_flags;
  int(NSISCALL *ExecuteCodeSegment)(int, HWND);
  void(NSISCALL *validate_filename)(LPWSTR);
  int(NSISCALL *RegisterPluginCallback)(
      HMODULE, NSISPLUGINCALLBACK);  // returns 0 on success, 1 if already
                                     // registered and < 0 on errors
} extra_parameters;

// ***** Section: StartBitsThread
UINT_PTR __cdecl NSISPluginCallback(NSPIM msg);

static struct {
  HANDLE thread;
  bool shutdown_requested;
  CRITICAL_SECTION cs;
  CONDITION_VARIABLE cv;
} gStartBitsThread = {nullptr, false, 0, 0};

// This thread connects to the BackgroundCopyManager, which may take some time
// if the BITS service is not already running. It also holds open the connection
// until gStartBitsThread.shutdown_requested becomes true.
DWORD WINAPI StartBitsThreadProc(LPVOID) {
  EnterCriticalSection(&gStartBitsThread.cs);
  auto leaveCS =
      MakeScopeExit([] { LeaveCriticalSection(&gStartBitsThread.cs); });

  if (FAILED(CoInitializeEx(nullptr, COINIT_MULTITHREADED))) {
    return 0;
  }
  auto coUninit = MakeScopeExit([] { CoUninitialize(); });

  IBackgroundCopyManager *bcm = nullptr;
  if (FAILED(CoCreateInstance(
          __uuidof(BackgroundCopyManager), nullptr, CLSCTX_LOCAL_SERVER,
          __uuidof(IBackgroundCopyManager), (LPVOID *)&bcm)) ||
      !bcm) {
    return 0;
  }

  do {
    SleepConditionVariableCS(&gStartBitsThread.cv, &gStartBitsThread.cs,
                             INFINITE);
  } while (!gStartBitsThread.shutdown_requested);

  bcm->Release();
  return 1;
}

// Start up the thread
// returns true on success
bool StartBitsServiceBackgroundThreadImpl(extra_parameters *extra_params) {
  EnterCriticalSection(&gStartBitsThread.cs);
  auto leaveCS =
      MakeScopeExit([] { LeaveCriticalSection(&gStartBitsThread.cs); });

  if (gStartBitsThread.thread) {
    // Thread is already started, assumed to be still running.
    return true;
  }

  // Ensure the callback is registered so the thread can be stopped, and also so
  // NSIS doesn't unload this DLL.
  extra_params->RegisterPluginCallback(gHInst, NSISPluginCallback);

  gStartBitsThread.shutdown_requested = false;

  gStartBitsThread.thread =
      CreateThread(nullptr, 0, StartBitsThreadProc, nullptr, 0, 0);
  if (!gStartBitsThread.thread) {
    return false;
  }

  return true;
}

// Shut down the Start BITS thread, if it was started.
void ShutdownStartBitsThread() {
  EnterCriticalSection(&gStartBitsThread.cs);
  if (gStartBitsThread.thread) {
    gStartBitsThread.shutdown_requested = true;
    WakeAllConditionVariable(&gStartBitsThread.cv);
    LeaveCriticalSection(&gStartBitsThread.cs);

    // Give the thread a little time to clean up.
    if (WaitForSingleObject(gStartBitsThread.thread, 1000) == WAIT_OBJECT_0) {
      EnterCriticalSection(&gStartBitsThread.cs);
      gStartBitsThread.thread = nullptr;
      LeaveCriticalSection(&gStartBitsThread.cs);
    } else {
      // Don't attempt to recover if we didn't see the thread end,
      // the process will be exiting soon anyway.
    }

  } else {
    LeaveCriticalSection(&gStartBitsThread.cs);
  }
}

// ***** Section: CancelBitsJobsByName
#define MAX_JOB_NAME 256

bool CancelBitsJobsByNameImpl(LPWSTR matchJobName) {
  if (FAILED(CoInitialize(nullptr))) {
    return false;
  }
  auto coUninit = MakeScopeExit([] { CoUninitialize(); });

  IBackgroundCopyManager *bcm = nullptr;
  if (FAILED(CoCreateInstance(
          __uuidof(BackgroundCopyManager), nullptr, CLSCTX_LOCAL_SERVER,
          __uuidof(IBackgroundCopyManager), (LPVOID *)&bcm)) ||
      !bcm) {
    return false;
  }
  auto bcmRelease = MakeScopeExit([bcm] { bcm->Release(); });

  IEnumBackgroundCopyJobs *enumerator = nullptr;
  // Attempt to enumerate jobs for all users. If that fails,
  // try for only the current user.
  if (FAILED(bcm->EnumJobs(BG_JOB_ENUM_ALL_USERS, &enumerator))) {
    enumerator = nullptr;
    if (FAILED(bcm->EnumJobs(0, &enumerator))) {
      return false;
    }
  }
  if (!enumerator) {
    return false;
  }
  auto enumeratorRelease =
      MakeScopeExit([enumerator] { enumerator->Release(); });

  bool success = true;

  IBackgroundCopyJob *job = nullptr;
  HRESULT nextResult;
  while ((nextResult = enumerator->Next(1, &job, nullptr),
          SUCCEEDED(nextResult))) {
    if (nextResult == S_FALSE) {
      break;
    }
    if (!job) {
      success = false;
      break;
    }

    LPWSTR curJobName = nullptr;

    if (SUCCEEDED(job->GetDisplayName(&curJobName)) && curJobName) {
      if (lstrcmpW(curJobName, matchJobName) == 0) {
        if (!SUCCEEDED(job->Cancel())) {
          // If we can't cancel we can still try the other jobs.
          success = false;
        }
      }
      CoTaskMemFree((LPVOID)curJobName);
      curJobName = nullptr;
    } else {
      // We may not be able to access certain jobs, keep trying the rest.
      success = false;
    }

    job->Release();
    job = nullptr;
  }

  if (!SUCCEEDED(nextResult)) {
    success = false;
  }

  return success;
}

// ***** Section: DLL entry points
extern "C" {
// Cancel all BITS jobs with the given name.
void __declspec(dllexport)
    CancelBitsJobsByName(HWND hwndParent, int string_size, char *variables,
                         stack_t **stacktop, extra_parameters *) {
  g_stacktop = stacktop;
  g_stringsize = string_size;

  WCHAR matchJobName[MAX_JOB_NAME + 1];
  matchJobName[0] = L'\0';

  if (!popstringn(matchJobName, sizeof(matchJobName) / sizeof(WCHAR))) {
    if (CancelBitsJobsByNameImpl(matchJobName)) {
      pushstring(L"ok");
      return;
    }
  }

  pushstring(L"error");
}

// Start the BITS service in the background, and hold a reference to it until
// the (un)installer exits.
// Does not provide any feedback or touch the stack.
void __declspec(dllexport)
    StartBitsServiceBackground(HWND, int, char *, stack_t **,
                               extra_parameters *extra_params) {
  StartBitsServiceBackgroundThreadImpl(extra_params);
}
}

// Handle messages from NSIS
UINT_PTR __cdecl NSISPluginCallback(NSPIM msg) {
  if (msg == NSPIM_UNLOAD) {
    ShutdownStartBitsThread();
  }
  return 0;
}

BOOL APIENTRY DllMain(HINSTANCE instance, DWORD reason, LPVOID) {
  if (reason == DLL_PROCESS_ATTACH) {
    gHInst = instance;
    InitializeConditionVariable(&gStartBitsThread.cv);
    InitializeCriticalSection(&gStartBitsThread.cs);
  }
  return TRUE;
}