Repository: dotnet/runtime
File: src/libraries/System.Net.Security/src/System/Net/Security/SslStream.Implementation.cs (C#)

Code before:

                break;
            }

            if (status.ErrorCode != SecurityStatusPalErrorCode.OK)
            {
                ArrayPool<byte>.Shared.Return(rentedBuffer);
                return new ValueTask(Task.FromException(ExceptionDispatchInfo.SetCurrentStackTrace(new IOException(SR.net_io_encrypt, SslStreamPal.GetException(status)))));
            }

            ValueTask t = writeAdapter.WriteAsync(outBuffer, 0, encryptedBytes);
            if (t.IsCompletedSuccessfully)
            {
                ArrayPool<byte>.Shared.Return(rentedBuffer);
                return t;
            }
            else
            {
                return CompleteWriteAsync(t, rentedBuffer);
            }

            async ValueTask WaitAndWriteAsync(Task t, byte[] bufferToReturn)
            {
                try
                {
                    await t.ConfigureAwait(false);
                    while (true)
                    {
                        SecurityStatusPal status = EncryptData(buffer, ref outBuffer, out int encryptedBytes);
                        if (status.ErrorCode == SecurityStatusPalErrorCode.TryAgain)
                        {
                            TaskCompletionSource<bool>? waiter = _handshakeWaiter;
                            if (waiter != null)
                            {
                                await writeAdapter.WaitAsync(waiter).ConfigureAwait(false);
                            }

                            continue;
                        }
                        else if (status.ErrorCode != SecurityStatusPalErrorCode.OK)
                        {
                            throw new IOException(SR.net_io_encrypt, SslStreamPal.GetException(status));
                        }

                        break;
                    }

                    await writeAdapter.WriteAsync(outBuffer, 0, encryptedBytes).ConfigureAwait(false);
                }
                finally
                {

Review comment (bug):
A lot of this logic appears to be duplicated with the caller. If that's true, since we expect this to be rare, could we just await a recursive call back to the caller? If there's any concern about stack diving from the recursion, we could queue the call. It'll end up being more expensive when it occurs, but it would seem to simplify things (feel free to tell me I'm wrong :) and this is supposed to be very rare.

Code after:
            }

            if (status.ErrorCode != SecurityStatusPalErrorCode.OK)
            {
                ArrayPool<byte>.Shared.Return(rentedBuffer);
                return new ValueTask(Task.FromException(ExceptionDispatchInfo.SetCurrentStackTrace(new IOException(SR.net_io_encrypt, SslStreamPal.GetException(status)))));
            }

            ValueTask t = writeAdapter.WriteAsync(outBuffer, 0, encryptedBytes);
            if (t.IsCompletedSuccessfully)
            {
                ArrayPool<byte>.Shared.Return(rentedBuffer);
                return t;
            }
            else
            {
                return CompleteWriteAsync(t, rentedBuffer);
            }

            async ValueTask WaitAndWriteAsync(TIOAdapter writeAdapter, ReadOnlyMemory<byte> buffer, Task waitTask, byte[] rentedBuffer)
            {
                byte[]? bufferToReturn = rentedBuffer;
                byte[] outBuffer = rentedBuffer;
                try
                {
                    // Wait for renegotiation to finish.
                    await waitTask.ConfigureAwait(false);

                    SecurityStatusPal status = EncryptData(buffer, ref outBuffer, out int encryptedBytes);
                    if (status.ErrorCode == SecurityStatusPalErrorCode.TryAgain)
                    {
                        // No need to hold on the buffer any more.
                        ArrayPool<byte>.Shared.Return(bufferToReturn);
                        bufferToReturn = null;
                        // Call WriteSingleChunk() recursively to avoid code duplication.
                        // This should be extremely rare in cases when second renegotiation happens concurrently with Write.
                        await WriteSingleChunk(writeAdapter, buffer).ConfigureAwait(false);
                    }
                    else if (status.ErrorCode == SecurityStatusPalErrorCode.OK)
                    {
                        await writeAdapter.WriteAsync(outBuffer, 0, encryptedBytes).ConfigureAwait(false);
                    }
                    else
                    {
                        throw new IOException(SR.net_io_encrypt, SslStreamPal.GetException(status));
                    }
                }
                finally
                {
                    if (bufferToReturn != null)
                    {