[PATCH 1/2] Refactor monitor.c

[Owen Smith <owen.smith@xxxxxxxxxx>]

- Split reboot question from statement
- Move GetDisplayName to seperate function call

Signed-off-by: Owen Smith <owen.smith@xxxxxxxxxx>
---
 src/monitor/monitor.c | 155 ++++++++++++++++++++++++++++++------------
 src/xenbus.inf        |   4 +-
 2 files changed, 114 insertions(+), 45 deletions(-)

diff --git a/src/monitor/monitor.c b/src/monitor/monitor.c
index 3cf33ce..07bf8c1 100644
--- a/src/monitor/monitor.c
+++ b/src/monitor/monitor.c
@@ -59,6 +59,7 @@ typedef struct _MONITOR_CONTEXT {
     HKEY                    RequestKey;
     PTCHAR                  Title;
     PTCHAR                  Text;
+    PTCHAR                  Question;
     BOOL                    RebootPending;
 } MONITOR_CONTEXT, *PMONITOR_CONTEXT;
 
@@ -348,36 +349,20 @@ GetPromptTimeout(
     return Value;
 }
 
-static VOID
-PromptForReboot(
+static PTCHAR
+GetDisplayName(
     IN PTCHAR           DriverName
     )
 {
-    PMONITOR_CONTEXT    Context = &MonitorContext;
-    PTCHAR              Title;
-    DWORD               TitleLength;
     HRESULT             Result;
     TCHAR               ServiceKeyName[MAX_PATH];
     HKEY                ServiceKey;
     DWORD               MaxValueLength;
+    DWORD               Type;
     DWORD               DisplayNameLength;
     PTCHAR              DisplayName;
-    DWORD               Type;
-    PTCHAR              Description;
-    PTCHAR              Text;
-    DWORD               TextLength;
-    PWTS_SESSION_INFO   SessionInfo;
-    DWORD               Count;
-    DWORD               Index;
-    BOOL                Success;
     HRESULT             Error;
 
-    Log("====> (%s)", DriverName);
-
-    Title = Context->Title;
-    TitleLength = (DWORD)((_tcslen(Context->Title) +
-                           1) * sizeof (TCHAR));
-
     Result = StringCbPrintf(ServiceKeyName,
                             MAX_PATH,
                             SERVICES_KEY "\\%s",
@@ -433,6 +418,68 @@ PromptForReboot(
         goto fail5;
     }
 
+    RegCloseKey(ServiceKey);
+
+    return DisplayName;
+
+fail5:
+    Log("fail5");
+
+fail4:
+    Log("fail4");
+
+    free(DisplayName);
+
+fail3:
+    Log("fail3");
+
+fail2:
+    Log("fail2");
+
+    RegCloseKey(ServiceKey);
+
+fail1:
+    Error = GetLastError();
+
+    {
+        PTCHAR  Message;
+        Message = GetErrorMessage(Error);
+        Log("fail1 (%s)", Message);
+        LocalFree(Message);
+    }
+
+    return NULL;
+}
+
+static VOID
+PromptForReboot(
+    IN PTCHAR           DriverName
+    )
+{
+    PMONITOR_CONTEXT    Context = &MonitorContext;
+    HRESULT             Result;
+    PTCHAR              Title;
+    DWORD               TitleLength;
+    PTCHAR              DisplayName;
+    PTCHAR              Description;
+    PTCHAR              Text;
+    DWORD               TextLength;
+    PWTS_SESSION_INFO   SessionInfo;
+    DWORD               Count;
+    DWORD               Index;
+    BOOL                Success;
+    HRESULT             Error;
+
+    Log("====> (%s)", DriverName);
+
+    Title = Context->Title;
+    TitleLength = (DWORD)((_tcslen(Context->Title) +
+                           1) * sizeof (TCHAR));
+
+    DisplayName = GetDisplayName(DriverName);
+    if (DisplayName == NULL)
+        goto fail1;
+
     Description = _tcsrchr(DisplayName, ';');
     if (Description == NULL)
         Description = DisplayName;
@@ -442,17 +489,20 @@ PromptForReboot(
     TextLength = (DWORD)((_tcslen(Description) +
                           1 + // ' '
                           _tcslen(Context->Text) +
+                          1 + // ' '
+                          _tcslen(Context->Question) +
                           1) * sizeof (TCHAR));
 
     Text = calloc(1, TextLength);
     if (Text == NULL)
-        goto fail6;
+        goto fail2;
 
     Result = StringCbPrintf(Text,
                             TextLength,
-                            TEXT("%s %s"),
+                            TEXT("%s %s %s"),
                             Description,
-                            Context->Text);
+                            Context->Text,
+                            Context->Question);
     assert(SUCCEEDED(Result));
 
     Success = WTSEnumerateSessions(WTS_CURRENT_SERVER_HANDLE,
@@ -461,7 +511,7 @@ PromptForReboot(
                                    &SessionInfo,
                                    &Count);
     if (!Success)
-        goto fail7;
+        goto fail3;
 
     for (Index = 0; Index < Count; Index++) {
         DWORD                   SessionId = SessionInfo[Index].SessionId;
@@ -492,7 +542,7 @@ PromptForReboot(
                                  TRUE);
 
         if (!Success)
-            goto fail8;
+            goto fail4;
 
         Context->RebootPending = TRUE;
 
@@ -506,30 +556,14 @@ PromptForReboot(
 
     free(DisplayName);
 
-    RegCloseKey(ServiceKey);
-
     Log("<====");
 
     return;
 
-fail8:
-    Log("fail8");
-
-    WTSFreeMemory(SessionInfo);
-
-fail7:
-    Log("fail7");
-
-fail6:
-    Log("fail6");
-
-fail5:
-    Log("fail5");
-
 fail4:
     Log("fail4");
 
-    free(DisplayName);
+    WTSFreeMemory(SessionInfo);
 
 fail3:
     Log("fail3");
@@ -537,7 +571,7 @@ fail3:
 fail2:
     Log("fail2");
 
-    RegCloseKey(ServiceKey);
+    free(DisplayName);
 
 fail1:
     Error = GetLastError();
@@ -862,6 +896,7 @@ GetDialogParameters(
     DWORD               MaxValueLength;
     DWORD               TitleLength;
     DWORD               TextLength;
+    DWORD               QuestionLength;
     DWORD               Type;
     HRESULT             Error;
 
@@ -926,8 +961,41 @@ GetDialogParameters(
         goto fail7;
     }
 
+    QuestionLength = MaxValueLength + sizeof (TCHAR);
+
+    Context->Question = calloc(1, QuestionLength);
+    if (Context == NULL)
+        goto fail8;
+
+    Error = RegQueryValueEx(Context->ParametersKey,
+                            "DialogQuestion",
+                            NULL,
+                            &Type,
+                            (LPBYTE)Context->Question,
+                            &QuestionLength);
+    if (Error != ERROR_SUCCESS) {
+        SetLastError(Error);
+        goto fail9;
+    }
+
+    if (Type != REG_SZ) {
+        SetLastError(ERROR_BAD_FORMAT);
+        goto fail10;
+    }
+
     return TRUE;
 
+fail10:
+    Log("fail10");
+
+fail9:
+    Log("fail9");
+
+    free(Context->Question);
+
+fail8:
+    Log("fail8");
+
 fail7:
     Log("fail7");
 
@@ -963,8 +1031,6 @@ fail1:
     return FALSE;
 }
 
-
-
 VOID WINAPI
 MonitorMain(
     _In_    DWORD       argc,
@@ -1080,6 +1146,7 @@ MonitorMain(
 done:
     (VOID) RegDeleteTree(Context->RequestKey, NULL);
 
+    free(Context->Question);
     free(Context->Text);
     free(Context->Title);
     CloseHandle(Context->RequestKey);
diff --git a/src/xenbus.inf b/src/xenbus.inf
index bf3f0a1..f9a5475 100644
--- a/src/xenbus.inf
+++ b/src/xenbus.inf
@@ -134,6 +134,7 @@ HKR,"Parameters",,0x00000010
 HKR,"Parameters","RequestKey",0x00000000,%RequestKey%
 HKR,"Parameters","DialogTitle",0x00000000,%DialogTitle%
 HKR,"Parameters","DialogText",0x00000000,%DialogText%
+HKR,"Parameters","DialogQuestion",0x00000000,%DialogQuestion%
 
 [Monitor_EventLog]
 AddReg=Monitor_EventLog_AddReg
@@ -159,7 +160,8 @@ MonitorName="@PRODUCT_NAME@ PV Driver Monitor"
 MonitorDesc="Provides support for @PRODUCT_NAME@ PV drivers"
 RequestKey="SYSTEM\CurrentControlSet\Services\xenbus_monitor\Request"
 DialogTitle="@PRODUCT_NAME@"
-DialogText="needs to restart the system to complete installation. Press 'Yes' 
to restart the system now or 'No' if you plan to restart the system later."
+DialogText="needs to restart the system to complete installation."
+DialogQuestion="Press 'Yes' to restart the system now or 'No' if you plan to 
restart the system later."
 
 SERVICE_BOOT_START=0x0 
 SERVICE_SYSTEM_START=0x1 
-- 
2.28.0.windows.1