Author: ion
Date: Sun Oct 30 01:07:09 2011
New Revision: 54274

URL: http://svn.reactos.org/svn/reactos?rev=54274&view=rev
Log:
[KERNEL32]: Rewrite (in some cases, simply clean-up) environment handling 
functions for better string handling, error codes, and performance. Part of 
ongoing kernel32 work.

Modified:
    trunk/reactos/dll/win32/kernel32/client/environ.c

Modified: trunk/reactos/dll/win32/kernel32/client/environ.c
URL: 
http://svn.reactos.org/svn/reactos/trunk/reactos/dll/win32/kernel32/client/environ.c?rev=54274&r1=54273&r2=54274&view=diff
==============================================================================
--- trunk/reactos/dll/win32/kernel32/client/environ.c [iso-8859-1] (original)
+++ trunk/reactos/dll/win32/kernel32/client/environ.c [iso-8859-1] Sun Oct 30 
01:07:09 2011
@@ -24,489 +24,506 @@
  */
 DWORD
 WINAPI
-GetEnvironmentVariableA (
-       LPCSTR  lpName,
-       LPSTR   lpBuffer,
-       DWORD   nSize
-       )
-{
-       ANSI_STRING VarName;
-       ANSI_STRING VarValue;
-       UNICODE_STRING VarNameU;
-       UNICODE_STRING VarValueU;
-       NTSTATUS Status;
-
-       /* initialize unicode variable name string */
-       RtlInitAnsiString (&VarName,
-                          (LPSTR)lpName);
-       RtlAnsiStringToUnicodeString (&VarNameU,
-                                     &VarName,
-                                     TRUE);
-
-       /* initialize ansi variable value string */
-       VarValue.Length = 0;
-       VarValue.MaximumLength = (USHORT)nSize;
-       VarValue.Buffer = lpBuffer;
-
-       /* initialize unicode variable value string and allocate buffer */
-       VarValueU.Length = 0;
-       if (nSize != 0)
-       {
-           VarValueU.MaximumLength = (USHORT)(nSize - 1) * sizeof(WCHAR);
-           VarValueU.Buffer = RtlAllocateHeap (RtlGetProcessHeap (),
-                                               0,
-                                               nSize * sizeof(WCHAR));
-            if (VarValueU.Buffer != NULL)
-            {
-                /* NULL-terminate the buffer in any case! 
RtlQueryEnvironmentVariable_U
-                   only terminates it if MaximumLength < Length! */
-                VarValueU.Buffer[nSize - 1] = L'\0';
-            }
-        }
-       else
-       {
-            VarValueU.MaximumLength = 0;
-            VarValueU.Buffer = NULL;
-       }
-
-        if (VarValueU.Buffer != NULL || nSize == 0)
-        {
-            /* get unicode environment variable */
-           Status = RtlQueryEnvironmentVariable_U (NULL,
-                                                   &VarNameU,
-                                                   &VarValueU);
-           if (!NT_SUCCESS(Status))
-           {
-               /* free unicode buffer */
-               RtlFreeHeap (RtlGetProcessHeap (),
-                            0,
-                            VarValueU.Buffer);
-
-               /* free unicode variable name string */
-               RtlFreeUnicodeString (&VarNameU);
-
-               BaseSetLastNTError (Status);
-               if (Status == STATUS_BUFFER_TOO_SMALL)
-               {
-                       return (VarValueU.Length / sizeof(WCHAR)) + 1;
-               }
-               else
-               {
-                       return 0;
-               }
-           }
-
-           /* convert unicode value string to ansi */
-           RtlUnicodeStringToAnsiString (&VarValue,
-                                         &VarValueU,
-                                         FALSE);
-
-            if (VarValueU.Buffer != NULL)
-            {
-                /* free unicode buffer */
-               RtlFreeHeap (RtlGetProcessHeap (),
-                            0,
-                            VarValueU.Buffer);
-            }
-
-           /* free unicode variable name string */
-           RtlFreeUnicodeString (&VarNameU);
-
-           return (VarValueU.Length / sizeof(WCHAR));
-        }
-        else
-        {
-            SetLastError (ERROR_NOT_ENOUGH_MEMORY);
-            return 0;
-        }
-}
-
+GetEnvironmentVariableA(IN LPCSTR lpName,
+                        IN LPSTR lpBuffer,
+                        IN DWORD nSize)
+{
+    ANSI_STRING VarName, VarValue;
+    UNICODE_STRING VarNameU, VarValueU;
+    PWSTR Buffer;
+    ULONG Result = 0, UniSize = 0;
+    NTSTATUS Status;
+
+    /* Initialize all the strings */
+    RtlInitAnsiString(&VarName, lpName);
+    RtlInitUnicodeString(&VarNameU, NULL);
+    RtlInitUnicodeString(&VarValueU, NULL);
+    Status = RtlAnsiStringToUnicodeString(&VarNameU, &VarName, TRUE);
+    if (!NT_SUCCESS(Status)) goto Quickie;
+
+    /* Check if the size is too big to fit */
+    if (nSize <= UNICODE_STRING_MAX_BYTES)
+    {
+        /* Keep the given size, minus a NULL-char */
+        if (nSize) UniSize = nSize * sizeof(WCHAR) - sizeof(UNICODE_NULL);
+    }
+    else
+    {
+        /* Set the maximum possible */
+        UniSize = UNICODE_STRING_MAX_BYTES - sizeof(UNICODE_NULL);
+    }
+
+    /* Allocate the value string buffer */
+    Buffer = RtlAllocateHeap(RtlGetProcessHeap(), 0, UniSize);
+    if (!Buffer)
+    {
+        Status = STATUS_NO_MEMORY;
+        goto Quickie;
+    }
+
+    /* And initialize its string */
+    RtlInitEmptyUnicodeString(&VarValueU, Buffer, UniSize);
+
+    /* Acquire the PEB lock since we'll be querying variables now */
+    RtlAcquirePebLock();
+
+    /* Query the variable */
+    Status = RtlQueryEnvironmentVariable_U(NULL, &VarNameU, &VarValueU);
+    if ((NT_SUCCESS(Status)) && !(nSize)) Status = STATUS_BUFFER_TOO_SMALL;
+
+    /* Check if we didn't have enough space */
+    if (Status == STATUS_BUFFER_TOO_SMALL)
+    {
+        /* Fixup the length that the API returned */
+        VarValueU.MaximumLength = VarValueU.Length + 2;
+
+        /* Free old Unicode buffer */
+        RtlFreeHeap(RtlGetProcessHeap(), 0, VarValueU.Buffer);
+
+        /* Allocate new one */
+        Buffer = RtlAllocateHeap(RtlGetProcessHeap(), 0, 
VarValueU.MaximumLength);
+        if (Buffer)
+        {
+            /* Query the variable so we can know its size */
+            VarValueU.Buffer = Buffer;
+            Status = RtlQueryEnvironmentVariable_U(NULL, &VarNameU, 
&VarValueU);
+            if (NT_SUCCESS(Status))
+            {
+                /* Get the ASCII length of the variable */
+                Result = RtlUnicodeStringToAnsiSize(&VarValueU);
+            }
+        }
+        else
+        {
+            /* Set failure status */
+            Status = STATUS_NO_MEMORY;
+        }
+    }
+    else
+    {
+        /* Check if the size is too big to fit */
+        if (nSize <= MAXULONG)
+        {
+            /* Keep the given size, minus a NULL-char */
+            if (nSize) nSize = nSize - sizeof(ANSI_NULL);
+        }
+        else
+        {
+            /* Set the maximum possible */
+            nSize = MAXULONG - sizeof(ANSI_NULL);
+        }
+
+        /* Check the size */
+        Result = RtlUnicodeStringToAnsiSize(&VarValueU);
+        if (Result <= nSize)
+        {
+            /* Convert the string */
+            RtlInitEmptyAnsiString(&VarValue, lpBuffer, nSize);
+            Status = RtlUnicodeStringToAnsiString(&VarValue, &VarValueU, 
FALSE);
+            if (NT_SUCCESS(Status))
+            {
+                /* NULL-terminate and set the final length */
+                lpBuffer[VarValue.Length] = ANSI_NULL;
+                Result = VarValue.Length;
+            }
+        }
+    }
+
+    /* Release the lock */
+    RtlReleasePebLock();
+
+Quickie:
+    /* Free the strings */
+    RtlFreeUnicodeString(&VarNameU);
+    if (VarValueU.Buffer) RtlFreeHeap(RtlGetProcessHeap(), 0, 
VarValueU.Buffer);
+
+    /* Check if we suceeded */
+    if (!NT_SUCCESS(Status))
+    {
+        /* We did not, clear the result and set the error code */
+        BaseSetLastNTError(Status);
+        Result = 0;
+    }
+
+    /* Return the result */
+    return Result;
+}
 
 /*
  * @implemented
  */
 DWORD
 WINAPI
