[PATCH v2] Add AutoReboot capability

[Owen Smith <owen.smith@xxxxxxxxxx>]

Setting HKLM\System\CCS\Services\xenbus_monitor\AutoReboot to non-zero
will allow xenbus_monitor to trigger a reboot when another driver
requests a reboot. AutoReboot is set to the maximum number of reboots
to perform. Auto reboots display a message with a 60 second timeout.

This setting can be used to allow headless/unmonitored VMs to complete
the neccessary number of reboots to return to PV disks/networks. Without
this capability its possible to update the driver on a parent device,
which may prompt for a reboot. After this reboot, its likely that
emulated devices are used whilst drivers are rebound to the device
nodes. This can leave headless/unmonitored VMs in a state where emulated
devices are in use with a pending reboot. If network settings have been
changed for PV devices (e.g. static IP addressing), then the VM may not
be accessible over RDP or similar connections.

RebootCount is cleared when no pending required reboots are detected.

Signed-off-by: Owen Smith <owen.smith@xxxxxxxxxx>
---
 src/monitor/monitor.c | 138 ++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 134 insertions(+), 4 deletions(-)

diff --git a/src/monitor/monitor.c b/src/monitor/monitor.c
index 07bf8c1..7d8689f 100644
--- a/src/monitor/monitor.c
+++ b/src/monitor/monitor.c
@@ -301,7 +301,8 @@ WTSStateName(
 
 static VOID
 DoReboot(
-    VOID
+    IN PTCHAR   Message,
+    IN DWORD    Timeout
     )
 {
     Log("waiting for pending install events...");
@@ -312,8 +313,8 @@ DoReboot(
 
 #pragma prefast(suppress:28159)
     (VOID) InitiateSystemShutdownEx(NULL,
-                                    NULL,
-                                    0,
+                                    Message,
+                                    Timeout,
                                     TRUE,
                                     TRUE,
                                     SHTDN_REASON_MAJOR_OPERATINGSYSTEM |
@@ -451,6 +452,125 @@ fail1:
     return NULL;
 }
 
+static BOOL
+TryAutoReboot(
+    IN PTCHAR           DriverName
+    )
+{
+    PMONITOR_CONTEXT    Context = &MonitorContext;
+    HRESULT             Result;
+    DWORD               Type;
+    DWORD               AutoReboot;
+    DWORD               RebootCount;
+    DWORD               Length;
+    PTCHAR              DisplayName;
+    PTCHAR              Description;
+    PTCHAR              Text;
+    DWORD               TextLength;
+    HRESULT             Error;
+
+    Length = sizeof (DWORD);
+
+    Error = RegQueryValueEx(Context->ParametersKey,
+                            "AutoReboot",
+                            NULL,
+                            &Type,
+                            (LPBYTE)&AutoReboot,
+                            &Length);
+    if (Error != ERROR_SUCCESS ||
+        Type != REG_DWORD)
+        AutoReboot = 0;
+
+    if (AutoReboot == 0)
+        goto done;
+
+    Length = sizeof (DWORD);
+
+    Error = RegQueryValueEx(Context->ParametersKey,
+                            "RebootCount",
+                            NULL,
+                            &Type,
+                            (LPBYTE)&RebootCount,
+                            &Length);
+    if (Error != ERROR_SUCCESS ||
+        Type != REG_DWORD)
+        RebootCount = 0;
+
+    if (RebootCount >= AutoReboot)
+        goto done;
+
+    Log("AutoRebooting (reboot %u of %u)\n",
+        RebootCount,
+        AutoReboot);
+
+    ++RebootCount;
+
+    (VOID) RegSetValueEx(Context->ParametersKey,
+                         "RebootCount",
+                         0,
+                         REG_DWORD,
+                         (const BYTE*)&RebootCount,
+                         (DWORD) sizeof(DWORD));
+
+    (VOID) RegFlushKey(Context->ParametersKey);
+
+    Context->RebootPending = TRUE;
+
+    DisplayName = GetDisplayName(DriverName);
+    if (DisplayName == NULL)
+        goto fail1;
+
+    Description = _tcsrchr(DisplayName, ';');
+    if (Description == NULL)
+        Description = DisplayName;
+    else
+        Description++;
+
+    TextLength = (DWORD)((_tcslen(Description) +
+                          1 + // ' '
+                          _tcslen(Context->Text) +
+                          1) * sizeof (TCHAR));
+
+    Text = calloc(1, TextLength);
+    if (Text == NULL)
+        goto fail2;
+
+    Result = StringCbPrintf(Text,
+                            TextLength,
+                            TEXT("%s %s"),
+                            Description,
+                            Context->Text);
+    assert(SUCCEEDED(Result));
+
+    free(DisplayName);
+
+    DoReboot(Text, 60);
+
+    free(Text);
+
+    return TRUE;
+
+done:
+    return FALSE;
+
+fail2:
+    Log("fail2");
+
+    free(DisplayName);
+
+fail1:
+    Error = GetLastError();
+
+    {
+        PTCHAR  Message;
+        Message = GetErrorMessage(Error);
+        Log("fail1 (%s)", Message);
+        LocalFree(Message);
+    }
+
+    return FALSE;
+}
+
 static VOID
 PromptForReboot(
     IN PTCHAR           DriverName
@@ -476,6 +596,10 @@ PromptForReboot(
     TitleLength = (DWORD)((_tcslen(Context->Title) +
                            1) * sizeof (TCHAR));
 
+    // AutoReboot is set, DoReboot has been called
+    if (TryAutoReboot(DriverName))
+        goto done;
+
     DisplayName = GetDisplayName(DriverName);
     if (DisplayName == NULL)
         goto fail1;
@@ -547,7 +671,7 @@ PromptForReboot(
         Context->RebootPending = TRUE;
 
         if (Response == IDYES || Response == IDTIMEOUT)
-            DoReboot();
+            DoReboot(NULL, 0);
 
         break;
     }
@@ -556,6 +680,7 @@ PromptForReboot(
 
     free(DisplayName);
 
+done:
     Log("<====");
 
     return;
@@ -672,6 +797,11 @@ loop:
         RegCloseKey(SubKey);
     }
 
+    Error = RegDeleteValue(Context->ParametersKey,
+                           "RebootCount");
+    if (Error == ERROR_SUCCESS)
+        (VOID) RegFlushKey(Context->ParametersKey);
+
     goto done;
 
 found:
-- 
2.28.0.windows.1