From a861f7c282ed95e5a79a055ad6f4f99eddccaf02 Mon Sep 17 00:00:00 2001
From: Rafal Harabien <rafalh@reactos.org>
Date: Wed, 20 Apr 2011 21:31:41 +0000
Subject: [PATCH] [KERNEL32] * Properly check if buffer given to
 GetComputerName is too small. Fixes hostname.exe if computer name is
 MAX_COMPUTERNAME_LENGTH long (the default for bootcd since it's generated
 randomly in Setup). * Simplify it a bit

svn path=/trunk/; revision=51412
---
 .../dll/win32/kernel32/misc/computername.c    | 39 +++++++++++--------
 1 file changed, 22 insertions(+), 17 deletions(-)

diff --git a/reactos/dll/win32/kernel32/misc/computername.c b/reactos/dll/win32/kernel32/misc/computername.c
index fcaf9f38a59..d04172d6a6a 100644
--- a/reactos/dll/win32/kernel32/misc/computername.c
+++ b/reactos/dll/win32/kernel32/misc/computername.c
@@ -39,7 +39,7 @@ BOOL
 GetComputerNameFromRegistry(LPWSTR RegistryKey,
                             LPWSTR ValueNameStr,
                             LPWSTR lpBuffer,
-                            LPDWORD nSize )
+                            LPDWORD nSize)
 {
     PKEY_VALUE_PARTIAL_INFORMATION KeyInfo;
     OBJECT_ATTRIBUTES ObjectAttributes;
@@ -50,7 +50,7 @@ GetComputerNameFromRegistry(LPWSTR RegistryKey,
     ULONG ReturnSize;
     NTSTATUS Status;
 
-    RtlInitUnicodeString(&KeyName,RegistryKey);
+    RtlInitUnicodeString(&KeyName, RegistryKey);
     InitializeObjectAttributes(&ObjectAttributes,
                                &KeyName,
                                OBJ_CASE_INSENSITIVE,
@@ -75,7 +75,7 @@ GetComputerNameFromRegistry(LPWSTR RegistryKey,
         return FALSE;
     }
 
-    RtlInitUnicodeString(&ValueName,ValueNameStr);
+    RtlInitUnicodeString(&ValueName, ValueNameStr);
 
     Status = ZwQueryValueKey(KeyHandle,
                              &ValueName,
@@ -83,35 +83,40 @@ GetComputerNameFromRegistry(LPWSTR RegistryKey,
                              KeyInfo,
                              KeyInfoSize,
                              &ReturnSize);
+
+    ZwClose(KeyHandle);
+
     if (!NT_SUCCESS(Status))
     {
-        RtlFreeHeap(RtlGetProcessHeap(), 0, KeyInfo);
-        ZwClose(KeyHandle);
         *nSize = ReturnSize;
-        SetLastErrorByStatus(Status);
-        return FALSE;
+        goto failed;
     }
 
-    if (lpBuffer && *nSize > (KeyInfo->DataLength / sizeof(WCHAR)))
+    if (KeyInfo->Type != REG_SZ)
     {
-        *nSize = KeyInfo->DataLength / sizeof(WCHAR) - 1;
-        lpBuffer[*nSize] = 0;
+        Status = STATUS_UNSUCCESSFUL;
+        goto failed;
     }
-    else
+
+    if (!lpBuffer || *nSize < (KeyInfo->DataLength / sizeof(WCHAR)))
     {
-        RtlFreeHeap(RtlGetProcessHeap(), 0, KeyInfo);
-        ZwClose(KeyHandle);
         *nSize = ReturnSize;
-        SetLastErrorByStatus(STATUS_BUFFER_OVERFLOW);
-        return FALSE;
+        Status = STATUS_BUFFER_OVERFLOW;
+        goto failed;
     }
 
-    RtlCopyMemory(lpBuffer, KeyInfo->Data, *nSize * sizeof(WCHAR));
+    *nSize = KeyInfo->DataLength / sizeof(WCHAR) - 1;
+    RtlCopyMemory(lpBuffer, KeyInfo->Data, KeyInfo->DataLength);
+    lpBuffer[*nSize] = 0;
 
     RtlFreeHeap(RtlGetProcessHeap(), 0, KeyInfo);
-    ZwClose(KeyHandle);
 
     return TRUE;
+
+failed:
+    RtlFreeHeap(RtlGetProcessHeap(), 0, KeyInfo);
+    SetLastErrorByStatus(Status);
+    return FALSE;
 }
 
 /*
-- 
2.17.1