src/Servers/Kestrel/shared/test/PassThroughConnectionMiddleware.cs - main - githubmirror/aspnetcore

src/Servers/Kestrel/shared/test/PassThroughConnectionMiddleware.cs
92 строки3 KB
William Godbe
Rename Microsoft.Aspnetcore.Testing to Microsoft.AspNetCore.InternalTesting (#51713)
02 ноя 2023, 03:15
02 ноя 2023, 03:15fc8deca
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System;
using System.IO.Pipelines;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.AspNetCore.Connections;

namespace Microsoft.AspNetCore.InternalTesting;

public class PassThroughConnectionMiddleware
{
    private readonly ConnectionDelegate _next;

    public PassThroughConnectionMiddleware(ConnectionDelegate next)
    {
        _next = next;
    }

    public Task OnConnectionAsync(ConnectionContext context)
    {
        context.Transport = new PassThroughDuplexPipe(context.Transport);
        return _next(context);
    }

    private class PassThroughDuplexPipe : IDuplexPipe
    {
        public PassThroughDuplexPipe(IDuplexPipe duplexPipe)
        {
            Input = new PassThroughPipeReader(duplexPipe.Input);
            Output = new PassThroughPipeWriter(duplexPipe.Output);
        }

        public PipeReader Input { get; }

        public PipeWriter Output { get; }

        private class PassThroughPipeWriter : PipeWriter
        {
            private readonly PipeWriter _output;

            public PassThroughPipeWriter(PipeWriter output)
            {
                _output = output;
            }

            public override void Advance(int bytes) => _output.Advance(bytes);

            public override void CancelPendingFlush() => _output.CancelPendingFlush();

            public override void Complete(Exception exception = null) => _output.Complete(exception);

            public override ValueTask<FlushResult> FlushAsync(CancellationToken cancellationToken = default) => _output.FlushAsync(cancellationToken);

            public override Memory<byte> GetMemory(int sizeHint = 0) => _output.GetMemory(sizeHint);

            public override Span<byte> GetSpan(int sizeHint = 0) => _output.GetSpan(sizeHint);
        }

        private class PassThroughPipeReader : PipeReader
        {
            private readonly PipeReader _input;

            public PassThroughPipeReader(PipeReader input)
            {
                _input = input;
            }

            public override void AdvanceTo(SequencePosition consumed) => _input.AdvanceTo(consumed);

            public override void AdvanceTo(SequencePosition consumed, SequencePosition examined) => _input.AdvanceTo(consumed, examined);

            public override void CancelPendingRead() => _input.CancelPendingRead();

            public override void Complete(Exception exception = null) => _input.Complete(exception);

            public override ValueTask<ReadResult> ReadAsync(CancellationToken cancellationToken = default) => _input.ReadAsync(cancellationToken);

            public override bool TryRead(out ReadResult result) => _input.TryRead(out result);
        }
    }
}

public static class PassThroughConnectionMiddlewareExtensions
{
    public static TBuilder UsePassThrough<TBuilder>(this TBuilder builder) where TBuilder : IConnectionBuilder
    {
        builder.Use(next => new PassThroughConnectionMiddleware(next).OnConnectionAsync);
        return builder;
    }
}