-GetEnvironmentVariableW (
-       LPCWSTR lpName,
-       LPWSTR  lpBuffer,
-       DWORD   nSize
-       )
-{
-       UNICODE_STRING VarName;
-       UNICODE_STRING VarValue;
-       NTSTATUS Status;
-
-       RtlInitUnicodeString (&VarName,
-                             lpName);
-
-       VarValue.Length = 0;
-       VarValue.MaximumLength = (USHORT) (nSize ? nSize - 1 : 0) * 
sizeof(WCHAR);
-       VarValue.Buffer = lpBuffer;
-
-       Status = RtlQueryEnvironmentVariable_U (NULL,
-                                               &VarName,
-                                               &VarValue);
-       if (!NT_SUCCESS(Status))
-       {
-               if (Status == STATUS_BUFFER_TOO_SMALL)
-               {
-                       return (VarValue.Length / sizeof(WCHAR)) + 1;
-               }
-               else
-               {
-                       BaseSetLastNTError (Status);
-                       return 0;
-               }
-       }
-       
-        if (nSize != 0)
-        {
-            /* make sure the string is NULL-terminated! 
RtlQueryEnvironmentVariable_U
-               only terminates it if MaximumLength < Length */
-           lpBuffer[VarValue.Length / sizeof(WCHAR)] = L'\0';
-       }
-
-       return (VarValue.Length / sizeof(WCHAR));
-}
-
-
-/*
- * @implemented
- */
-BOOL
-WINAPI
-SetEnvironmentVariableA (
-       LPCSTR  lpName,
-       LPCSTR  lpValue
-       )
-{
-       ANSI_STRING VarName;
-       ANSI_STRING VarValue;
-       UNICODE_STRING VarNameU;
-       UNICODE_STRING VarValueU;
-       NTSTATUS Status;
-
-       DPRINT("SetEnvironmentVariableA(Name '%s', Value '%s')\n", lpName, 
lpValue);
-
-       RtlInitAnsiString (&VarName,
-                          (LPSTR)lpName);
-       RtlAnsiStringToUnicodeString (&VarNameU,
-                                     &VarName,
-                                     TRUE);
-
-       if (lpValue)
-       {
-               RtlInitAnsiString (&VarValue,
-                                  (LPSTR)lpValue);
-               RtlAnsiStringToUnicodeString (&VarValueU,
-                                             &VarValue,
-                                             TRUE);
-
-               Status = RtlSetEnvironmentVariable (NULL,
-                                                   &VarNameU,
-                                                   &VarValueU);
-
-               RtlFreeUnicodeString (&VarValueU);
-       }
-       else
-       {
-               Status = RtlSetEnvironmentVariable (NULL,
-                                                   &VarNameU,
-                                                   NULL);
-       }
-       RtlFreeUnicodeString (&VarNameU);
-
-       if (!NT_SUCCESS(Status))
-       {
-               BaseSetLastNTError (Status);
-               return FALSE;
-       }
-
-       return TRUE;
-}
-
-
-/*
- * @implemented
- */
-BOOL
-WINAPI
-SetEnvironmentVariableW (
-       LPCWSTR lpName,
-       LPCWSTR lpValue
-       )
-{
-       UNICODE_STRING VarName;
-       UNICODE_STRING VarValue;
-       NTSTATUS Status;
-
-       DPRINT("SetEnvironmentVariableW(Name '%S', Value '%S')\n", lpName, 
lpValue);
-
-       RtlInitUnicodeString (&VarName,
-                             lpName);
-
-       RtlInitUnicodeString (&VarValue,
-                             lpValue);
-
-       Status = RtlSetEnvironmentVariable (NULL,
-                                           &VarName,
-                                           &VarValue);
-
-       if (!NT_SUCCESS(Status))
-       {
-               BaseSetLastNTError (Status);
-               return FALSE;
-       }
-
-       return TRUE;
-}
-
-
-/*
- * @implemented
- */
-LPSTR
-WINAPI
-GetEnvironmentStringsA (
-       VOID
-       )
-{
-       UNICODE_STRING UnicodeString;
-       ANSI_STRING AnsiString;
-       PWCHAR EnvU;
-       PWCHAR PtrU;
-       ULONG  Length;
-       PCHAR EnvPtr = NULL;
-
-       EnvU = (PWCHAR)(NtCurrentPeb ()->ProcessParameters->Environment);
-
-       if (EnvU == NULL)
-               return NULL;
-
-       if (*EnvU == 0)
-               return NULL;
-
-       /* get environment size */
-       PtrU = EnvU;
-       while (*PtrU)
-       {
-               while (*PtrU)
-                       PtrU++;
-               PtrU++;
-       }
-       Length = (ULONG)(PtrU - EnvU);
-       DPRINT("Length %lu\n", Length);
-
-       /* allocate environment buffer */
-       EnvPtr = RtlAllocateHeap (RtlGetProcessHeap (),
-                                 0,
-                                 Length + 1);
-        if (EnvPtr == NULL)
-        {
-            SetLastError(ERROR_NOT_ENOUGH_MEMORY);
-            return NULL;
-        }
-       DPRINT("EnvPtr %p\n", EnvPtr);
-
-       /* convert unicode environment to ansi */
-       UnicodeString.MaximumLength = (USHORT)Length * sizeof(WCHAR) + 
sizeof(WCHAR);
-       UnicodeString.Buffer = EnvU;
-
-       AnsiString.MaximumLength = (USHORT)Length + 1;
-       AnsiString.Length = 0;
-       AnsiString.Buffer = EnvPtr;
-
-       DPRINT ("UnicodeString.Buffer \'%S\'\n", UnicodeString.Buffer);
-
-       while (*(UnicodeString.Buffer))
-       {
-               UnicodeString.Length = wcslen (UnicodeString.Buffer) * 
sizeof(WCHAR);
-               UnicodeString.MaximumLength = UnicodeString.Length + 
sizeof(WCHAR);
-               if (UnicodeString.Length > 0)
-               {
-                       AnsiString.Length = 0;
-                       AnsiString.MaximumLength = (USHORT)Length + 1 - 
(AnsiString.Buffer - EnvPtr);
-
-                       RtlUnicodeStringToAnsiString (&AnsiString,
-                                                     &UnicodeString,
-                                                     FALSE);
-
-                       AnsiString.Buffer += (AnsiString.Length + 1);
-                       UnicodeString.Buffer += ((UnicodeString.Length / 
sizeof(WCHAR)) + 1);
-               }
-       }
-       *(AnsiString.Buffer) = 0;
-
-       return EnvPtr;
-}
-
-
-/*
- * @implemented
- */
-LPWSTR
-WINAPI
-GetEnvironmentStringsW (
-       VOID
-       )
-{
-       return (LPWSTR)(NtCurrentPeb ()->ProcessParameters->Environment);
-}
-
-
-/*
- * @implemented
- */
-BOOL
-WINAPI
-FreeEnvironmentStringsA (
-       LPSTR   EnvironmentStrings
-       )
-{
-       if (EnvironmentStrings == NULL)
-               return FALSE;
-
-       RtlFreeHeap (RtlGetProcessHeap (),
-                    0,
-                    EnvironmentStrings);
-
-       return TRUE;
-}
-
-
-/*
- * @implemented
- */
-BOOL
-WINAPI
-FreeEnvironmentStringsW (
-       LPWSTR  EnvironmentStrings
-       )
-{
- return TRUE;
-}
-
-
-/*
- * @implemented
- */
-DWORD
-WINAPI
-ExpandEnvironmentStringsA (
-       LPCSTR  lpSrc,
-       LPSTR   lpDst,
-       DWORD   nSize
-       )
-{
-       ANSI_STRING Source;
-       ANSI_STRING Destination;
-       UNICODE_STRING SourceU;
-       UNICODE_STRING DestinationU;
-       NTSTATUS Status;
-       ULONG Length = 0;
-
-       RtlInitAnsiString (&Source,
-                          (LPSTR)lpSrc);
-       Status = RtlAnsiStringToUnicodeString (&SourceU,
-                                              &Source,
-                                              TRUE);
+GetEnvironmentVariableW(IN LPCWSTR lpName,
+                        IN LPWSTR lpBuffer,
+                        IN DWORD nSize)
+{
+    UNICODE_STRING VarName, VarValue;
+    NTSTATUS Status;
+
+    if (nSize <= UNICODE_STRING_MAX_BYTES)
+    {
+        if (nSize) nSize = nSize * sizeof(WCHAR) - sizeof(UNICODE_NULL);
+    }
+    else
+    {
+        nSize = UNICODE_STRING_MAX_BYTES - sizeof(UNICODE_NULL);
+    }
+
+    Status = RtlInitUnicodeStringEx(&VarName, lpName);
+    if (NT_SUCCESS(Status))
+    {
+        RtlInitEmptyUnicodeString(&VarValue, lpBuffer, nSize);
+
+        Status = RtlQueryEnvironmentVariable_U(NULL, &VarName, &VarValue);
         if (!NT_SUCCESS(Status))
         {
+            if (Status == STATUS_BUFFER_TOO_SMALL)
+            {
+                return (VarValue.Length / sizeof(WCHAR)) + sizeof(ANSI_NULL);
+            }
             BaseSetLastNTError (Status);
             return 0;
         }
 
-    /* make sure we don't overflow the maximum ANSI_STRING size */
-    if (nSize > 0x7fff)
-        nSize = 0x7fff;
-
-       Destination.Length = 0;
-       Destination.MaximumLength = (USHORT)nSize;
-       Destination.Buffer = lpDst;
-
-       DestinationU.Length = 0;
-       DestinationU.MaximumLength = (USHORT)nSize * sizeof(WCHAR);
-       DestinationU.Buffer = RtlAllocateHeap (RtlGetProcessHeap (),
-                                              0,
-                                              DestinationU.MaximumLength);
-        if (DestinationU.Buffer == NULL)
-        {
-            RtlFreeUnicodeString(&SourceU);
-            SetLastError(ERROR_NOT_ENOUGH_MEMORY);
-            return 0;
-        }
-
-       Status = RtlExpandEnvironmentStrings_U (NULL,
-                                               &SourceU,
-                                               &DestinationU,
-                                               &Length);
-
-       RtlFreeUnicodeString (&SourceU);
-
-       if (!NT_SUCCESS(Status))
-       {
-               BaseSetLastNTError (Status);
-               if (Status != STATUS_BUFFER_TOO_SMALL)
-               {
-                       RtlFreeHeap (RtlGetProcessHeap (),
-                                    0,
-                                    DestinationU.Buffer);
-                       return 0;
-               }
-       }
-
-       RtlUnicodeStringToAnsiString (&Destination,
-                                     &DestinationU,
-                                     FALSE);
-
-       RtlFreeHeap (RtlGetProcessHeap (),
-                    0,
-                    DestinationU.Buffer);
-
-       return (Length / sizeof(WCHAR));
-}
-
+        lpBuffer[VarValue.Length / sizeof(WCHAR)] = UNICODE_NULL;
+    }
+
+    return (VarValue.Length / sizeof(WCHAR));
+}
+
+/*
+ * @implemented
+ */
+BOOL
+WINAPI
+SetEnvironmentVariableA(IN LPCSTR lpName,
+                        IN LPCSTR lpValue)
+{
+    ANSI_STRING VarName, VarValue;
+    UNICODE_STRING VarNameU, VarValueU;
+    NTSTATUS Status;
+
+    RtlInitAnsiString(&VarName, (LPSTR)lpName);
+    Status = RtlAnsiStringToUnicodeString(&VarNameU, &VarName, TRUE);
+    if (NT_SUCCESS(Status))
+    {
+        if (lpValue)
+        {
+            RtlInitAnsiString(&VarValue, (LPSTR)lpValue);
+            Status = RtlAnsiStringToUnicodeString(&VarValueU, &VarValue, TRUE);
+            if (NT_SUCCESS(Status))
+            {
+                Status = RtlSetEnvironmentVariable(NULL, &VarNameU, 
&VarValueU);
+                RtlFreeUnicodeString(&VarValueU);
+            }
+        }
+        else
+        {
+            Status = RtlSetEnvironmentVariable(NULL, &VarNameU, NULL);
+        }
+
+        RtlFreeUnicodeString(&VarNameU);
+
+        if (NT_SUCCESS(Status)) return TRUE;
+    }
+
+    BaseSetLastNTError(Status);
+    return FALSE;
+}
+
+/*
+ * @implemented
+ */
+BOOL
+WINAPI
+SetEnvironmentVariableW(IN LPCWSTR lpName,
+                        IN LPCWSTR lpValue)
+{
+    UNICODE_STRING VarName, VarValue;
+    NTSTATUS Status;
+
+    Status = RtlInitUnicodeStringEx(&VarName, lpName);
+    if (NT_SUCCESS(Status))
+    {
+        if (lpValue)
+        {
+            Status = RtlInitUnicodeStringEx(&VarValue, lpValue);
+            if (NT_SUCCESS(Status))
+            {
+                Status = RtlSetEnvironmentVariable(NULL, &VarName, &VarValue);
+            }
+        }
+        else
+        {
+            Status = RtlSetEnvironmentVariable(NULL, &VarName, NULL);
+        }
+
+        if (NT_SUCCESS(Status)) return TRUE;
+    }
+
+    BaseSetLastNTError(Status);
+    return FALSE;
+}
+
+/*
+ * @implemented
+ */
+LPSTR
+WINAPI
+GetEnvironmentStringsA(VOID)
+{
+    ULONG Length, Size;
+    NTSTATUS Status;
+    PWCHAR Environment, p;
+    PCHAR Buffer = NULL;
+
+    RtlAcquirePebLock();
+    p = Environment = NtCurrentPeb()->ProcessParameters->Environment;
+
+    do
+    {
+        p += wcslen(Environment) + 1;
+    } while (*p);
+
+    Length = p - Environment + sizeof(UNICODE_NULL);
+
+    Status = RtlUnicodeToMultiByteSize(&Size, Environment, Length);
+    if (NT_SUCCESS(Status))
+    {
+        Buffer = RtlAllocateHeap(RtlGetProcessHeap(), 0, Size);
+        if (Buffer)
+        {
+            Status = RtlUnicodeToOemN(Buffer, Size, 0, Environment, Length);
+            if (!NT_SUCCESS(Status))
+            {
+                RtlFreeHeap(RtlGetProcessHeap(), 0, Buffer);
+                Buffer = NULL;
+
+                BaseSetLastNTError(Status);
+            }
+        }
+        else
+        {
+            BaseSetLastNTError(STATUS_NO_MEMORY);
+        }
+    }
+    else
+    {
+        BaseSetLastNTError(Status);
+    }
+
+    RtlReleasePebLock();
+    return Buffer;
+}
+
+/*
+ * @implemented
+ */
+LPWSTR
+WINAPI
+GetEnvironmentStringsW(VOID)
+{
+    PWCHAR Environment, p;
+    ULONG Length;
+
+    RtlAcquirePebLock();
+
+    p = Environment = NtCurrentPeb()->ProcessParameters->Environment;
+
+    do
+    {
+        p += wcslen(Environment) + 1;
+    } while (*p);
+
+    Length = p - Environment + sizeof(UNICODE_NULL);
+
+    p = RtlAllocateHeap(RtlGetProcessHeap(), 0, Length);
+    if (p)
+    {
+        RtlCopyMemory(p, Environment, Length);
+    }
+    else
+    {
+        BaseSetLastNTError(STATUS_NO_MEMORY);
+    }
+
+    RtlReleasePebLock();
+    return p;
+}
+
+/*
+ * @implemented
+ */
+BOOL
+WINAPI
+FreeEnvironmentStringsA(IN LPSTR EnvironmentStrings)
+{
+    return (BOOL)RtlFreeHeap(RtlGetProcessHeap(), 0, EnvironmentStrings);
+}
+
+/*
+ * @implemented
+ */
+BOOL
+WINAPI
+FreeEnvironmentStringsW(IN LPWSTR EnvironmentStrings)
+{
+    return (BOOL)RtlFreeHeap(RtlGetProcessHeap(), 0, EnvironmentStrings);
+}
 
 /*
  * @implemented
  */
 DWORD
 WINAPI
