LumaTunnel

using System.Net;
using System.Text;
using LumaTunnel.Shared.Models;

namespace LumaTunnel.Client.Core.Proxy;

internal sealed record HttpProxyRequest(
    string Method,
    string RequestTarget,
    string Version,
    IReadOnlyDictionary<string, string> Headers,
    ReadOnlyMemory<byte> OriginalHeader,
    ReadOnlyMemory<byte> Remaining,
    ReadOnlyMemory<byte> RewrittenHeader)
{
    public bool IsConnect => Method.Equals("CONNECT", StringComparison.OrdinalIgnoreCase);
}

internal static class HttpProxyParser
{
    private const int MaxHeaderBytes = 64 * 1024;

    public static async Task<HttpProxyRequest?> ReadRequestAsync(Stream stream, CancellationToken cancellationToken)
    {
        var buffer = new byte[4096];
        var received = new byte[MaxHeaderBytes];
        var length = 0;
        var headerEnd = -1;
        while (length < received.Length && headerEnd < 0)
        {
            var count = await stream.ReadAsync(buffer, cancellationToken).ConfigureAwait(false);
            if (count == 0)
                return null;
            if (length + count > received.Length)
                throw new InvalidDataException("HTTP proxy request header exceeds 64 KiB.");
            buffer.AsSpan(0, count).CopyTo(received.AsSpan(length));
            length += count;
            headerEnd = FindHeaderEnd(received.AsSpan(0, length));
        }
        if (headerEnd < 0)
            throw new InvalidDataException("Incomplete HTTP proxy request header.");

        var headerBytes = received.AsMemory(0, headerEnd);
        var headerText = Encoding.Latin1.GetString(headerBytes.Span);
        var lines = headerText.Split("\r\n", StringSplitOptions.None);
        var firstLine = lines[0].Split(' ', 3, StringSplitOptions.RemoveEmptyEntries);
        if (firstLine.Length != 3)
            throw new InvalidDataException("Invalid HTTP request line.");

        var headers = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
        foreach (var line in lines.Skip(1))
        {
            var separator = line.IndexOf(':');
            if (separator > 0)
                headers[line[..separator].Trim()] = line[(separator + 1)..].Trim();
        }

        return new HttpProxyRequest(
            firstLine[0],
            firstLine[1],
            firstLine[2],
            headers,
            headerBytes,
            received.AsMemory(headerEnd, length - headerEnd).ToArray(),
            headerBytes);
    }

    public static bool TryParseAuthority(string value, int defaultPort, out ProxyTarget target)
    {
        target = null!;
        value = value.Trim();
        string host;
        var port = defaultPort;
        if (value.StartsWith('['))
        {
            var end = value.IndexOf(']');
            if (end < 2)
                return false;
            host = value[1..end];
            if (end + 1 < value.Length && (value[end + 1] != ':' || !int.TryParse(value[(end + 2)..], out port)))
                return false;
        }
        else
        {
            var separator = value.LastIndexOf(':');
            if (separator > 0 && value.IndexOf(':') == separator)
            {
                host = value[..separator];
                if (!int.TryParse(value[(separator + 1)..], out port))
                    return false;
            }
            else
            {
                host = value;
            }
        }

        if (string.IsNullOrWhiteSpace(host) || port is < 1 or > 65535)
            return false;
        target = new ProxyTarget(host, port, IPAddress.TryParse(host, out var address) ? address : null);
        return true;
    }

    public static byte[] RewriteForOrigin(HttpProxyRequest request, Uri? absoluteUri)
    {
        var target = absoluteUri is null
            ? request.RequestTarget
            : string.IsNullOrEmpty(absoluteUri.PathAndQuery) ? "/" : absoluteUri.PathAndQuery;
        var builder = new StringBuilder().Append(request.Method).Append(' ').Append(target).Append(' ').Append(request.Version).Append("\r\n");
        foreach (var header in request.Headers)
        {
            if (header.Key.Equals("Proxy-Connection", StringComparison.OrdinalIgnoreCase)
                || header.Key.Equals("Proxy-Authorization", StringComparison.OrdinalIgnoreCase))
                continue;
            builder.Append(header.Key).Append(": ").Append(header.Value).Append("\r\n");
        }
        builder.Append("\r\n");
        return Encoding.Latin1.GetBytes(builder.ToString());
    }

    public static Task WriteErrorAsync(Stream stream, int status, string reason, CancellationToken cancellationToken)
    {
        var bytes = Encoding.ASCII.GetBytes($"HTTP/1.1 {status} {reason}\r\nConnection: close\r\nContent-Length: 0\r\nProxy-Agent: LumaTunnel/0.1\r\n\r\n");
        return stream.WriteAsync(bytes, cancellationToken).AsTask();
    }

    private static int FindHeaderEnd(ReadOnlySpan<byte> value)
    {
        for (var i = 3; i < value.Length; i++)
        {
            if (value[i - 3] == '\r' && value[i - 2] == '\n' && value[i - 1] == '\r' && value[i] == '\n')
                return i + 1;
        }
        return -1;
    }
}