From 378ce9fd72a15da9e20c9c01f41dcf47337e6cda Mon Sep 17 00:00:00 2001
From: Travis Plunk <travis.plunk@microsoft.com>
Date: Thu, 18 Aug 2022 09:47:09 -0700
Subject: [PATCH] Make Register Microsoft Update timeout (#17910)

* Make Register MU timeout

* Update Product.wxs

* Update powershell-win-core.csproj

* Rename src/PowerShell.Core.Instrumentation/RegisterMicrosoftUpdate.ps1 to assets/MicrosoftUpdate/RegisterMicrosoftUpdate.ps1

* Update src/powershell-win-core/powershell-win-core.csproj

* Update RegisterMicrosoftUpdate.ps1

* Update RegisterMicrosoftUpdate.ps1

* Update files.wxs

* Update assets/MicrosoftUpdate/RegisterMicrosoftUpdate.ps1

Co-authored-by: Aditya Patwardhan <adityap@microsoft.com>

* Update assets/MicrosoftUpdate/RegisterMicrosoftUpdate.ps1

Co-authored-by: Aditya Patwardhan <adityap@microsoft.com>

* Update RegisterMicrosoftUpdate.ps1

* Update RegisterMicrosoftUpdate.ps1

* Update RegisterMicrosoftUpdate.ps1

* Update RegisterMicrosoftUpdate.ps1

* Update assets/MicrosoftUpdate/RegisterMicrosoftUpdate.ps1

Co-authored-by: Aditya Patwardhan <adityap@microsoft.com>
---
 .../RegisterMicrosoftUpdate.ps1               | 70 +++++++++++++++++++
 assets/wix/Product.wxs                        |  2 +-
 assets/wix/files.wxs                          |  4 ++
 .../powershell-win-core.csproj                |  2 +-
 4 files changed, 76 insertions(+), 2 deletions(-)
 create mode 100644 assets/MicrosoftUpdate/RegisterMicrosoftUpdate.ps1

diff --git a/assets/MicrosoftUpdate/RegisterMicrosoftUpdate.ps1 b/assets/MicrosoftUpdate/RegisterMicrosoftUpdate.ps1
new file mode 100644
index 000000000..fd341d4fb
--- /dev/null
+++ b/assets/MicrosoftUpdate/RegisterMicrosoftUpdate.ps1
@@ -0,0 +1,70 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+param(
+    [ValidateSet('Hang', 'Fail')]
+    $TestHook
+)
+
+$waitTimeoutSeconds = 300
+switch ($TestHook) {
+    'Hang' {
+        $waitTimeoutSeconds = 10
+        $jobScript = { Start-Sleep -Seconds 600 }
+    }
+    'Fail' {
+        $jobScript = { throw "This job script should fail" }
+    }
+    default {
+        $jobScript = {
+            # This registers Microsoft Update via a predifened GUID with the Windows Update Agent.
+            # https://docs.microsoft.com/en-us/windows/win32/wua_sdk/opt-in-to-microsoft-update
+
+            $serviceManager = (New-Object -ComObject Microsoft.Update.ServiceManager)
+            $isRegistered = $serviceManager.QueryServiceRegistration('7971f918-a847-4430-9279-4a52d1efe18d').Service.IsRegisteredWithAu
+
+            if (!$isRegistered) {
+                Write-Verbose -Verbose "Opting into Microsoft Update as the Autmatic Update Service"
+                # 7 is the combination of asfAllowPendingRegistration, asfAllowOnlineRegistration, asfRegisterServiceWithAU
+                # AU means Automatic Updates
+                $null = $serviceManager.AddService2('7971f918-a847-4430-9279-4a52d1efe18d', 7, '')
+            }
+            else {
+                Write-Verbose -Verbose "Microsoft Update is already registered for Automatic Updates"
+            }
+
+            $isRegistered = $serviceManager.QueryServiceRegistration('7971f918-a847-4430-9279-4a52d1efe18d').Service.IsRegisteredWithAu
+
+            # Return if it was successful, which is the opposite of Pending.
+            return $isRegistered
+        }
+    }
+}
+
+Write-Verbose "Running job script: $jobScript" -Verbose
+$job = Start-ThreadJob -ScriptBlock $jobScript
+
+Write-Verbose "Waiting on Job for $waitTimeoutSeconds seconds" -Verbose
+$null = Wait-Job -Job $job -Timeout $waitTimeoutSeconds
+
+if ($job.State -ne 'Running') {
+    Write-Verbose "Job finished.  State: $($job.State)" -Verbose
+    $result = Receive-Job -Job $job -Verbose
+    Write-Verbose "Result: $result" -Verbose
+    if ($result) {
+        Write-Verbose "Registration succeeded" -Verbose
+        exit 0
+    }
+    else {
+        Write-Verbose "Registration failed" -Verbose
+        # at the time this was written, the MSI is ignoring the exit code
+        exit 1
+    }
+}
+else {
+    Write-Verbose "Job timed out" -Verbose
+    Write-Verbose "Stopping Job.  State: $($job.State)" -Verbose
+    Stop-Job -Job $job
+    # at the time this was written, the MSI is ignoring the exit code
+    exit 258
+}
diff --git a/assets/wix/Product.wxs b/assets/wix/Product.wxs
index bf6ab5d7c..85cc0011d 100644
--- a/assets/wix/Product.wxs
+++ b/assets/wix/Product.wxs
@@ -87,7 +87,7 @@
     <SetProperty Id="EnableMU"
                  Before="EnableMU"
                  Sequence="execute"
-                 Value="&quot;[VersionFolder]pwsh.exe&quot; -NoProfile -ExecutionPolicy Bypass -Command &quot;$null=(New-Object -ComObject Microsoft.Update.ServiceManager).AddService2('7971f918-a847-4430-9279-4a52d1efe18d', 7, '')&quot;" />
+                 Value="&quot;[VersionFolder]pwsh.exe&quot; -NoProfile -ExecutionPolicy Bypass -File &quot;[VersionFolder]RegisterMicrosoftUpdate.ps1&quot;" />
     <CustomAction Id="EnableMU"
                   BinaryKey="WixCA"
                   DllEntry="CAQuietExec"
diff --git a/assets/wix/files.wxs b/assets/wix/files.wxs
index 6f7d0e6cb..4da5773c5 100644
--- a/assets/wix/files.wxs
+++ b/assets/wix/files.wxs
@@ -3150,6 +3150,9 @@
       <Component Id="cmp0721DF55B34849B78C98AE00D863CBFF">
         <File Id="fil72E7B79A9B10498CBE884076BC863617" KeyPath="yes" Source="$(var.ProductSourcePath)\clrgc.dll" />
       </Component>
+      <Component Id="cmp42A739F2AAC24D729DDBCAA096B8C55A">
+        <File Id="fil48314E2AC49B4E208C5667CEC50CA78C" KeyPath="yes" Source="$(var.ProductSourcePath)\RegisterMicrosoftUpdate.ps1" />
+      </Component>
     </DirectoryRef>
   </Fragment>
   <Fragment>
@@ -4162,6 +4165,7 @@
       <ComponentRef Id="cmpD182C2F62588471F82252B0038C75327" />
       <ComponentRef Id="cmpF3EB02BB62A34F509B22CC370C430A20" />
       <ComponentRef Id="cmp16FC1857F01A4E388616BA7AA8CC5F73" />
+      <ComponentRef Id="cmp42A739F2AAC24D729DDBCAA096B8C55A" />
     </ComponentGroup>
   </Fragment>
 </Wix>
diff --git a/src/powershell-win-core/powershell-win-core.csproj b/src/powershell-win-core/powershell-win-core.csproj
index 16f625827..9879e8485 100644
--- a/src/powershell-win-core/powershell-win-core.csproj
+++ b/src/powershell-win-core/powershell-win-core.csproj
@@ -29,7 +29,7 @@
       <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
       <CopyToPublishDirectory>PreserveNewest</CopyToPublishDirectory>
     </Content>
-    <Content Include="..\..\LICENSE.txt;..\..\ThirdPartyNotices.txt;..\powershell-native\Install-PowerShellRemoting.ps1;..\PowerShell.Core.Instrumentation\PowerShell.Core.Instrumentation.man;..\PowerShell.Core.Instrumentation\RegisterManifest.ps1;..\..\assets\GroupPolicy\PowerShellCoreExecutionPolicy.admx;..\..\assets\GroupPolicy\PowerShellCoreExecutionPolicy.adml;..\..\assets\GroupPolicy\InstallPSCorePolicyDefinitions.ps1">
+    <Content Include="..\..\LICENSE.txt;..\..\ThirdPartyNotices.txt;..\powershell-native\Install-PowerShellRemoting.ps1;..\PowerShell.Core.Instrumentation\PowerShell.Core.Instrumentation.man;..\PowerShell.Core.Instrumentation\RegisterManifest.ps1;..\..\assets\MicrosoftUpdate\RegisterMicrosoftUpdate.ps1;..\..\assets\GroupPolicy\PowerShellCoreExecutionPolicy.admx;..\..\assets\GroupPolicy\PowerShellCoreExecutionPolicy.adml;..\..\assets\GroupPolicy\InstallPSCorePolicyDefinitions.ps1">
       <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
       <CopyToPublishDirectory>PreserveNewest</CopyToPublishDirectory>
     </Content>