-ExpandEnvironmentStringsW (
-       LPCWSTR lpSrc,
-       LPWSTR  lpDst,
-       DWORD   nSize
-       )
-{
-       UNICODE_STRING Source;
-       UNICODE_STRING Destination;
-       NTSTATUS Status;
-       ULONG Length = 0;
-
-       RtlInitUnicodeString (&Source,
-                             (LPWSTR)lpSrc);
+ExpandEnvironmentStringsA(IN LPCSTR lpSrc,
+                          IN LPSTR lpDst,
+                          IN DWORD nSize)
+{
+    ANSI_STRING Source, Dest;
+    UNICODE_STRING SourceU, DestU;
+    PWSTR Buffer;
+    ULONG Result = 0, UniSize = 0, Length;
+    NTSTATUS Status;
+
+    /* Initialize all the strings */
+    RtlInitAnsiString(&Source, lpSrc);
+    RtlInitUnicodeString(&SourceU, NULL);
+    RtlInitUnicodeString(&DestU, NULL);
+    Status = RtlAnsiStringToUnicodeString(&SourceU, &Source, TRUE);
+    if (!NT_SUCCESS(Status)) goto Quickie;
+
+    /* Check if the size is too big to fit */
+    if (nSize <= UNICODE_STRING_MAX_BYTES)
+    {
+        /* Keep the given size, minus a NULL-char */
+        if (nSize) UniSize = nSize * sizeof(WCHAR) - sizeof(UNICODE_NULL);
+    }
+    else
+    {
+        /* Set the maximum possible */
+        UniSize = UNICODE_STRING_MAX_BYTES - sizeof(UNICODE_NULL);
+    }
+
+    /* Allocate the value string buffer */
+    Buffer = RtlAllocateHeap(RtlGetProcessHeap(), 0, UniSize);
+    if (!Buffer)
+    {
+        Status = STATUS_NO_MEMORY;
+        goto Quickie;
+    }
+
+    /* And initialize its string */
+    RtlInitEmptyUnicodeString(&DestU, Buffer, UniSize);
+
+    /* Query the variable */
+    Length = 0;
+    Status = RtlExpandEnvironmentStrings_U(NULL, &SourceU, &DestU, &Length);
+
+    /* Check if we didn't have enough space */
+    if (Status == STATUS_BUFFER_TOO_SMALL)
+    {
+        /* Fixup the length that the API returned */
+        DestU.MaximumLength = Length;
+
+        /* Free old Unicode buffer */
+        RtlFreeHeap(RtlGetProcessHeap(), 0, DestU.Buffer);
+
+        /* Allocate new one */
+        Buffer = RtlAllocateHeap(RtlGetProcessHeap(), 0, Length);
+        if (Buffer)
+        {
+            /* Query the variable so we can know its size */
+            DestU.Buffer = Buffer;
+            Status = RtlExpandEnvironmentStrings_U(NULL, &SourceU, &DestU, 
&Length);
+            if (NT_SUCCESS(Status))
+            {
+                /* Get the ASCII length of the variable */
+                Result = RtlUnicodeStringToAnsiSize(&DestU);
+            }
+        }
+        else
+        {
+            /* Set failure status */
+            Status = STATUS_NO_MEMORY;
+        }
+    }
+    else
+    {
+        /* Check if the size is too big to fit */
+        if (nSize <= MAXULONG)
+        {
+            /* Keep the given size, minus a NULL-char */
+            if (nSize) nSize = nSize - sizeof(ANSI_NULL);
+        }
+        else
+        {
+            /* Set the maximum possible */
+            nSize = MAXULONG - sizeof(ANSI_NULL);
+        }
+
+        /* Check the size */
+        Result = RtlUnicodeStringToAnsiSize(&DestU);
+        if (Result <= nSize)
+        {
+            /* Convert the string */
+            RtlInitEmptyAnsiString(&Dest, lpDst, nSize);
+            Status = RtlUnicodeStringToAnsiString(&Dest, &DestU, FALSE);
+            if (!NT_SUCCESS(Status))
+            {
+                /* Clear the destination */
+                *lpDst = ANSI_NULL;
+            }
+        }
+    }
+Quickie:
+    /* Free the strings */
+    RtlFreeUnicodeString(&SourceU);
+    if (DestU.Buffer) RtlFreeHeap(RtlGetProcessHeap(), 0, DestU.Buffer);
+
+    /* Check if we suceeded */
+    if (!NT_SUCCESS(Status))
+    {
+        /* We did not, clear the result and set the error code */
+        BaseSetLastNTError(Status);
+        Result = 0;
+    }
+
+    /* Return the result */
+    return Result;
+}
+
+/*
+ * @implemented
+ */
+DWORD
+WINAPI
+ExpandEnvironmentStringsW(IN LPCWSTR lpSrc,
+                          IN LPWSTR lpDst,
+                          IN DWORD nSize)
+{
+    UNICODE_STRING Source, Destination;
+    NTSTATUS Status;;
+
+    RtlInitUnicodeString(&Source, (LPWSTR)lpSrc);
 
     /* make sure we don't overflow the maximum UNICODE_STRING size */
-    if (nSize > 0x7fff)
-        nSize = 0x7fff;
-
-       Destination.Length = 0;
-       Destination.MaximumLength = (USHORT)nSize * sizeof(WCHAR);
-       Destination.Buffer = lpDst;
-
-       Status = RtlExpandEnvironmentStrings_U (NULL,
-                                               &Source,
-                                               &Destination,
-                                               &Length);
-       if (!NT_SUCCESS(Status))
-       {
-               BaseSetLastNTError (Status);
-               if (Status != STATUS_BUFFER_TOO_SMALL)
-                       return 0;
-       }
-
-       return (Length / sizeof(WCHAR));
+    if (nSize > UNICODE_STRING_MAX_BYTES) nSize = UNICODE_STRING_MAX_BYTES;
+
+    RtlInitEmptyUnicodeString(&Destination, lpDst, nSize * sizeof(WCHAR));
+    Status = RtlExpandEnvironmentStrings_U(NULL,
+                                           &Source,
+                                           &Destination,
+                                           &nSize);
+    if ((NT_SUCCESS(Status)) || (Status == STATUS_BUFFER_TOO_SMALL))
+    {
+        return nSize / sizeof(WCHAR);
+    }
+
+    BaseSetLastNTError (Status);
+    return 0;
+
 }
 
 /*


Reply via email to