[DISKPART] Support quoted options in commands

base/system/diskpart/diskpart.h

@@ -394,6 +394,14 @@ RoundingDivide(
     _In_ ULONGLONG Dividend,
     _In_ ULONGLONG Divisor);
 
+PWSTR
+DuplicateQuotedString(
+    _In_ PWSTR pszInString);
+
+PWSTR
+DuplicateString(
+    _In_ PWSTR pszInString);
+
 /* offline.c */
 BOOL offline_main(INT argc, LPWSTR *argv);
 

base/system/diskpart/interpreter.c

@@ -229,6 +229,7 @@ InterpretMain(VOID)
     LPWSTR args_vector[MAX_ARGS_COUNT];
     INT args_count = 0;
     BOOL bWhiteSpace = TRUE;
+    BOOL bQuote = FALSE;
     BOOL bRun = TRUE;
     LPWSTR ptr;
 
@@ -243,17 +244,21 @@ InterpretMain(VOID)
         /* Get input from the user. */
         fgetws(input_line, MAX_STRING_SIZE, stdin);
 
+        bQuote = FALSE;
         ptr = input_line;
         while (*ptr != 0)
         {
-            if (iswspace(*ptr) || *ptr == L'\n')
+            if (*ptr == L'"')
+                bQuote = !bQuote;
+
+            if ((iswspace(*ptr) && (bQuote == FALSE))|| *ptr == L'\n')
             {
                 *ptr = 0;
                 bWhiteSpace = TRUE;
             }
             else
             {
-                if ((bWhiteSpace != FALSE) && (args_count < MAX_ARGS_COUNT))
+                if ((bWhiteSpace != FALSE) && (bQuote == FALSE) && (args_count < MAX_ARGS_COUNT))
                 {
                     args_vector[args_count] = ptr;
                     args_count++;

base/system/diskpart/misc.c

@@ -78,3 +78,71 @@ RoundingDivide(
 {
     return (Dividend + Divisor / 2) / Divisor;
 }
+
+
+PWSTR
+DuplicateQuotedString(
+    _In_ PWSTR pszInString)
+{
+    PWSTR pszOutString = NULL;
+    PWSTR pStart, pEnd;
+    INT nLength;
+
+    if ((pszInString == NULL) || (pszInString[0] == UNICODE_NULL))
+        return NULL;
+
+    if (pszInString[0] == L'"')
+    {
+        if (pszInString[1] == UNICODE_NULL)
+            return NULL;
+
+        pStart = &pszInString[1];
+        pEnd = wcschr(pStart, '"');
+        if (pEnd == NULL)
+        {
+            nLength = wcslen(pStart);
+        }
+        else
+        {
+            nLength = (pEnd - pStart);
+        }
+    }
+    else
+    {
+        pStart = pszInString;
+        nLength = wcslen(pStart);
+    }
+
+    pszOutString = RtlAllocateHeap(RtlGetProcessHeap(),
+                                   HEAP_ZERO_MEMORY,
+                                   (nLength + 1) * sizeof(WCHAR));
+    if (pszOutString == NULL)
+        return NULL;
+
+    wcsncpy(pszOutString, pStart, nLength);
+
+    return pszOutString;
+}
+
+
+PWSTR
+DuplicateString(
+    _In_ PWSTR pszInString)
+{
+    PWSTR pszOutString = NULL;
+    INT nLength;
+
+    if ((pszInString == NULL) || (pszInString[0] == UNICODE_NULL))
+        return NULL;
+
+    nLength = wcslen(pszInString);
+    pszOutString = RtlAllocateHeap(RtlGetProcessHeap(),
+                                   HEAP_ZERO_MEMORY,
+                                   (nLength + 1) * sizeof(WCHAR));
+    if (pszOutString == NULL)
+        return NULL;
+
+    wcscpy(pszOutString, pszInString);
+
+    return pszOutString;
+}