Harden CHEF_POWERSHELL_BIN/architecture resolution and fix a native assembly-resolve crash

Chef.PowerShell.Wrapper/Chef.PowerShell.Wrapper.cpp

@@ -13,21 +13,39 @@ using namespace System::Reflection;
 // for assemblies.
 Assembly^ currentDomain_AssemblyResolve(Object^ sender, ResolveEventArgs^ args)
 {
+    AssemblyName^ name = gcnew AssemblyName(args->Name);
+
+    // Prefer CHEF_POWERSHELL_BIN when it actually has the assembly, so operators can
+    // still point at a different/updated DLL set without moving files around.
     String^ prefix = Environment::GetEnvironmentVariable("CHEF_POWERSHELL_BIN");
     if (prefix) {
         try
         {
-            AssemblyName^ name = gcnew AssemblyName(args->Name);
             String^ finalPath = Path::Combine(prefix, name->Name + ".dll");
-            Assembly^ retval = Assembly::LoadFrom(finalPath);
-            return retval;
+            return Assembly::LoadFrom(finalPath);
         }
         catch (FileNotFoundException^)
         {
+            // Fall through -- CHEF_POWERSHELL_BIN may be stale or unrelated to wherever
+            // this wrapper assembly actually got loaded from (see fallback below).
         }
     }
-    else if (args->Name->ToLower()->StartsWith("chef.powershell")) {
-        throw gcnew FileNotFoundException("Unable to load Chef.Powershell.dll. Make sure the file is located in the same directory as ruby.exe or in CHEF_POWERSHELL_BIN.");
+
+    // Fall back to the directory this wrapper assembly itself was loaded from. Without
+    // this, an unset/stale/wrong CHEF_POWERSHELL_BIN breaks assembly resolution even
+    // though Chef.PowerShell.dll is sitting right next to this wrapper DLL.
+    try
+    {
+        String^ ownDir = Path::GetDirectoryName(Assembly::GetExecutingAssembly()->Location);
+        String^ finalPath = Path::Combine(ownDir, name->Name + ".dll");
+        return Assembly::LoadFrom(finalPath);
+    }
+    catch (FileNotFoundException^)
+    {
+    }
+
+    if (name->Name->ToLower()->StartsWith("chef.powershell")) {
+        throw gcnew FileNotFoundException("Unable to load " + name->Name + ".dll. Make sure the file is located in the same directory as this wrapper assembly or in CHEF_POWERSHELL_BIN.");
     }
 
     return nullptr;

chef-powershell/ext/chef-powershell/extconf.rb

@@ -59,7 +59,7 @@ def find_source_dlls
     # Look for DLLs in the gem's bin directory structure
     # They should be in bin/ruby_bin_folder/<ARCH>/
     gem_bin_dir = File.expand_path("../../bin", __dir__)
-    arch = ENV["PROCESSOR_ARCHITECTURE"] || "AMD64"
+    arch = ENV["PROCESSOR_ARCHITECTURE"].to_s.empty? ? "AMD64" : ENV["PROCESSOR_ARCHITECTURE"]
 
     possible_locations = [
       File.join(gem_bin_dir, "ruby_bin_folder", arch),
@@ -109,7 +109,7 @@ def install_dlls
     source_dir = find_source_dlls
 
     unless source_dir
-      arch = ENV["PROCESSOR_ARCHITECTURE"] || "AMD64"
+      arch = ENV["PROCESSOR_ARCHITECTURE"].to_s.empty? ? "AMD64" : ENV["PROCESSOR_ARCHITECTURE"]
       log "WARNING: Could not locate source DLL files in gem installation."
       log "Expected location: bin/ruby_bin_folder/#{arch}/"
       log "DLL files will not be installed."

chef-powershell/lib/chef-powershell/powershell.rb

@@ -48,7 +48,7 @@ class PowerShell
     attr_reader :verbose
 
     def self.resolve_wrapper_dll
-      arch = ENV["PROCESSOR_ARCHITECTURE"] || "AMD64"
+      arch = ENV["PROCESSOR_ARCHITECTURE"].to_s.empty? ? "AMD64" : ENV["PROCESSOR_ARCHITECTURE"]
       searched = []
 
       # Gem path

chef-powershell/lib/chef-powershell/powershell_exec.rb

@@ -102,22 +102,26 @@ module PowerShellExec
       # We don't want to move files around so we're setting the variable here to keep everything tidy.
       # file_path = Gem.loaded_specs["chef-powershell"].full_gem_path + "/bin/ruby_bin_folder/#{ENV["PROCESSOR_ARCHITECTURE"]}/"
       # ENV["CHEF_POWERSHELL_BIN"] = file_path
-      gem_spec = Gem.loaded_specs["chef-powershell"]
-      if gem_spec
-        arch = ENV["PROCESSOR_ARCHITECTURE"] || "AMD64"
-        file_path = File.join(gem_spec.full_gem_path, "bin", "ruby_bin_folder", arch)
-        if File.directory?(file_path)
-          ENV["CHEF_POWERSHELL_BIN"] = file_path
-        else
-          # Habitat fallback: try locating the built package and use its bin directory
-          begin
-            hab_path = `hab pkg path chef/chef-powershell-shim 2>NUL`.strip
-            if hab_path != "" && File.directory?(hab_path)
-              candidate = File.join(hab_path, "bin")
-              ENV["CHEF_POWERSHELL_BIN"] = candidate if File.directory?(candidate)
+      # Respect an explicit override (e.g. pointing at a freshly built Habitat package) --
+      # don't clobber it just because the loaded gem's own bundled DLLs also happen to exist.
+      if ENV["CHEF_POWERSHELL_BIN"].to_s.empty?
+        gem_spec = Gem.loaded_specs["chef-powershell"]
+        if gem_spec
+          arch = ENV["PROCESSOR_ARCHITECTURE"].to_s.empty? ? "AMD64" : ENV["PROCESSOR_ARCHITECTURE"]
+          file_path = File.join(gem_spec.full_gem_path, "bin", "ruby_bin_folder", arch)
+          if File.directory?(file_path)
+            ENV["CHEF_POWERSHELL_BIN"] = file_path
+          else
+            # Habitat fallback: try locating the built package and use its bin directory
+            begin
+              hab_path = `hab pkg path chef/chef-powershell-shim 2>NUL`.strip
+              if hab_path != "" && File.directory?(hab_path)
+                candidate = File.join(hab_path, "bin")
+                ENV["CHEF_POWERSHELL_BIN"] = candidate if File.directory?(candidate)
+              end
+            rescue
+              # ignore habitat lookup failures
             end
-          rescue
-            # ignore habitat lookup failures
           end
         end
       end

chef-powershell/lib/chef-powershell/pwsh.rb

@@ -20,7 +20,7 @@ class Pwsh < ChefPowerShell::PowerShell
 
     def self.resolve_core_wrapper_dll
       gem_spec = Gem.loaded_specs["chef-powershell"]
-      arch = ENV["PROCESSOR_ARCHITECTURE"] || "AMD64"
+      arch = ENV["PROCESSOR_ARCHITECTURE"].to_s.empty? ? "AMD64" : ENV["PROCESSOR_ARCHITECTURE"]
 
       if gem_spec
         base = gem_spec.full_gem_path
@@ -66,7 +66,7 @@ def exec(script, timeout: -1)
       ENV["DOTNET_MULTILEVEL_LOOKUP"] = "0"
       gem_spec = Gem.loaded_specs["chef-powershell"]
       if gem_spec
-        arch = ENV["PROCESSOR_ARCHITECTURE"] || "AMD64"
+        arch = ENV["PROCESSOR_ARCHITECTURE"].to_s.empty? ? "AMD64" : ENV["PROCESSOR_ARCHITECTURE"]
         arch_root = File.join(gem_spec.full_gem_path, "bin", "ruby_bin_folder", arch)
         ENV["DOTNET_ROOT"] = arch_root
       end

chef-powershell/spec/integration/dll_integration_spec.rb

@@ -18,9 +18,9 @@
 require "open3"
 require "tempfile"
 
-# Snapshot CHEF_POWERSHELL_BIN *before* require "chef-powershell", because
-# powershell_exec.rb unconditionally overwrites it with the gem's own bin path
-# at module load time.
+# Snapshot CHEF_POWERSHELL_BIN *before* require "chef-powershell". powershell_exec.rb
+# no longer clobbers an existing override at module load time, but this snapshot is
+# kept as a belt-and-suspenders guard against a regression of that behavior.
 # rubocop:disable Lint/UnderscorePrefixedVariableName
 _hab_bin_override = ENV["CHEF_POWERSHELL_BIN"].dup if ENV["CHEF_POWERSHELL_BIN"] && !ENV["CHEF_POWERSHELL_BIN"].empty?
 # rubocop:enable Lint/UnderscorePrefixedVariableName
@@ -42,7 +42,7 @@
 
 # Re-pin CHEF_POWERSHELL_BIN to the Hab package dir so the C++/CLI assembly
 # resolver (currentDomain_AssemblyResolve in Wrapper.cpp) finds Chef.PowerShell.dll
-# in the right place. powershell_exec.rb overwrote it with the gem's own path.
+# in the right place.
 ENV["CHEF_POWERSHELL_BIN"] = DLL_BIN_DIR
 
 # ---------------------------------------------------------------------------
@@ -456,4 +456,33 @@ def raw_pwsh(script, timeout: -1)
   ensure
     ENV["CHEF_POWERSHELL_BIN"] = orig_bin
   end
+
+  # Regression coverage for a real bug: powershell_exec.rb used to unconditionally
+  # overwrite CHEF_POWERSHELL_BIN with the loaded gem's own bin path at require time,
+  # silently clobbering an explicit override (e.g. pointing at a freshly built Habitat
+  # package) whenever the gem's own bundled ruby_bin_folder also happened to exist.
+  it "does not clobber an explicit CHEF_POWERSHELL_BIN override when requiring chef-powershell" do
+    gem_root = File.expand_path("../..", __dir__)
+    sentinel = File.join("C:", "explicit_chef_powershell_bin_override")
+
+    script = <<~'RUBY'
+      require "chef-powershell"
+      puts ENV["CHEF_POWERSHELL_BIN"]
+    RUBY
+
+    env = { "CHEF_POWERSHELL_BIN" => sentinel }
+
+    stdout, stderr, status = Tempfile.create(["chef_powershell_bin_override_check", ".rb"]) do |file|
+      file.write(script)
+      file.close
+
+      Bundler.with_unbundled_env do
+        Open3.capture3(env, "bundle", "exec", "ruby", file.path, chdir: gem_root)
+      end
+    end
+
+    expect(status.success?).to be(true),
+      "subprocess failed (exit #{status.exitstatus}):\nSTDOUT:\n#{stdout}\nSTDERR:\n#{stderr}"
+    expect(stdout.strip).to eq(sentinel)
+  end
 end