Skip to content

Commit

Permalink
fix: env var loading order in WindowsExec (#1136)
Browse files Browse the repository at this point in the history
  • Loading branch information
tilo-chen authored and MikeDombo committed Nov 23, 2021
1 parent 2e9e326 commit ca80f96
Show file tree
Hide file tree
Showing 5 changed files with 156 additions and 45 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,113 @@
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/

package com.aws.greengrass.integrationtests.util;

import com.aws.greengrass.integrationtests.BaseITCase;
import com.aws.greengrass.util.Exec;
import com.aws.greengrass.util.platforms.Platform;
import com.sun.jna.platform.win32.Advapi32Util;
import org.apache.commons.lang3.SystemUtils;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledOnOs;
import org.junit.jupiter.api.condition.OS;

import java.io.IOException;
import java.util.concurrent.TimeUnit;

import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.containsString;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.fail;

public class ExecIntegTest extends BaseITCase {

private static final String TEST_ENV_ENTRY = "integ_test_env";
private static final String SYSTEM_ENV_REG_KEY = "HKLM\\SYSTEM\\CurrentControlSet\\Control\\Session Manager\\Environment";

private static String testUserEnvRegKey;

@AfterAll
static void cleanup() throws IOException, InterruptedException {
if (SystemUtils.IS_OS_WINDOWS) {
deleteRegistryEntry(SYSTEM_ENV_REG_KEY, TEST_ENV_ENTRY);
if (testUserEnvRegKey != null) {
deleteRegistryEntry(testUserEnvRegKey, TEST_ENV_ENTRY);
}
}
}

private static void deleteRegistryEntry(String key, String entry) throws IOException, InterruptedException {
// reg command reference https://docs.microsoft.com/en-us/windows-server/administration/windows-commands/reg
Process p = new ProcessBuilder("reg", "delete", key, "/f", "/v", entry).start();
if (!p.waitFor(20, TimeUnit.SECONDS)) {
p.destroyForcibly();
fail("delete registry entry timeout");
}
}

@Test
@EnabledOnOs(OS.WINDOWS)
void GIVEN_windows_exec_with_user_WHEN_set_env_vars_from_multiple_sources_THEN_precedence_is_correct()
throws IOException, InterruptedException {
// Check setting default env works
// setDefaultEnv is static and will persist throughout this test.
String expectedVal = "Exec default";
Exec.setDefaultEnv(TEST_ENV_ENTRY, expectedVal);
try (Exec exec = Platform.getInstance().createNewProcessRunner()) {
String output = exec.cd("C:\\") // cd in case test user doesn't have permission to access current directory
.withUser(WINDOWS_TEST_UESRNAME)
.withShell("echo", "%" + TEST_ENV_ENTRY + "%")
.execAndGetStringOutput();
assertEquals(expectedVal, output);
}

// Set system-level env var via registry. should override the default
expectedVal = "system env var";
try (Exec exec = Platform.getInstance().createNewProcessRunner()) {
String output = exec
.withShell("reg", "add", SYSTEM_ENV_REG_KEY,
"/f", "/v", TEST_ENV_ENTRY, "/t", "REG_SZ", "/d", expectedVal)
.execAndGetStringOutput();
assertThat(output, containsString("completed successfully"));
}

try (Exec exec = Platform.getInstance().createNewProcessRunner()) {
String output = exec.cd("C:\\").withUser(WINDOWS_TEST_UESRNAME)
.withShell("echo", "%" + TEST_ENV_ENTRY + "%").execAndGetStringOutput();
assertEquals(expectedVal, output);
}

// Set user-level env var via registry. should override the system-level setting
// FYI: Usually, user-level setting has higher precedence than system-level. PATH is special.
// User PATH is appended to system-level PATH by windows.
expectedVal = "user account env var";
String testUserSid = Advapi32Util.getAccountByName(WINDOWS_TEST_UESRNAME).sidString;
testUserEnvRegKey = String.format("HKU\\%s\\Environment", testUserSid);
try (Exec exec = Platform.getInstance().createNewProcessRunner()) {
String output = exec
.cd("C:\\")
.withUser(WINDOWS_TEST_UESRNAME)
.withShell("reg", "add", testUserEnvRegKey, "/f", "/v", TEST_ENV_ENTRY, "/t", "REG_SZ", "/d", expectedVal)
.execAndGetStringOutput();
assertThat(output, containsString("completed successfully"));
}

try (Exec exec = Platform.getInstance().createNewProcessRunner()) {
String output = exec.cd("C:\\").withUser(WINDOWS_TEST_UESRNAME)
.withShell("echo", "%" + TEST_ENV_ENTRY + "%").execAndGetStringOutput();
assertEquals(expectedVal, output);
}

// Use setenv, which overrides everything
expectedVal = "setenv override";
try (Exec exec = Platform.getInstance().createNewProcessRunner()) {
String output = exec.cd("C:\\").withUser(WINDOWS_TEST_UESRNAME).setenv(TEST_ENV_ENTRY, expectedVal)
.withShell("echo", "%" + TEST_ENV_ENTRY + "%").execAndGetStringOutput();
assertEquals(expectedVal, output);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -6,17 +6,12 @@
package com.aws.greengrass.integrationtests.util;

import com.aws.greengrass.config.PlatformResolver;
import com.aws.greengrass.lifecyclemanager.Kernel;
import com.aws.greengrass.lifecyclemanager.KernelAlternatives;
import com.aws.greengrass.util.Exec;
import com.aws.greengrass.util.platforms.Platform;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.DisabledOnOs;
import org.junit.jupiter.api.condition.EnabledOnOs;
import org.junit.jupiter.api.condition.OS;
import org.junit.jupiter.api.io.TempDir;

import java.io.File;
import java.io.IOException;
Expand All @@ -37,29 +32,10 @@
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.Mockito.lenient;
import static org.mockito.Mockito.mock;

class ExecTest {

@TempDir
protected Path tempDir;
private static Kernel kernel;

@BeforeAll
static void setup() {
KernelAlternatives kernelAlts = mock(KernelAlternatives.class);
lenient().when(kernelAlts.getBinDir()).thenReturn(Paths.get("scripts"));
kernel = new Kernel();
kernel.getContext().put(KernelAlternatives.class, kernelAlts);
}

@AfterAll
static void cleanup() {
kernel.shutdown();
}

private String readLink(String path) throws IOException {
private static String readLink(String path) throws IOException {
Path p = Paths.get(path);
if (Files.isSymbolicLink(p)) {
return Files.readSymbolicLink(p).toString();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,6 @@ public class WindowsExec extends Exec {
super();
// Windows env var keys are case-insensitive. Use case-insensitive map to avoid collision
environment = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
environment.putAll(defaultEnvironment);
String pathExt = System.getenv("PATHEXT");
pathext = Arrays.asList(pathExt.split(File.pathSeparator));
}
Expand Down Expand Up @@ -120,6 +119,7 @@ protected Process createProcess() throws IOException {
// Clear the env in this case because later we'll load the given user's env instead
winPb.environment().clear();
}
winPb.setDefaultEnvironment(defaultEnvironment);
winPb.environment().putAll(environment);
winPb.directory(dir).command(commands);
synchronized (Kernel32.INSTANCE) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
package vendored.org.apache.dolphinscheduler.common.utils.process;

import com.sun.jna.platform.win32.WinBase;
import lombok.Setter;

import java.io.File;
import java.io.IOException;
Expand Down Expand Up @@ -182,6 +183,9 @@ public class ProcessBuilderForWin32 {
* Begin Amazon addition.
*/

// Separate out the defaultEnvironment so that we can control env var precedence later in ProcessImplForWin32
@Setter
private Map<String,String> defaultEnvironment;
private static final int PROCESS_CREATION_FLAGS_DEFAULT = WinBase.CREATE_UNICODE_ENVIRONMENT // use unicode
| WinBase.CREATE_NEW_CONSOLE; // create new console to send ctrl-c to the process
private int processCreationFlags = PROCESS_CREATION_FLAGS_DEFAULT;
Expand Down Expand Up @@ -1111,6 +1115,7 @@ public Process start() throws IOException {
username,
password,
cmdarray,
defaultEnvironment,
environment,
dir,
redirects,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -172,7 +172,8 @@ public FileOutputStream run() {
static Process start(String username,
char[] password,
String[] cmdarray,
java.util.Map<String,String> envMap,
java.util.Map<String,String> defaultEnv,
java.util.Map<String,String> overrideEnv,
String dir,
ProcessBuilderForWin32.Redirect[] redirects,
boolean redirectErrorStream,
Expand Down Expand Up @@ -220,7 +221,8 @@ static Process start(String username,
}
}

return new ProcessImplForWin32(username, password, cmdarray, envMap, dir, stdHandles, redirectErrorStream, processCreationFlags);
return new ProcessImplForWin32(username, password, cmdarray, defaultEnv, overrideEnv, dir, stdHandles,
redirectErrorStream, processCreationFlags);
} finally {
// In theory, close() can throw IOException
// (although it is rather unlikely to happen here)
Expand Down Expand Up @@ -467,7 +469,8 @@ private ProcessImplForWin32(
String username,
char[] password,
String[] cmd,
final java.util.Map<String,String> envMap,
final java.util.Map<String,String> defaultEnv,
final java.util.Map<String,String> overrideEnv,
final String path,
final long[] stdHandles,
final boolean redirectErrorStream,
Expand Down Expand Up @@ -534,7 +537,8 @@ private ProcessImplForWin32(
cmd);
}

handle = create(username, password, cmdstr, envMap, path, stdHandles, redirectErrorStream, processCreationFlags);
handle = create(username, password, cmdstr, defaultEnv, overrideEnv, path, stdHandles, redirectErrorStream,
processCreationFlags);

AccessController.doPrivileged(
new PrivilegedAction<Void>() {
Expand Down Expand Up @@ -828,20 +832,28 @@ private WinNT.HANDLE processCreate(String username,
private synchronized WinNT.HANDLE create(String username,
char[] password,
String cmd,
java.util.Map<String,String> envMap,
java.util.Map<String,String> defaultEnv,
java.util.Map<String,String> overrideEnv,
final String path,
final long[] stdHandles,
final boolean redirectErrorStream,
final int processCreationFlags) throws ProcessCreationException {
String envblock;
ProcessCreationExtras extraInfo = new ProcessCreationExtras();
if (envMap == null) {
envMap = Collections.emptyMap();
if (defaultEnv == null) {
defaultEnv = Collections.emptyMap();
}
if (overrideEnv == null) {
overrideEnv = Collections.emptyMap();
}
if (username == null) {
envblock = Advapi32Util.getEnvironmentBlock(envMap);
// Windows env var keys are case-insensitive. Use case-insensitive map to avoid collision
Map<String, String> mergedEnv = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
mergedEnv.putAll(defaultEnv);
mergedEnv.putAll(overrideEnv);
envblock = Advapi32Util.getEnvironmentBlock(mergedEnv);
} else {
envblock = setupRunAsAnotherUser(username, password, envMap, extraInfo);
envblock = setupRunAsAnotherUser(username, password, defaultEnv, overrideEnv, extraInfo);
}

// init handles
Expand Down Expand Up @@ -951,8 +963,9 @@ private static void waitForTimeoutInterruptibly(WinNT.HANDLE handle, long timeou
* Preparation for running process as another user. Populates the given extraInfo if applicable.
* @return environment block for CreateProcess* APIs
*/
private String setupRunAsAnotherUser(String username, char[] password, Map<String, String> envMap,
ProcessCreationExtras extraInfo) throws ProcessCreationException {
private String setupRunAsAnotherUser(String username, char[] password, Map<String, String> defaultEnv,
Map<String, String> overrideEnv, ProcessCreationExtras extraInfo)
throws ProcessCreationException {
// Get and cache process token and isService state
synchronized (Advapi32.INSTANCE) {
if (processToken.get() == null) {
Expand Down Expand Up @@ -1028,7 +1041,7 @@ private String setupRunAsAnotherUser(String username, char[] password, Map<Strin
lastErrorRuntimeException());
}

String envblock = computeEnvironmentBlock(userTokenHandle.getValue(), envMap);
String envblock = computeEnvironmentBlock(userTokenHandle.getValue(), defaultEnv, overrideEnv);
Kernel32Util.closeHandleRefs(userTokenHandle);
return envblock;
}
Expand Down Expand Up @@ -1098,8 +1111,8 @@ private static boolean checkIsService(WinNT.HANDLEByReference processToken) thro
* @return environment block for starting a process
* @see <a href="https://docs.microsoft.com/en-us/windows/win32/api/userenv/nf-userenv-createenvironmentblock">docs</a>
*/
private static String computeEnvironmentBlock(WinNT.HANDLE userTokenHandle, Map<String, String> additionalEnv)
throws ProcessCreationException {
private static String computeEnvironmentBlock(WinNT.HANDLE userTokenHandle, Map<String, String> defaultEnv,
Map<String, String> overrideEnv) throws ProcessCreationException {
PointerByReference lpEnv = new PointerByReference();
// Get user's env vars, inheriting current process env
// It returns pointer to a block of null-terminated strings. It ends with two nulls (\0\0).
Expand All @@ -1108,7 +1121,12 @@ private static String computeEnvironmentBlock(WinNT.HANDLE userTokenHandle, Map<
}

// Windows env var keys are case-insensitive. Use case-insensitive map to avoid collision
Map<String, String> userEnvMap = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
// The resulting env is merged from defaultEnv, the given user's env (returned by CreateEnvironmentBlock),
// and overrideEnv. They are merged in that order so that later envs have higher precedence in case
// a key is defined in multiple places.
Map<String, String> mergedEnv = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
mergedEnv.putAll(defaultEnv);

int offset = 0;
while (true) {
String s = lpEnv.getValue().getWideString(offset);
Expand All @@ -1118,17 +1136,16 @@ private static String computeEnvironmentBlock(WinNT.HANDLE userTokenHandle, Map<
// wide string uses 2 bytes per char. +2 to skip the terminating null
offset += s.length() * 2 + 2;
int splitInd = s.indexOf('=');
userEnvMap.put(s.substring(0, splitInd), s.substring(splitInd + 1));
mergedEnv.put(s.substring(0, splitInd), s.substring(splitInd + 1));
}

if (!UserEnv.INSTANCE.DestroyEnvironmentBlock(lpEnv.getValue())) {
throw lastErrorProcessCreationException("DestroyEnvironmentBlock");
}

// Set additional envs on top of user default env
userEnvMap.putAll(additionalEnv);
mergedEnv.putAll(overrideEnv);

return Advapi32Util.getEnvironmentBlock(userEnvMap);
return Advapi32Util.getEnvironmentBlock(mergedEnv);
}

private static class ProcessCreationExtras {
Expand Down

0 comments on commit ca80f96

Please sign in to comment.