From f0224548c98bbea41dbc3a84f73302c8003ac980 Mon Sep 17 00:00:00 2001 From: Tony Di Nucci <35197667+tdinucci@users.noreply.github.com> Date: Sat, 19 May 2018 20:01:14 +0100 Subject: [PATCH] Override CopyTo on DeflateStream (dotnet/corefx#29751) * Override CopyTo on DeflateStream * Address PR comments Commit migrated from https://github.com/dotnet/corefx/commit/3b65271ea8e59b085b6a3e374ab15bb7e8745231 --- .../Compression/CompressionStreamUnitTestBase.cs | 15 ++++ .../ref/System.IO.Compression.cs | 1 + .../IO/Compression/DeflateZLib/DeflateStream.cs | 98 ++++++++++++++++++++-- 3 files changed, 106 insertions(+), 8 deletions(-) diff --git a/src/libraries/Common/tests/System/IO/Compression/CompressionStreamUnitTestBase.cs b/src/libraries/Common/tests/System/IO/Compression/CompressionStreamUnitTestBase.cs index 1a82bdf..dea719c 100644 --- a/src/libraries/Common/tests/System/IO/Compression/CompressionStreamUnitTestBase.cs +++ b/src/libraries/Common/tests/System/IO/Compression/CompressionStreamUnitTestBase.cs @@ -618,6 +618,21 @@ namespace System.IO.Compression } } + [Theory] + [InlineData(CompressionMode.Compress)] + [InlineData(CompressionMode.Decompress)] + public void CopyTo_ArgumentValidation(CompressionMode mode) + { + using (Stream compressor = CreateStream(new MemoryStream(), mode)) + { + AssertExtensions.Throws("destination", () => { compressor.CopyTo(null); }); + AssertExtensions.Throws("bufferSize", () => { compressor.CopyTo(new MemoryStream(), 0); }); + Assert.Throws(() => { compressor.CopyTo(new MemoryStream(new byte[1], writable: false)); }); + compressor.Dispose(); + Assert.Throws(() => { compressor.CopyTo(new MemoryStream()); }); + } + } + public enum ReadWriteMode { SyncArray, diff --git a/src/libraries/System.IO.Compression/ref/System.IO.Compression.cs b/src/libraries/System.IO.Compression/ref/System.IO.Compression.cs index 6676bf3..6731463 100644 --- a/src/libraries/System.IO.Compression/ref/System.IO.Compression.cs +++ b/src/libraries/System.IO.Compression/ref/System.IO.Compression.cs @@ -49,6 +49,7 @@ namespace System.IO.Compression public override void Write(System.ReadOnlySpan buffer) { } public override System.Threading.Tasks.Task WriteAsync(byte[] array, int offset, int count, System.Threading.CancellationToken cancellationToken) { throw null; } public override System.Threading.Tasks.ValueTask WriteAsync(System.ReadOnlyMemory buffer, System.Threading.CancellationToken cancellationToken = default) { throw null; } + public override void CopyTo(Stream destination, int bufferSize) { } } public partial class GZipStream : System.IO.Stream { diff --git a/src/libraries/System.IO.Compression/src/System/IO/Compression/DeflateZLib/DeflateStream.cs b/src/libraries/System.IO.Compression/src/System/IO/Compression/DeflateZLib/DeflateStream.cs index 853afde..1b287ee 100644 --- a/src/libraries/System.IO.Compression/src/System/IO/Compression/DeflateZLib/DeflateStream.cs +++ b/src/libraries/System.IO.Compression/src/System/IO/Compression/DeflateZLib/DeflateStream.cs @@ -706,6 +706,16 @@ namespace System.IO.Compression } } + public override void CopyTo(Stream destination, int bufferSize) + { + StreamHelpers.ValidateCopyToArgs(this, destination, bufferSize); + + EnsureDecompressionMode(); + EnsureNotDisposed(); + + new CopyToStream(this, destination, bufferSize).CopyFromSourceToDestination(); + } + public override Task CopyToAsync(Stream destination, int bufferSize, CancellationToken cancellationToken) { // Validation as base CopyToAsync would do @@ -723,17 +733,22 @@ namespace System.IO.Compression } // Do the copy - return new CopyToAsyncStream(this, destination, bufferSize, cancellationToken).CopyFromSourceToDestination(); + return new CopyToStream(this, destination, bufferSize, cancellationToken).CopyFromSourceToDestinationAsync(); } - private sealed class CopyToAsyncStream : Stream + private sealed class CopyToStream : Stream { private readonly DeflateStream _deflateStream; private readonly Stream _destination; private readonly CancellationToken _cancellationToken; private byte[] _arrayPoolBuffer; - public CopyToAsyncStream(DeflateStream deflateStream, Stream destination, int bufferSize, CancellationToken cancellationToken) + public CopyToStream(DeflateStream deflateStream, Stream destination, int bufferSize) : + this(deflateStream, destination, bufferSize, CancellationToken.None) + { + } + + public CopyToStream(DeflateStream deflateStream, Stream destination, int bufferSize, CancellationToken cancellationToken) { Debug.Assert(deflateStream != null); Debug.Assert(destination != null); @@ -745,7 +760,7 @@ namespace System.IO.Compression _arrayPoolBuffer = ArrayPool.Shared.Rent(bufferSize); } - public async Task CopyFromSourceToDestination() + public async Task CopyFromSourceToDestinationAsync() { _deflateStream.AsyncOperationStarting(); try @@ -758,7 +773,10 @@ namespace System.IO.Compression { await _destination.WriteAsync(new ReadOnlyMemory(_arrayPoolBuffer, 0, bytesRead), _cancellationToken).ConfigureAwait(false); } - else break; + else + { + break; + } } // Now, use the source stream's CopyToAsync to push directly to our inflater via this helper stream @@ -773,6 +791,34 @@ namespace System.IO.Compression } } + public void CopyFromSourceToDestination() + { + try + { + // Flush any existing data in the inflater to the destination stream. + while (true) + { + int bytesRead = _deflateStream._inflater.Inflate(_arrayPoolBuffer, 0, _arrayPoolBuffer.Length); + if (bytesRead > 0) + { + _destination.Write(_arrayPoolBuffer, 0, bytesRead); + } + else + { + break; + } + } + + // Now, use the source stream's CopyToAsync to push directly to our inflater via this helper stream + _deflateStream._stream.CopyTo(this, _arrayPoolBuffer.Length); + } + finally + { + ArrayPool.Shared.Return(_arrayPoolBuffer); + _arrayPoolBuffer = null; + } + } + public override async Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken) { // Validate inputs @@ -800,14 +846,50 @@ namespace System.IO.Compression { await _destination.WriteAsync(new ReadOnlyMemory(_arrayPoolBuffer, 0, bytesRead), cancellationToken).ConfigureAwait(false); } - else break; + else + { + break; + } + } + } + + public override void Write(byte[] buffer, int offset, int count) + { + // Validate inputs + Debug.Assert(buffer != _arrayPoolBuffer); + _deflateStream.EnsureNotDisposed(); + + if (count <= 0) + { + return; + } + else if (count > buffer.Length - offset) + { + // The buffer stream is either malicious or poorly implemented and returned a number of + // bytes larger than the buffer supplied to it. + throw new InvalidDataException(SR.GenericInvalidData); + } + + // Feed the data from base stream into the decompression engine. + _deflateStream._inflater.SetInput(buffer, offset, count); + + // While there's more decompressed data available, forward it to the buffer stream. + while (true) + { + int bytesRead = _deflateStream._inflater.Inflate(new Span(_arrayPoolBuffer)); + if (bytesRead > 0) + { + _destination.Write(_arrayPoolBuffer, 0, bytesRead); + } + else + { + break; + } } } - public override void Write(byte[] buffer, int offset, int count) => WriteAsync(buffer, offset, count, default(CancellationToken)).GetAwaiter().GetResult(); public override bool CanWrite => true; public override void Flush() { } - public override bool CanRead => false; public override bool CanSeek => false; public override long Length { get { throw new NotSupportedException(); } } -- 2.7.4