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,17 +48,33 @@ 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.Buffer = RtlAllocateHeap (RtlGetProcessHeap (),
0,
VarValueU.MaximumLength);
/* get unicode environment variable */
Status = RtlQueryEnvironmentVariable_U (NULL,
&VarNameU,
&VarValueU);
if (!NT_SUCCESS(Status))
{ {
VarValueU.MaximumLength = (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 */ /* free unicode buffer */
RtlFreeHeap (RtlGetProcessHeap (), RtlFreeHeap (RtlGetProcessHeap (),
0, 0,
@ -70,28 +86,37 @@ 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
{ {
return 0; return 0;
} }
} }
/* convert unicode value string to ansi */ /* convert unicode value string to ansi */
RtlUnicodeStringToAnsiString (&VarValue, RtlUnicodeStringToAnsiString (&VarValue,
&VarValueU, &VarValueU,
FALSE); FALSE);
/* free unicode buffer */ if (VarValueU.Buffer != NULL)
RtlFreeHeap (RtlGetProcessHeap (), {
0, /* free unicode buffer */
VarValueU.Buffer); RtlFreeHeap (RtlGetProcessHeap (),
0,
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,
@ -125,7 +150,7 @@ GetEnvironmentVariableW (
SetLastErrorByStatus (Status); SetLastErrorByStatus (Status);
if (Status == STATUS_BUFFER_TOO_SMALL) if (Status == STATUS_BUFFER_TOO_SMALL)
{ {
return (VarValue.Length / sizeof(WCHAR)) + 1; return (VarValue.Length / sizeof(WCHAR)) + 1;
} }
else else
{ {
@ -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));
} }