From e3645baeee4558cd6f1838c2d5383b7c736d4bfd Mon Sep 17 00:00:00 2001
From: Josh Cooper <joshcooper@users.noreply.github.com>
Date: Tue, 9 Jan 2024 17:48:20 -0800
Subject: [PATCH 1/2] Refactor wchar size

(cherry picked from commit 493a8b74b46140881f144fa6a6e70e0a8d9f834f)
---
 lib/puppet/util/windows/registry.rb | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/lib/puppet/util/windows/registry.rb b/lib/puppet/util/windows/registry.rb
index f37a8d848fe..4f5b74fe681 100644
--- a/lib/puppet/util/windows/registry.rb
+++ b/lib/puppet/util/windows/registry.rb
@@ -15,6 +15,8 @@ module Registry
 
     ERROR_NO_MORE_ITEMS = 259
 
+    WCHAR_SIZE = FFI.type_size(:wchar)
+
     def root(name)
       Win32::Registry.const_get(name)
     rescue NameError
@@ -235,7 +237,7 @@ def read(key, name_ptr, *rtype)
 
         string_length = 0
         # buffer is raw bytes, *not* chars - less a NULL terminator
-        string_length = (byte_length / FFI.type_size(:wchar)) - 1 if byte_length > 0
+        string_length = (byte_length / WCHAR_SIZE) - 1 if byte_length > 0
 
         begin
           case type
@@ -278,7 +280,7 @@ def query_value_ex(key, name_ptr, &block)
 
             if result != FFI::ERROR_SUCCESS
               # buffer is raw bytes, *not* chars - less a NULL terminator
-              name_length = (name_ptr.size / FFI.type_size(:wchar)) - 1 if name_ptr.size > 0
+              name_length = (name_ptr.size / WCHAR_SIZE) - 1 if name_ptr.size > 0
               msg = _("Failed to read registry value %{value} at %{key}") % { value: name_ptr.read_wide_string(name_length), key: key.keyname }
               raise Puppet::Util::Windows::Error.new(msg, result)
             end

From bcdfbf0ad4e826e2474fcd0e95ad30ed94a66f77 Mon Sep 17 00:00:00 2001
From: Josh Cooper <joshcooper@users.noreply.github.com>
Date: Tue, 9 Jan 2024 17:52:22 -0800
Subject: [PATCH 2/2] (PUP-11122) Ensure reg values are null terminated

RegQueryValueExW doesn't guarantee the returned buffer for the `lpData`
parameter is null terminated, so ensure that it is.

When retrieving a registry value of type REG_SZ or REG_EXPAND_SZ extend the
buffer by 1 wchar (2 bytes) so we can always write a wchar null terminator that
is guaranteed not to overwrite user data.

The resulting wchar array is then guaranteed to be properly wchar null
terminated:

    \u0041
    \u0042
    \u0000

which when UTF-16LE encoded will result in the byte array:

    \x41 \x00
    \x42 \x00
    \x00 \x00

If Windows does null terminate, then we will add an additional wchar null
terminator, which won't hurt anything.

The same applies to REG_MULTI_SZ, except it is supposed to have two wchar null
terminators, for a total of 4 bytes. One terminator is for the last string in
the array, and one terminator is for the entire array. For example, if the array
contains the strings ['A', 'B'] and Windows does not null terminate the
`lpData` buffer, then the resulting UTF-16LE encoded byte array will contain:

    \x41 \x00
    \x00 \x00
    \x42 \x00
    \x00 \x00
    \x00 \x00

(cherry picked from commit 25a33dfd6afc17a083a254457debdd2206a73326)
---
 lib/puppet/util/windows/registry.rb | 37 +++++++++++++++++++++++++++--
 1 file changed, 35 insertions(+), 2 deletions(-)

diff --git a/lib/puppet/util/windows/registry.rb b/lib/puppet/util/windows/registry.rb
index 4f5b74fe681..8560175598a 100644
--- a/lib/puppet/util/windows/registry.rb
+++ b/lib/puppet/util/windows/registry.rb
@@ -273,12 +273,45 @@ def query_value_ex(key, name_ptr, &block)
             FFI::Pointer::NULL, type_ptr,
             FFI::Pointer::NULL, length_ptr)
 
-          FFI::MemoryPointer.new(:byte, length_ptr.read_dword) do |buffer_ptr|
+          # The call to RegQueryValueExW below is potentially unsafe:
+          # https://learn.microsoft.com/en-us/windows/win32/api/winreg/nf-winreg-regqueryvalueexw
+          #
+          #   "If the data has the REG_SZ, REG_MULTI_SZ or REG_EXPAND_SZ type,
+          #   the string may not have been stored with the proper terminating
+          #   null characters. Therefore, even if the function returns
+          #   ERROR_SUCCESS, the application should ensure that the string is
+          #   properly terminated before using it; otherwise, it may overwrite a
+          #   buffer. (Note that REG_MULTI_SZ strings should have two
+          #   terminating null characters.)"
+          #
+          # Since we don't know if the values will be properly null terminated,
+          # extend the buffer to guarantee we can append one or two wide null
+          # characters, without overwriting any data.
+          base_bytes_len = length_ptr.read_dword
+          pad_bytes_len = case type_ptr.read_dword
+                          when Win32::Registry::REG_SZ, Win32::Registry::REG_EXPAND_SZ
+                            WCHAR_SIZE
+                          when Win32::Registry::REG_MULTI_SZ
+                            WCHAR_SIZE * 2
+                          else
+                            0
+                          end
+
+          FFI::MemoryPointer.new(:byte, base_bytes_len + pad_bytes_len) do |buffer_ptr|
             result = RegQueryValueExW(key.hkey, name_ptr,
               FFI::Pointer::NULL, type_ptr,
               buffer_ptr, length_ptr)
 
-            if result != FFI::ERROR_SUCCESS
+            # Ensure buffer is null terminated with 1 or 2 wchar nulls, depending on the type
+            if result == FFI::ERROR_SUCCESS
+              case type_ptr.read_dword
+              when Win32::Registry::REG_SZ, Win32::Registry::REG_EXPAND_SZ
+                buffer_ptr.put_uint16(base_bytes_len, 0)
+              when Win32::Registry::REG_MULTI_SZ
+                buffer_ptr.put_uint16(base_bytes_len, 0)
+                buffer_ptr.put_uint16(base_bytes_len + WCHAR_SIZE, 0)
+              end
+            else
               # buffer is raw bytes, *not* chars - less a NULL terminator
               name_length = (name_ptr.size / WCHAR_SIZE) - 1 if name_ptr.size > 0
               msg = _("Failed to read registry value %{value} at %{key}") % { value: name_ptr.read_wide_string(name_length), key: key.keyname }