Allow appending to the service environment.

[nssm.git] / registry.cpp

diff --git a/registry.cpp b/registry.cpp
index 169cd08..382cb0f 100644 (file)
--- a/registry.cpp
+++ b/registry.cpp
@@ -113,42 +113,43 @@ int create_exit_action(char *service_name, const char *action_string) {
   return 0;\r
 }\r
 \r
-int set_environment(char *service_name, HKEY key, char **env) {\r
+int set_environment(char *service_name, HKEY key, char *value, char **env, unsigned long *envlen) {\r
   unsigned long type = REG_MULTI_SZ;\r
-  unsigned long envlen = 0;\r
 \r
   /* Dummy test to find buffer size */\r
-  unsigned long ret = RegQueryValueEx(key, NSSM_REG_ENV, 0, &type, NULL, &envlen);\r
+  unsigned long ret = RegQueryValueEx(key, value, 0, &type, NULL, envlen);\r
   if (ret != ERROR_SUCCESS) {\r
+    *envlen = 0;\r
     /* The service probably doesn't have any environment configured */\r
     if (ret == ERROR_FILE_NOT_FOUND) return 0;\r
-    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_QUERYVALUE_FAILED, NSSM_REG_ENV, error_string(GetLastError()), 0);\r
+    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_QUERYVALUE_FAILED, value, error_string(GetLastError()), 0);\r
     return 1;\r
   }\r
 \r
   if (type != REG_MULTI_SZ) {\r
-    log_event(EVENTLOG_WARNING_TYPE, NSSM_EVENT_INVALID_ENVIRONMENT_STRING_TYPE, NSSM_REG_ENV, service_name, 0);\r
+    log_event(EVENTLOG_WARNING_TYPE, NSSM_EVENT_INVALID_ENVIRONMENT_STRING_TYPE, value, service_name, 0);\r
     return 2;\r
   }\r
 \r
   /* Probably not possible */\r
-  if (! envlen) return 0;\r
+  if (! *envlen) return 0;\r
 \r
   /* Previously initialised? */\r
   if (*env) HeapFree(GetProcessHeap(), 0, *env);\r
 \r
-  *env = (char *) HeapAlloc(GetProcessHeap(), 0, envlen);\r
+  *env = (char *) HeapAlloc(GetProcessHeap(), 0, *envlen);\r
   if (! *env) {\r
-    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_OUT_OF_MEMORY, "environment registry", "set_environment()", 0);\r
+    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_OUT_OF_MEMORY, value, "set_environment()", 0);\r
     return 3;\r
   }\r
 \r
   /* Actually get the strings */\r
-  ret = RegQueryValueEx(key, NSSM_REG_ENV, 0, &type, (unsigned char *) *env, &envlen);\r
+  ret = RegQueryValueEx(key, value, 0, &type, (unsigned char *) *env, envlen);\r
   if (ret != ERROR_SUCCESS) {\r
+    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_QUERYVALUE_FAILED, value, error_string(GetLastError()), 0);\r
     HeapFree(GetProcessHeap(), 0, *env);\r
     *env = 0;\r
-    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_QUERYVALUE_FAILED, NSSM_REG_ENV, error_string(GetLastError()), 0);\r
+    *envlen = 0;\r
     return 4;\r
   }\r
 \r
@@ -326,7 +327,42 @@ int get_parameters(nssm_service_t *service, STARTUPINFO *si) {
   }\r
 \r
   /* Try to get environment variables - may fail */\r
-  set_environment(service->name, key, &service->env);\r
+  set_environment(service->name, key, NSSM_REG_ENV, &service->env, &service->envlen);\r
+  /* Environment variables to add to existing rather than replace - may fail. */\r
+  set_environment(service->name, key, NSSM_REG_ENV_EXTRA, &service->env_extra, &service->env_extralen);\r
+\r
+  if (service->env_extra) {\r
+    /* Append these to any other environment variables set. */\r
+    if (service->env) {\r
+      /* Append extra variables to configured variables. */\r
+      unsigned long envlen = service->envlen + service->env_extralen - 1;\r
+      char *env = (char *) HeapAlloc(GetProcessHeap(), 0, envlen);\r
+      if (env) {\r
+        memmove(env, service->env, service->envlen - 1);\r
+        memmove(env + service->envlen - 1, service->env_extra, service->env_extralen);\r
+\r
+        HeapFree(GetProcessHeap(), 0, service->env);\r
+        service->env = env;\r
+        service->envlen = envlen;\r
+      }\r
+      else log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_OUT_OF_MEMORY, "environment", "get_parameters()", 0);\r
+    }\r
+    else {\r
+      /* Append extra variables to our environment. */\r
+      char *env, *s;\r
+      size_t envlen, len;\r
+\r
+      env = service->env_extra;\r
+      len = 0;\r
+      while (*env) {\r
+        envlen = strlen(env) + 1;\r
+        for (s = env; *s && *s != '='; s++);\r
+        if (*s == '=') *s++ = '\0';\r
+        if (! SetEnvironmentVariable(env, s)) log_event(EVENTLOG_WARNING_TYPE, NSSM_EVENT_SETENVIRONMENTVARIABLE_FAILED, env, s, error_string(GetLastError()));\r
+        env += envlen;\r
+      }\r
+    }\r
+  }\r
 \r
   /* Try to get stdout and stderr */\r
   if (get_output_handles(key, si)) {\r