PowerShell · daxian-dbw · Apr 28, 2023 · Aug 5, 2022 · Oct 7, 2022 · Nov 30, 2022
diff --git a/src/System.Management.Automation/engine/AsyncByteStreamTransfer.cs b/src/System.Management.Automation/engine/AsyncByteStreamTransfer.cs
@@ -0,0 +1,83 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#nullable enable
+
+using System.Buffers;
+using System.IO;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace System.Management.Automation;
+
+/// <summary>
+/// Represents the transfer of bytes from one <see cref="Stream" /> to another
+/// asynchronously.
+/// </summary>
+internal sealed class AsyncByteStreamTransfer : IDisposable
+{
+    private const int DefaultBufferSize = 1024;
+
+    private readonly BytePipe _bytePipe;
+
+    private readonly BytePipe _destinationPipe;
+
+    private readonly Memory<byte> _buffer;
+
+    private readonly CancellationTokenSource _cts = new();
+
+    private Task? _readToBufferTask;
+
+    public AsyncByteStreamTransfer(
+        BytePipe bytePipe,
+        BytePipe destinationPipe)
+    {
+        _bytePipe = bytePipe;
+        _destinationPipe = destinationPipe;
+        _buffer = new byte[DefaultBufferSize];
+    }
+
+    public Task EOF => _readToBufferTask ?? Task.CompletedTask;
+
+    public void BeginReadChunks()
+    {
+        _readToBufferTask = Task.Run(ReadBufferAsync);
+    }
+
+    public void Dispose() => _cts.Cancel();
+
+    private async Task ReadBufferAsync()
+    {
+        Stream stream;
+        Stream? destinationStream = null;
+        try
+        {
+            stream = await _bytePipe.GetStream(_cts.Token);
+            destinationStream = await _destinationPipe.GetStream(_cts.Token);
+
+            while (true)
+            {
+                int bytesRead;
+                bytesRead = await stream.ReadAsync(_buffer, _cts.Token);
+                if (bytesRead is 0)
+                {
+                    break;
+                }
+
+                destinationStream.Write(_buffer.Span.Slice(0, bytesRead));
+            }
+        }
+        catch (IOException)
+        {
+            return;
+        }
+        catch (OperationCanceledException)
+        {
+            return;
+        }
+        finally
+        {
+            destinationStream?.Close();
+        }
+    }
+}
diff --git a/src/System.Management.Automation/engine/BytePipe.cs b/src/System.Management.Automation/engine/BytePipe.cs
@@ -0,0 +1,115 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#nullable enable
+
+using System.Diagnostics;
+using System.IO;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.PowerShell.Telemetry;
+
+namespace System.Management.Automation;
+
+/// <summary>
+/// Represents a lazily retrieved <see cref="Stream" /> for transfering bytes
+/// to or from.
+/// </summary>
+internal abstract class BytePipe
+{
+    public abstract Task<Stream> GetStream(CancellationToken cancellationToken);
+
+    internal AsyncByteStreamTransfer Bind(BytePipe bytePipe)
+    {
+        Debug.Assert(bytePipe is not null);
+        return new AsyncByteStreamTransfer(bytePipe, destinationPipe: this);
+    }
+}
+
+/// <summary>
+/// Represents a <see cref="Stream" /> lazily retrieved from the underlying
+/// <see cref="NativeCommandProcessor" />.
+/// </summary>
+internal sealed class NativeCommandProcessorBytePipe : BytePipe
+{
+    private readonly NativeCommandProcessor _nativeCommand;
+
+    private readonly bool _stdout;
+
+    internal NativeCommandProcessorBytePipe(
+        NativeCommandProcessor nativeCommand,
+        bool stdout)
+    {
+        Debug.Assert(nativeCommand is not null);
+        _nativeCommand = nativeCommand;
+        _stdout = stdout;
+    }
+
+    public override async Task<Stream> GetStream(CancellationToken cancellationToken)
+    {
+        // If the native command we're wrapping is the upstream command then
+        // NativeCommandProcessor.Prepare will have already been called before
+        // the creation of this BytePipe.
+        if (_stdout)
+        {
+            return _nativeCommand.GetStream(stdout: true);
+        }
+
+        await _nativeCommand.WaitForProcessInitializationAsync(cancellationToken);
+        return _nativeCommand.GetStream(stdout: false);
+    }
+}
+
+/// <summary>
+/// Provides an byte pipe implementation representing a <see cref="FileStream" />.
+/// </summary>
+internal sealed class FileBytePipe : BytePipe
+{
+    private readonly Stream _stream;
+
+    private FileBytePipe(Stream stream)
+    {
+        Debug.Assert(stream is not null);
+        _stream = stream;
+    }
+
+    internal static FileBytePipe Create(string fileName, bool append)
+    {
+        FileStream fileStream;
+        try
+        {
+            PathUtils.MasterStreamOpen(
+                fileName,
+                resolvedEncoding: null,
+                defaultEncoding: false,
+                append,
+                Force: true,
+                NoClobber: false,
+                out fileStream,
+                streamWriter: out _,
+                readOnlyFileInfo: out _,
+                isLiteralPath: true);
+        }
+        catch (Exception e) when (e.Data.Contains(typeof(ErrorRecord)))
+        {
+            // The error record is attached to the exception when thrown to preserve
+            // the call stack.
+            ErrorRecord? errorRecord = e.Data[typeof(ErrorRecord)] as ErrorRecord;
+            if (errorRecord is null)
+            {
+                throw;
+            }
+
+            e.Data.Remove(typeof(ErrorRecord));
+            throw new RuntimeException(null, e, errorRecord);
+        }
+
+        ApplicationInsightsTelemetry.SendExperimentalUseData(
+            ExperimentalFeature.PSNativeCommandPreserveBytePipe,
+            "f");
+
+        return new FileBytePipe(fileStream);
+    }
+
+    public override Task<Stream> GetStream(CancellationToken cancellationToken) => Task.FromResult(_stream);
+}
diff --git a/src/System.Management.Automation/engine/ExperimentalFeature/ExperimentalFeature.cs b/src/System.Management.Automation/engine/ExperimentalFeature/ExperimentalFeature.cs
@@ -22,6 +22,7 @@ public class ExperimentalFeature
 
         internal const string EngineSource = "PSEngine";
         internal const string PSNativeCommandErrorActionPreferenceFeatureName = "PSNativeCommandErrorActionPreference";
+        internal const string PSNativeCommandPreserveBytePipe = "PSNativeCommandPreserveBytePipe";
         internal const string PSModuleAutoLoadSkipOfflineFilesFeatureName = "PSModuleAutoLoadSkipOfflineFiles";
         internal const string PSCustomTableHeaderLabelDecoration = "PSCustomTableHeaderLabelDecoration";
         internal const string PSFeedbackProvider = "PSFeedbackProvider";
@@ -126,6 +127,9 @@ static ExperimentalFeature()
                 new ExperimentalFeature(
                     name: PSCustomTableHeaderLabelDecoration,
                     description: "Formatting differentiation for table header labels that aren't property members"),
+                new ExperimentalFeature(
+                    name: PSNativeCommandPreserveBytePipe,
+                    description: "Byte output is retained when piping between two or more native commands"),
                 new ExperimentalFeature(
                     name: PSFeedbackProvider,
                     description: "Replace the hard-coded suggestion framework with the extensible feedback provider"),