An implementation of the PROXY-protocol (version 2) for .NET/C# Stream objects

StreamExtensions.cs

/*
Copyright 2023 by Addison Schuhardt

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

using System.Buffers.Binary;
using System.Net;

namespace Proxy;

[Flags]
public enum ProxyStatus
{
    /// <summary>
    ///     The proxy header is invalid.  Drop the request and do not use <see cref="ProxyHeader.Source" /> or
    ///     <see cref="ProxyHeader.Destination" />.
    /// </summary>
    Invalid = 0,

    /// <summary>
    ///     The request is coming directly from the proxy server, rather than being proxied from a remote client.  Use the
    ///     connection's actual source and destination addresses rather than these ones.  Do not use
    ///     <see cref="ProxyHeader.Source" /> or <see cref="ProxyHeader.Destination" />.  This usually pertains to
    ///     health-checks and the like.
    /// </summary>
    Local = 1,

    /// <summary>
    ///     The request originated from a remote client and was delivered via a proxy.  This is the typical case.
    /// </summary>
    Proxy = 1 << 1,

    /// <summary>
    ///     The proxy header doesn't specify the type of data it contains.  Ignore it and use the request at your discretion.
    ///     Do not use <see cref="ProxyHeader.Source" /> or <see cref="ProxyHeader.Destination" />.
    /// </summary>
    Unspecified = 1 << 2,

    /// <summary>
    ///     The connection type is unsupported by this implementation.  Currently it only supports TCP via IPv4 and IPv6.
    ///     Ignore the header and use the request at your discretion.  Do not use <see cref="ProxyHeader.Source" /> or
    ///     <see cref="ProxyHeader.Destination" />.
    /// </summary>
    Unsupported = 1 << 3
}

public class ProxyHeader
{
    private readonly IPAddress _destAddress;
    private readonly ushort _destPort;
    private readonly IPAddress _sourceAddress;
    private readonly ushort _sourcePort;

    internal ProxyHeader(IPAddress sourceAddress, IPAddress destAddress, ushort sourcePort, ushort destPort,
        ProxyStatus status)
    {
        _sourceAddress = sourceAddress;
        _destAddress = destAddress;
        _sourcePort = sourcePort;
        _destPort = destPort;
        Status = status;
    }

    public EndPoint Source => new IPEndPoint(_sourceAddress, _sourcePort);
    public EndPoint Destination => new IPEndPoint(_destAddress, _destPort);
    public ProxyStatus Status { get; }
}

internal static class ProxyProtocolExtensions
{
    /// <summary>
    ///     Attempts to read a PROXY-protocol header (version 2) at the start of the data stream.  If a valid header is found,
    ///     it will be recorded in the <see cref="header" /> parameter and the stream will be advanced by the size of the
    ///     header.  If no valid header was found and the stream supports seeking, then the stream will be reset to its
    ///     original position prior to calling this method.
    /// </summary>
    /// <remarks>
    ///     Find the specification for this protocol at
    ///     <a href="https://github.com/haproxy/haproxy/blob/master/doc/proxy-protocol.txt"></a>
    /// </remarks>
    /// <param name="stream">The stream from which to attempt to read the header</param>
    /// <param name="header">The resulting header</param>
    /// <returns>True if a valid header was present in the stream and the stream was advanced, otherwise false</returns>
    public static bool TryReadProxyHeader(this Stream stream, out ProxyHeader header)
    {
        // signature[12] | version[1] | family[1] | size[2] | address block...
        Span<byte> prelude = stackalloc byte[16];
        Span<byte> signature = stackalloc byte[]
        {
            0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A
        };

        stream.ReadExactly(prelude);

        // verify the signature
        for (var i = 0; i < signature.Length; i++)
        {
            if (prelude[i] != signature[i])
            {
                if (stream.CanSeek)
                    stream.Seek(-prelude.Length, SeekOrigin.Current);

                header = BuildEmptyHeader(ProxyStatus.Invalid);
                return false;
            }
        }

        var versionAndType = prelude[12];
        var size = BinaryPrimitives.ReverseEndianness(BitConverter.ToUInt16(prelude[14..]));

        // check the protocol version and type
        switch (versionAndType)
        {
            case 0x20:
                // LOCAL (i.e. health-checking from the proxy server)
                // discard the address block
                header = BuildEmptyHeader(ProxyStatus.Local);
                stream.Seek(size, SeekOrigin.Current);
                return true;
            case 0x21:
                // PROXY (i.e. an actual proxied request; continue on...)
                header = BuildProxyHeader(stream, prelude[13], size);
                return true;
            default:
                // invalid
                header = BuildEmptyHeader(ProxyStatus.Invalid);
                return false;
        }
    }

    private static ProxyHeader BuildProxyHeader(Stream stream, byte family, ushort blockSize)
    {
        ProxyHeader header;

        var skip = 0;
        switch (family)
        {
            case 0x00:
                // unspecified; still valid, up to the caller to decide what to do
                header = BuildEmptyHeader(ProxyStatus.Proxy | ProxyStatus.Unspecified);
                skip = blockSize;
                break;
            case 0x11:
                // TCP over IPv4
                Span<byte> ipv4Block = stackalloc byte[12];
                stream.ReadExactly(ipv4Block);
                header = BuildIPv4Header(ProxyStatus.Proxy, ipv4Block);
                skip = blockSize - ipv4Block.Length;
                break;
            case 0x21:
                // TCP over IPv6
                Span<byte> ipv6Block = stackalloc byte[36];
                stream.ReadExactly(ipv6Block);
                header = BuildIPv6Header(ProxyStatus.Proxy, ipv6Block);
                skip = blockSize - ipv6Block.Length;
                break;
            default:
                // valid, but unsupported by this implementation; discard the remaining buffer
                header = BuildEmptyHeader(ProxyStatus.Proxy | ProxyStatus.Unsupported);
                skip = blockSize;
                break;
        }

        if (skip > 0)
        {
            if (stream.CanSeek)
                stream.Seek(skip, SeekOrigin.Current);
            else
                Skip(stream, skip);
        }

        return header;
    }

    private static ProxyHeader BuildIPv4Header(ProxyStatus status, ReadOnlySpan<byte> block)
    {
        return new ProxyHeader(
            new IPAddress(block[..4]),
            new IPAddress(block.Slice(4, 4)),
            BitConverter.ToUInt16(block.Slice(8, 2)),
            BitConverter.ToUInt16(block.Slice(10, 2)), status);
    }

    private static ProxyHeader BuildIPv6Header(ProxyStatus status, ReadOnlySpan<byte> block)
    {
        return new ProxyHeader(
            new IPAddress(block[..16]),
            new IPAddress(block.Slice(16, 16)),
            BitConverter.ToUInt16(block.Slice(32, 2)),
            BitConverter.ToUInt16(block.Slice(34, 2)), status);
    }

    private static void Skip(Stream stream, int count)
    {
        for (var i = 0; i < count; i++)
            stream.ReadByte();
    }

    private static ProxyHeader BuildEmptyHeader(ProxyStatus status)
    {
        return new ProxyHeader(IPAddress.None, IPAddress.None, IPEndPoint.MaxPort, IPEndPoint.MaxPort, status);
    }
}