imports.cpp

   1 #include "nssm.h"
   2
   3 imports_t imports;
   4
   5 /*
   6   Try to set up function pointers.
   7   In this first implementation it is not an error if we can't load them
   8   because we aren't currently trying to load any functions which we
   9   absolutely need.  If we later add some indispensible imports we can
  10   return non-zero here to force an application exit.
  11 */
  12 HMODULE get_dll(const TCHAR *dll, unsigned long *error) {
  13   *error = 0;
  14
  15   HMODULE ret = LoadLibrary(dll);
  16   if (! ret) {
  17     *error = GetLastError();
  18     log_event(EVENTLOG_WARNING_TYPE, NSSM_EVENT_LOADLIBRARY_FAILED, dll, error_string(*error), 0);
  19   }
  20
  21   return ret;
  22 }
  23
  24 FARPROC get_import(HMODULE library, const char *function, unsigned long *error) {
  25   *error = 0;
  26
  27   FARPROC ret = GetProcAddress(library, function);
  28   if (! ret) {
  29     *error = GetLastError();
  30     TCHAR *function_name;
  31 #ifdef UNICODE
  32     size_t buflen;
  33     mbstowcs_s(&buflen, NULL, 0, function, _TRUNCATE);
  34     function_name = (TCHAR *) HeapAlloc(GetProcessHeap(), 0, buflen * sizeof(TCHAR));
  35     if (function_name) mbstowcs_s(&buflen, function_name, buflen * sizeof(TCHAR), function, _TRUNCATE);
  36 #else
  37     function_name = (TCHAR *) function;
  38 #endif
  39     if (*error != ERROR_PROC_NOT_FOUND) log_event(EVENTLOG_WARNING_TYPE, NSSM_EVENT_GETPROCADDRESS_FAILED, function_name, error_string(*error), 0);
  40 #ifdef UNICODE
  41     if (function_name) HeapFree(GetProcessHeap(), 0, function_name);
  42 #endif
  43   }
  44
  45   return ret;
  46 }
  47
  48 int get_imports() {
  49   unsigned long error;
  50
  51   ZeroMemory(&imports, sizeof(imports));
  52
  53   imports.kernel32 = get_dll(_T("kernel32.dll"), &error);
  54   if (imports.kernel32) {
  55     imports.AttachConsole = (AttachConsole_ptr) get_import(imports.kernel32, "AttachConsole", &error);
  56     if (! imports.AttachConsole) {
  57       if (error != ERROR_PROC_NOT_FOUND) return 2;
  58     }
  59
  60     imports.SleepConditionVariableCS = (SleepConditionVariableCS_ptr) get_import(imports.kernel32, "SleepConditionVariableCS", &error);
  61     if (! imports.SleepConditionVariableCS) {
  62       if (error != ERROR_PROC_NOT_FOUND) return 3;
  63     }
  64
  65     imports.WakeConditionVariable = (WakeConditionVariable_ptr) get_import(imports.kernel32, "WakeConditionVariable", &error);
  66     if (! imports.WakeConditionVariable) {
  67       if (error != ERROR_PROC_NOT_FOUND) return 4;
  68     }
  69   }
  70   else if (error != ERROR_MOD_NOT_FOUND) return 1;
  71
  72   imports.advapi32 = get_dll(_T("advapi32.dll"), &error);
  73   if (imports.advapi32) {
  74     imports.CreateWellKnownSid = (CreateWellKnownSid_ptr) get_import(imports.advapi32, "CreateWellKnownSid", &error);
  75     if (! imports.CreateWellKnownSid) {
  76       if (error != ERROR_PROC_NOT_FOUND) return 6;
  77     }
  78     imports.IsWellKnownSid = (IsWellKnownSid_ptr) get_import(imports.advapi32, "IsWellKnownSid", &error);
  79     if (! imports.IsWellKnownSid) {
  80       if (error != ERROR_PROC_NOT_FOUND) return 7;
  81     }
  82   }
  83   else if (error != ERROR_MOD_NOT_FOUND) return 5;
  84
  85   imports.user32 = get_dll(_T("user32.dll"), &error);
  86   if (imports.user32) {
  87     imports.RegisterPowerSettingNotification = (RegisterPowerSettingNotification_ptr) get_import(imports.user32, "RegisterPowerSettingNotification", &error);
  88     if (! imports.RegisterPowerSettingNotification) {
  89       if (error != ERROR_PROC_NOT_FOUND) return 9;
  90     }
  91   }
  92   else if (error != ERROR_MOD_NOT_FOUND) return 8;
  93
  94   return 0;
  95 }
  96
  97 void free_imports() {
  98   if (imports.kernel32) FreeLibrary(imports.kernel32);
  99   if (imports.advapi32) FreeLibrary(imports.advapi32);
 100   ZeroMemory(&imports, sizeof(imports));
 101 }