aspnetcore/src/Kestrel.Transport.Abstractions/Internal/TransportConnection.Features.cs

using System;
using System.Buffers;
using System.Collections;
using System.Collections.Generic;
using System.IO.Pipelines;
using System.Net;
using Microsoft.AspNetCore.Http.Features;
using Microsoft.AspNetCore.Protocols.Features;

namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal
{
    public partial class TransportConnection : IFeatureCollection,
                                               IHttpConnectionFeature,
                                               IConnectionIdFeature,
                                               IConnectionTransportFeature,
                                               IMemoryPoolFeature,
                                               IApplicationTransportFeature,
                                               ITransportSchedulerFeature
    {
        private static readonly Type IHttpConnectionFeatureType = typeof(IHttpConnectionFeature);
        private static readonly Type IConnectionIdFeatureType = typeof(IConnectionIdFeature);
        private static readonly Type IConnectionTransportFeatureType = typeof(IConnectionTransportFeature);
        private static readonly Type IMemoryPoolFeatureType = typeof(IMemoryPoolFeature);
        private static readonly Type IApplicationTransportFeatureType = typeof(IApplicationTransportFeature);
        private static readonly Type ITransportSchedulerFeatureType = typeof(ITransportSchedulerFeature);

        private object _currentIHttpConnectionFeature;
        private object _currentIConnectionIdFeature;
        private object _currentIConnectionTransportFeature;
        private object _currentIMemoryPoolFeature;
        private object _currentIApplicationTransportFeature;
        private object _currentITransportSchedulerFeature;

        private int _featureRevision;

        private List<KeyValuePair<Type, object>> MaybeExtra;

        private object ExtraFeatureGet(Type key)
        {
            if (MaybeExtra == null)
            {
                return null;
            }
            for (var i = 0; i < MaybeExtra.Count; i++)
            {
                var kv = MaybeExtra[i];
                if (kv.Key == key)
                {
                    return kv.Value;
                }
            }
            return null;
        }

        private void ExtraFeatureSet(Type key, object value)
        {
            if (MaybeExtra == null)
            {
                MaybeExtra = new List<KeyValuePair<Type, object>>(2);
            }

            for (var i = 0; i < MaybeExtra.Count; i++)
            {
                if (MaybeExtra[i].Key == key)
                {
                    MaybeExtra[i] = new KeyValuePair<Type, object>(key, value);
                    return;
                }
            }
            MaybeExtra.Add(new KeyValuePair<Type, object>(key, value));
        }

        bool IFeatureCollection.IsReadOnly => false;

        int IFeatureCollection.Revision => _featureRevision;

        string IHttpConnectionFeature.ConnectionId
        {
            get => ConnectionId;
            set => ConnectionId = value;
        }

        IPAddress IHttpConnectionFeature.RemoteIpAddress
        {
            get => RemoteAddress;
            set => RemoteAddress = value;
        }

        IPAddress IHttpConnectionFeature.LocalIpAddress
        {
            get => LocalAddress;
            set => LocalAddress = value;
        }

        int IHttpConnectionFeature.RemotePort
        {
            get => RemotePort;
            set => RemotePort = value;
        }

        int IHttpConnectionFeature.LocalPort
        {
            get => LocalPort;
            set => LocalPort = value;
        }

        MemoryPool<byte> IMemoryPoolFeature.MemoryPool => MemoryPool;

        IDuplexPipe IConnectionTransportFeature.Transport
        {
            get => Transport;
            set => Transport = value;
        }

        IDuplexPipe IApplicationTransportFeature.Application
        {
            get => Application;
            set => Application = value;
        }

        PipeScheduler ITransportSchedulerFeature.InputWriterScheduler => InputWriterScheduler;
        PipeScheduler ITransportSchedulerFeature.OutputReaderScheduler => OutputReaderScheduler;

        object IFeatureCollection.this[Type key]
        {
            get
            {
                if (key == IHttpConnectionFeatureType)
                {
                    return _currentIHttpConnectionFeature;
                }

                if (key == IConnectionIdFeatureType)
                {
                    return _currentIConnectionIdFeature;
                }

                if (key == IConnectionTransportFeatureType)
                {
                    return _currentIConnectionTransportFeature;
                }

                if (key == IMemoryPoolFeatureType)
                {
                    return _currentIMemoryPoolFeature;
                }

                if (key == IApplicationTransportFeatureType)
                {
                    return _currentIApplicationTransportFeature;
                }

                if (key == ITransportSchedulerFeatureType)
                {
                    return _currentITransportSchedulerFeature;
                }

                if (MaybeExtra != null)
                {
                    return ExtraFeatureGet(key);
                }

                return null;
            }
            set
            {
                _featureRevision++;

                if (key == IHttpConnectionFeatureType)
                {
                    _currentIHttpConnectionFeature = value;
                }
                else if (key == IConnectionIdFeatureType)
                {
                    _currentIConnectionIdFeature = value;
                }
                else if (key == IConnectionTransportFeatureType)
                {
                    _currentIConnectionTransportFeature = value;
                }
                else if (key == IMemoryPoolFeatureType)
                {
                    _currentIMemoryPoolFeature = value;
                }
                else if (key == IApplicationTransportFeatureType)
                {
                    _currentIApplicationTransportFeature = value;
                }
                else if (key == ITransportSchedulerFeatureType)
                {
                    _currentITransportSchedulerFeature = value;
                }
                else
                {
                    ExtraFeatureSet(key, value);
                }
            }
        }

        TFeature IFeatureCollection.Get<TFeature>()
        {
            if (typeof(TFeature) == IHttpConnectionFeatureType)
            {
                return (TFeature)_currentIHttpConnectionFeature;
            }
            else if (typeof(TFeature) == IConnectionIdFeatureType)
            {
                return (TFeature)_currentIConnectionIdFeature;
            }
            else if (typeof(TFeature) == IConnectionTransportFeatureType)
            {
                return (TFeature)_currentIConnectionTransportFeature;
            }
            else if (typeof(TFeature) == IMemoryPoolFeatureType)
            {
                return (TFeature)_currentIMemoryPoolFeature;
            }
            else if (typeof(TFeature) == IApplicationTransportFeatureType)
            {
                return (TFeature)_currentIApplicationTransportFeature;
            }
            else if (typeof(TFeature) == ITransportSchedulerFeatureType)
            {
                return (TFeature)_currentITransportSchedulerFeature;
            }
            else if (MaybeExtra != null)
            {
                return (TFeature)ExtraFeatureGet(typeof(TFeature));
            }

            return default;
        }

        void IFeatureCollection.Set<TFeature>(TFeature instance)
        {
            _featureRevision++;

            if (typeof(TFeature) == IHttpConnectionFeatureType)
            {
                _currentIHttpConnectionFeature = instance;
            }
            else if (typeof(TFeature) == IConnectionIdFeatureType)
            {
                _currentIConnectionIdFeature = instance;
            }
            else if (typeof(TFeature) == IConnectionTransportFeatureType)
            {
                _currentIConnectionTransportFeature = instance;
            }
            else if (typeof(TFeature) == IMemoryPoolFeatureType)
            {
                _currentIMemoryPoolFeature = instance;
            }
            else if (typeof(TFeature) == IApplicationTransportFeatureType)
            {
                _currentIApplicationTransportFeature = instance;
            }
            else if (typeof(TFeature) == ITransportSchedulerFeatureType)
            {
                _currentITransportSchedulerFeature = instance;
            }
            else
            {
                ExtraFeatureSet(typeof(TFeature), instance);
            }
        }

        IEnumerator<KeyValuePair<Type, object>> IEnumerable<KeyValuePair<Type, object>>.GetEnumerator() => FastEnumerable().GetEnumerator();

        IEnumerator IEnumerable.GetEnumerator() => FastEnumerable().GetEnumerator();

        private IEnumerable<KeyValuePair<Type, object>> FastEnumerable()
        {
            if (_currentIHttpConnectionFeature != null)
            {
                yield return new KeyValuePair<Type, object>(IHttpConnectionFeatureType, _currentIHttpConnectionFeature);
            }

            if (_currentIConnectionIdFeature != null)
            {
                yield return new KeyValuePair<Type, object>(IConnectionIdFeatureType, _currentIConnectionIdFeature);
            }

            if (_currentIConnectionTransportFeature != null)
            {
                yield return new KeyValuePair<Type, object>(IConnectionTransportFeatureType, _currentIConnectionTransportFeature);
            }

            if (_currentIMemoryPoolFeature != null)
            {
                yield return new KeyValuePair<Type, object>(IMemoryPoolFeatureType, _currentIMemoryPoolFeature);
            }

            if (_currentIApplicationTransportFeature != null)
            {
                yield return new KeyValuePair<Type, object>(IApplicationTransportFeatureType, _currentIApplicationTransportFeature);
            }

            if (_currentITransportSchedulerFeature != null)
            {
                yield return new KeyValuePair<Type, object>(ITransportSchedulerFeatureType, _currentITransportSchedulerFeature);
            }

            if (MaybeExtra != null)
            {
                foreach (var item in MaybeExtra)
                {
                    yield return item;
                }
            }
        }
    }
}