forked from dotnet/aspnetcore
/
BufferingStream.cs
302 lines (245 loc) · 9.29 KB
/
BufferingStream.cs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
using System.Buffers;
using System.Diagnostics;
using System.IO.Pipelines;
using System.Runtime.CompilerServices;
using System.Text;
using Microsoft.Extensions.Logging;
namespace Microsoft.AspNetCore.HttpLogging;
internal abstract class BufferingStream : Stream, IBufferWriter<byte>
{
private const int MinimumBufferSize = 4096; // 4K
protected int _bytesBuffered;
private BufferSegment? _head;
private BufferSegment? _tail;
protected Memory<byte> _tailMemory; // remainder of tail memory
protected int _tailBytesBuffered;
protected ILogger _logger;
protected Stream _innerStream;
public BufferingStream(Stream innerStream, ILogger logger)
{
_logger = logger;
_innerStream = innerStream;
}
public override bool CanSeek => _innerStream.CanSeek;
public override bool CanRead => _innerStream.CanRead;
public override bool CanWrite => _innerStream.CanWrite;
public override long Length => _innerStream.Length;
public override long Position
{
get => _innerStream.Position;
set => _innerStream.Position = value;
}
public override int WriteTimeout
{
get => _innerStream.WriteTimeout;
set => _innerStream.WriteTimeout = value;
}
public string GetString(Encoding encoding)
{
try
{
if (_head == null || _tail == null)
{
// nothing written
return "";
}
// Only place where we are actually using the buffered data.
// update tail here.
_tail.End = _tailBytesBuffered;
var ros = new ReadOnlySequence<byte>(_head, 0, _tail, _tailBytesBuffered);
var bufferWriter = new ArrayBufferWriter<char>();
var decoder = encoding.GetDecoder();
// First calls convert on the entire ReadOnlySequence, with flush: false.
// flush: false is required as we don't want to write invalid characters that
// are spliced due to truncation. If we set flush: true, if effectively means
// we expect EOF in this array, meaning it will try to write any bytes at the end of it.
EncodingExtensions.Convert(decoder, ros, bufferWriter, flush: false, out var charUsed, out var completed);
// Afterwards, we need to call convert in a loop until complete is true.
// The first call to convert many return true, but if it doesn't, we call
// Convert with a empty ReadOnlySequence and flush: true until we get completed: true.
// This should never infinite due to the contract for decoders.
// But for safety, call this only 10 times, throwing a decode failure if it fails.
for (var i = 0; i < 10; i++)
{
if (completed)
{
return new string(bufferWriter.WrittenSpan);
}
else
{
EncodingExtensions.Convert(decoder, ReadOnlySequence<byte>.Empty, bufferWriter, flush: true, out charUsed, out completed);
}
}
throw new DecoderFallbackException("Failed to decode after 10 calls to Decoder.Convert");
}
catch (DecoderFallbackException ex)
{
_logger.DecodeFailure(ex);
return "<Decoder failure>";
}
finally
{
Reset();
}
}
public void Advance(int bytes)
{
if ((uint)bytes > (uint)_tailMemory.Length)
{
ThrowArgumentOutOfRangeException(nameof(bytes));
}
_tailBytesBuffered += bytes;
_bytesBuffered += bytes;
_tailMemory = _tailMemory.Slice(bytes);
}
public Memory<byte> GetMemory(int sizeHint = 0)
{
AllocateMemory(sizeHint);
return _tailMemory;
}
public Span<byte> GetSpan(int sizeHint = 0)
{
AllocateMemory(sizeHint);
return _tailMemory.Span;
}
private void AllocateMemory(int sizeHint)
{
if (_head is null)
{
// We need to allocate memory to write since nobody has written before
var newSegment = AllocateSegment(sizeHint);
// Set all the pointers
_head = _tail = newSegment;
_tailBytesBuffered = 0;
}
else
{
var bytesLeftInBuffer = _tailMemory.Length;
if (bytesLeftInBuffer == 0 || bytesLeftInBuffer < sizeHint)
{
Debug.Assert(_tail != null);
if (_tailBytesBuffered > 0)
{
// Flush buffered data to the segment
_tail.End += _tailBytesBuffered;
_tailBytesBuffered = 0;
}
var newSegment = AllocateSegment(sizeHint);
_tail.SetNext(newSegment);
_tail = newSegment;
}
}
}
private BufferSegment AllocateSegment(int sizeHint)
{
var newSegment = CreateSegment();
// We can't use the recommended pool so use the ArrayPool
newSegment.SetOwnedMemory(ArrayPool<byte>.Shared.Rent(GetSegmentSize(sizeHint)));
_tailMemory = newSegment.AvailableMemory;
return newSegment;
}
private static BufferSegment CreateSegment()
{
return new BufferSegment();
}
private static int GetSegmentSize(int sizeHint, int maxBufferSize = int.MaxValue)
{
// First we need to handle case where hint is smaller than minimum segment size
sizeHint = Math.Max(MinimumBufferSize, sizeHint);
// After that adjust it to fit into pools max buffer size
var adjustedToMaximumSize = Math.Min(maxBufferSize, sizeHint);
return adjustedToMaximumSize;
}
protected override void Dispose(bool disposing)
{
if (disposing)
{
Reset();
}
}
public void Reset()
{
var segment = _head;
while (segment != null)
{
var returnSegment = segment;
segment = segment.NextSegment;
// We haven't reached the tail of the linked list yet, so we can always return the returnSegment.
returnSegment.ResetMemory();
}
_head = _tail = null;
_bytesBuffered = 0;
_tailBytesBuffered = 0;
}
// Copied from https://github.com/dotnet/corefx/blob/de3902bb56f1254ec1af4bf7d092fc2c048734cc/src/System.Memory/src/System/ThrowHelper.cs
private static void ThrowArgumentOutOfRangeException(string argumentName) { throw CreateArgumentOutOfRangeException(argumentName); }
[MethodImpl(MethodImplOptions.NoInlining)]
private static Exception CreateArgumentOutOfRangeException(string argumentName) { return new ArgumentOutOfRangeException(argumentName); }
public override void Flush()
{
_innerStream.Flush();
}
public override Task FlushAsync(CancellationToken cancellationToken)
{
return _innerStream.FlushAsync(cancellationToken);
}
public override int Read(byte[] buffer, int offset, int count)
{
return _innerStream.Read(buffer, offset, count);
}
public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
{
return _innerStream.ReadAsync(buffer, offset, count, cancellationToken);
}
public override long Seek(long offset, SeekOrigin origin)
{
return _innerStream.Seek(offset, origin);
}
public override void SetLength(long value)
{
_innerStream.SetLength(value);
}
public override void Write(byte[] buffer, int offset, int count)
{
_innerStream.Write(buffer, offset, count);
}
public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
{
return _innerStream.WriteAsync(buffer, offset, count, cancellationToken);
}
public override ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
{
return _innerStream.WriteAsync(buffer, cancellationToken);
}
public override void Write(ReadOnlySpan<byte> buffer)
{
_innerStream.Write(buffer);
}
public override IAsyncResult BeginRead(byte[] buffer, int offset, int count, AsyncCallback? callback, object? state)
{
return _innerStream.BeginRead(buffer, offset, count, callback, state);
}
public override IAsyncResult BeginWrite(byte[] buffer, int offset, int count, AsyncCallback? callback, object? state)
{
return _innerStream.BeginWrite(buffer, offset, count, callback, state);
}
public override int EndRead(IAsyncResult asyncResult)
{
return _innerStream.EndRead(asyncResult);
}
public override void EndWrite(IAsyncResult asyncResult)
{
_innerStream.EndWrite(asyncResult);
}
// Do not override CopyTo/Async, they call Read/Async internally.
public override ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
{
return _innerStream.ReadAsync(buffer, cancellationToken);
}
public override ValueTask DisposeAsync()
{
return _innerStream.DisposeAsync();
}
}