handle memory allocation errors in GetEnvironmentVariable() and ensure the string is always null-terminated

svn path=/trunk/; revision=16553
This commit is contained in:
Thomas Bluemel 2005-07-13 15:09:52 +00:00
parent 90be5e4d9b
commit a56f2840b3

View file

@ -48,11 +48,27 @@ GetEnvironmentVariableA (
/* initialize unicode variable value string and allocate buffer */ /* initialize unicode variable value string and allocate buffer */
VarValueU.Length = 0; VarValueU.Length = 0;
VarValueU.MaximumLength = nSize * sizeof(WCHAR); if (nSize != 0)
{
VarValueU.MaximumLength = (nSize - 1) * sizeof(WCHAR);
VarValueU.Buffer = RtlAllocateHeap (RtlGetProcessHeap (), VarValueU.Buffer = RtlAllocateHeap (RtlGetProcessHeap (),
0, 0,
VarValueU.MaximumLength); 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 */ /* get unicode environment variable */
Status = RtlQueryEnvironmentVariable_U (NULL, Status = RtlQueryEnvironmentVariable_U (NULL,
&VarNameU, &VarNameU,
@ -70,7 +86,7 @@ GetEnvironmentVariableA (
SetLastErrorByStatus (Status); SetLastErrorByStatus (Status);
if (Status == STATUS_BUFFER_TOO_SMALL) if (Status == STATUS_BUFFER_TOO_SMALL)
{ {
return VarValueU.Length / sizeof(WCHAR) + 1; return (VarValueU.Length / sizeof(WCHAR)) + 1;
} }
else else
{ {
@ -83,15 +99,24 @@ GetEnvironmentVariableA (
&VarValueU, &VarValueU,
FALSE); FALSE);
if (VarValueU.Buffer != NULL)
{
/* free unicode buffer */ /* free unicode buffer */
RtlFreeHeap (RtlGetProcessHeap (), RtlFreeHeap (RtlGetProcessHeap (),
0, 0,
VarValueU.Buffer); VarValueU.Buffer);
}
/* free unicode variable name string */ /* free unicode variable name string */
RtlFreeUnicodeString (&VarNameU); RtlFreeUnicodeString (&VarNameU);
return (VarValueU.Length / sizeof(WCHAR)); return (VarValueU.Length / sizeof(WCHAR));
}
else
{
SetLastError (ERROR_NOT_ENOUGH_MEMORY);
return 0;
}
} }
@ -114,7 +139,7 @@ GetEnvironmentVariableW (
lpName); lpName);
VarValue.Length = 0; VarValue.Length = 0;
VarValue.MaximumLength = nSize * sizeof(WCHAR); VarValue.MaximumLength = (nSize != 0 ? (nSize - 1) * sizeof(WCHAR) : 0);
VarValue.Buffer = lpBuffer; VarValue.Buffer = lpBuffer;
Status = RtlQueryEnvironmentVariable_U (NULL, Status = RtlQueryEnvironmentVariable_U (NULL,
@ -133,6 +158,10 @@ GetEnvironmentVariableW (
} }
} }
/* make sure the string is NULL-terminated! RtlQueryEnvironmentVariable_U
only terminates it if MaximumLength < Length */
VarValue.Buffer[VarValue.Length / sizeof(WCHAR)] = L'\0';
return (VarValue.Length / sizeof(WCHAR)); return (VarValue.Length / sizeof(WCHAR));
} }