Update WriteConsole to not use stackalloc for buffer with too large size (#18084)

This commit is contained in:
Dongbo Wang
2022-09-16 22:11:54 +05:00
committed by GitHub
parent 1a82574743
commit 1847e86b9a
@@ -10,9 +10,9 @@
// On the use of DangerousGetHandle: If the handle has been invalidated, then the API we pass it to will return an error. These
// handles should not be exposed to recycling attacks (because they are not exposed at all), but if they were, the worse they
// could do is diddle with the console buffer.
#pragma warning disable 1634, 1691
using System;
using System.Buffers;
using System.Text;
using System.Runtime.InteropServices;
using System.Management.Automation;
@@ -1479,9 +1479,7 @@ namespace Microsoft.PowerShell
bSize.X++;
SMALL_RECT wRegion = writeRegion;
wRegion.Right++;
// Suppress the PreFAST warning about not using Marshal.GetLastWin32Error() to
// get the error code.
#pragma warning disable 56523
result = NativeMethods.WriteConsoleOutput(
consoleHandle.DangerousGetHandle(),
characterBuffer,
@@ -1491,9 +1489,6 @@ namespace Microsoft.PowerShell
}
else
{
// Suppress the PreFAST warning about not using Marshal.GetLastWin32Error() to
// get the error code.
#pragma warning disable 56523
result = NativeMethods.WriteConsoleOutput(
consoleHandle.DangerousGetHandle(),
characterBuffer,
@@ -1796,9 +1791,7 @@ namespace Microsoft.PowerShell
readRegion.Top = (short)origin.Y;
readRegion.Right = (short)(origin.X + bufferSize.X - 1);
readRegion.Bottom = (short)(origin.Y + bufferSize.Y - 1);
// Suppress the PreFAST warning about not using Marshal.GetLastWin32Error() to
// get the error code.
#pragma warning disable 56523
bool result = NativeMethods.ReadConsoleOutput(
consoleHandle.DangerousGetHandle(),
characterBuffer,
@@ -2474,9 +2467,6 @@ namespace Microsoft.PowerShell
DWORD result;
StringBuilder consoleTitle = new StringBuilder((int)bufferSize);
// Suppress the PreFAST warning about not using Marshal.GetLastWin32Error() to
// get the error code.
#pragma warning disable 56523
result = NativeMethods.GetConsoleTitle(consoleTitle, bufferSize);
// If the result is zero, it may mean and error but it may also mean
// that the window title has been set to null. Since we can't tell the
@@ -2561,12 +2551,16 @@ namespace Microsoft.PowerShell
return;
}
// Native WriteConsole doesn't support output buffer longer than 64K.
// We need to chop the output string if it is too long.
int cursor = 0; // This records the chopping position in output string
const int MaxBufferSize = 16383; // this is 64K/4 - 1 to account for possible width of each character.
// Native WriteConsole doesn't support output buffer longer than 64K, so we need to chop the output string if it is too long.
// This records the chopping position in output string.
int cursor = 0;
// This is 64K/4 - 1 to account for possible width of each character.
const int MaxBufferSize = 16383;
const int MaxStackAllocSize = 512;
ReadOnlySpan<char> outBuffer;
// In case that a new line is required, we try to write out the last chunk and the new-line string together,
// to avoid one extra call to 'WriteConsole' just for a new line string.
while (cursor + MaxBufferSize < output.Length)
{
outBuffer = output.Slice(cursor, MaxBufferSize);
@@ -2575,19 +2569,37 @@ namespace Microsoft.PowerShell
}
outBuffer = output.Slice(cursor);
if (newLine)
{
var endOfLine = Environment.NewLine.AsSpan();
var endOfLineLength = endOfLine.Length;
Span<char> outBufferLine = stackalloc char[outBuffer.Length + endOfLineLength];
outBuffer.CopyTo(outBufferLine);
endOfLine.CopyTo(outBufferLine.Slice(outBufferLine.Length - endOfLineLength));
WriteConsole(consoleHandle, outBufferLine);
}
else
if (!newLine)
{
WriteConsole(consoleHandle, outBuffer);
return;
}
char[] rentedArray = null;
string lineEnding = Environment.NewLine;
int size = outBuffer.Length + lineEnding.Length;
// We expect the 'size' will often be small, and thus optimize that case with 'stackalloc'.
Span<char> buffer = size <= MaxStackAllocSize ? stackalloc char[size] : default;
try
{
if (buffer.IsEmpty)
{
rentedArray = ArrayPool<char>.Shared.Rent(size);
buffer = rentedArray.AsSpan().Slice(0, size);
}
outBuffer.CopyTo(buffer);
lineEnding.CopyTo(buffer.Slice(outBuffer.Length));
WriteConsole(consoleHandle, buffer);
}
finally
{
if (rentedArray is not null)
{
ArrayPool<char>.Shared.Return(rentedArray);
}
}
}