diff --git a/src/Microsoft.AspNet.Server.Kestrel/Http/FrameRequestStream.cs b/src/Microsoft.AspNet.Server.Kestrel/Http/FrameRequestStream.cs index 652f2bbd4b..2c19b40dd8 100644 --- a/src/Microsoft.AspNet.Server.Kestrel/Http/FrameRequestStream.cs +++ b/src/Microsoft.AspNet.Server.Kestrel/Http/FrameRequestStream.cs @@ -11,8 +11,7 @@ namespace Microsoft.AspNet.Server.Kestrel.Http public class FrameRequestStream : Stream { private readonly MessageBody _body; - private bool _stopped; - private bool _aborted; + private StreamState _state; public FrameRequestStream(MessageBody body) { @@ -52,14 +51,7 @@ namespace Microsoft.AspNet.Server.Kestrel.Http public override int Read(byte[] buffer, int offset, int count) { - if (_stopped) - { - throw new ObjectDisposedException(nameof(FrameRequestStream)); - } - if (_aborted) - { - throw new IOException("The request has been aborted."); - } + ValidateState(); return ReadAsync(buffer, offset, count).GetAwaiter().GetResult(); } @@ -67,14 +59,7 @@ namespace Microsoft.AspNet.Server.Kestrel.Http #if NET451 public override IAsyncResult BeginRead(byte[] buffer, int offset, int count, AsyncCallback callback, object state) { - if (_stopped) - { - throw new ObjectDisposedException(nameof(FrameRequestStream)); - } - if (_aborted) - { - throw new IOException("The request has been aborted."); - } + ValidateState(); var task = ReadAsync(buffer, offset, count, CancellationToken.None, state); if (callback != null) @@ -91,14 +76,7 @@ namespace Microsoft.AspNet.Server.Kestrel.Http private Task ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken, object state) { - if (_stopped) - { - throw new ObjectDisposedException(nameof(FrameRequestStream)); - } - if (_aborted) - { - throw new IOException("The request has been aborted."); - } + ValidateState(); var tcs = new TaskCompletionSource(state); var task = _body.ReadAsync(new ArraySegment(buffer, offset, count), cancellationToken); @@ -124,14 +102,7 @@ namespace Microsoft.AspNet.Server.Kestrel.Http public override Task ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken) { - if (_stopped) - { - throw new ObjectDisposedException(nameof(FrameRequestStream)); - } - if (_aborted) - { - throw new IOException("The request has been aborted."); - } + ValidateState(); return _body.ReadAsync(new ArraySegment(buffer, offset, count), cancellationToken); } @@ -145,14 +116,37 @@ namespace Microsoft.AspNet.Server.Kestrel.Http { // Can't use dispose (or close) as can be disposed too early by user code // As exampled in EngineTests.ZeroContentLengthNotSetAutomaticallyForCertainStatusCodes - _stopped = true; + _state = StreamState.Disposed; } public void Abort() { // We don't want to throw an ODE until the app func actually completes. // If the request is aborted, we throw an IOException instead. - _aborted = true; + if (_state != StreamState.Disposed) + { + _state = StreamState.Aborted; + } + } + + private void ValidateState() + { + switch (_state) + { + case StreamState.Open: + return; + case StreamState.Disposed: + throw new ObjectDisposedException(nameof(FrameRequestStream)); + case StreamState.Aborted: + throw new IOException("The request has been aborted."); + } + } + + private enum StreamState + { + Open, + Disposed, + Aborted } } } diff --git a/src/Microsoft.AspNet.Server.Kestrel/Http/FrameResponseStream.cs b/src/Microsoft.AspNet.Server.Kestrel/Http/FrameResponseStream.cs index 25b8a72cfe..827234f872 100644 --- a/src/Microsoft.AspNet.Server.Kestrel/Http/FrameResponseStream.cs +++ b/src/Microsoft.AspNet.Server.Kestrel/Http/FrameResponseStream.cs @@ -11,8 +11,7 @@ namespace Microsoft.AspNet.Server.Kestrel.Http class FrameResponseStream : Stream { private readonly FrameContext _context; - private bool _stopped; - private bool _aborted; + private StreamState _state; public FrameResponseStream(FrameContext context) { @@ -37,28 +36,14 @@ namespace Microsoft.AspNet.Server.Kestrel.Http public override void Flush() { - if (_stopped) - { - throw new ObjectDisposedException(nameof(FrameResponseStream)); - } - if (_aborted) - { - throw new IOException("The request has been aborted."); - } + ValidateState(); _context.FrameControl.Flush(); } public override Task FlushAsync(CancellationToken cancellationToken) { - if (_stopped) - { - throw new ObjectDisposedException(nameof(FrameResponseStream)); - } - if (_aborted) - { - throw new IOException("The request has been aborted."); - } + ValidateState(); return _context.FrameControl.FlushAsync(cancellationToken); } @@ -80,28 +65,14 @@ namespace Microsoft.AspNet.Server.Kestrel.Http public override void Write(byte[] buffer, int offset, int count) { - if (_stopped) - { - throw new ObjectDisposedException(nameof(FrameResponseStream)); - } - if (_aborted) - { - throw new IOException("The request has been aborted."); - } + ValidateState(); _context.FrameControl.Write(new ArraySegment(buffer, offset, count)); } public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken) { - if (_stopped) - { - throw new ObjectDisposedException(nameof(FrameResponseStream)); - } - if (_aborted) - { - throw new IOException("The request has been aborted."); - } + ValidateState(); return _context.FrameControl.WriteAsync(new ArraySegment(buffer, offset, count), cancellationToken); } @@ -110,14 +81,37 @@ namespace Microsoft.AspNet.Server.Kestrel.Http { // Can't use dispose (or close) as can be disposed too early by user code // As exampled in EngineTests.ZeroContentLengthNotSetAutomaticallyForCertainStatusCodes - _stopped = true; + _state = StreamState.Disposed; } public void Abort() { // We don't want to throw an ODE until the app func actually completes. // If the request is aborted, we throw an IOException instead. - _aborted = true; + if (_state != StreamState.Disposed) + { + _state = StreamState.Aborted; + } + } + + private void ValidateState() + { + switch (_state) + { + case StreamState.Open: + return; + case StreamState.Disposed: + throw new ObjectDisposedException(nameof(FrameResponseStream)); + case StreamState.Aborted: + throw new IOException("The request has been aborted."); + } + } + + private enum StreamState + { + Open, + Disposed, + Aborted } } }