← All source files

V11 / ca6f28eff240

PublicWebProxy.cs

Lilith/Core/WebResearch/PublicWebProxy.cs

Unchanged from the preceding version
CS
using System.Net;
using System.Net.Sockets;
using System.Text;
using System.Collections.Concurrent;

namespace Lilith.Core.WebResearch;

// All browser traffic, including redirects and subresources, passes through this
// proxy. DNS results are checked and the socket connects to the checked address.
internal sealed class PublicWebProxy : IDisposable
{
    private readonly TcpListener _listener = new(IPAddress.Loopback, 0);
    private readonly CancellationTokenSource _stop = new();
    private readonly ConcurrentDictionary<int, TcpClient> _clients = new();
    private int _next;
    internal int Port => ((IPEndPoint)_listener.LocalEndpoint).Port;
    internal PublicWebProxy() { _listener.Start(32); _ = AcceptAsync(); }
    private async Task AcceptAsync()
    {
        try
        {
            while (!_stop.IsCancellationRequested)
            {
                var client = await _listener.AcceptTcpClientAsync(_stop.Token);
                if (_clients.Count >= 64) { client.Dispose(); continue; }
                int id = Interlocked.Increment(ref _next); _clients[id] = client;
                _ = HandleAsync(id, client);
            }
        }
        catch (Exception) when (_stop.IsCancellationRequested) { }
    }
    private async Task HandleAsync(int id, TcpClient client)
    {
        using (client)
        using (var deadline = CancellationTokenSource.CreateLinkedTokenSource(_stop.Token))
        {
            deadline.CancelAfter(TimeSpan.FromSeconds(45)); var ct = deadline.Token;
            try
            {
                var stream = client.GetStream(); var header = new List<byte>(); var one = new byte[1];
                while (header.Count < 16384)
                {
                    if (await stream.ReadAsync(one, ct) == 0) return;
                    header.Add(one[0]);
                    if (header.Count >= 4 && header[^4] == 13 && header[^3] == 10 && header[^2] == 13 && header[^1] == 10) break;
                }
                string[] lines = Encoding.ASCII.GetString(header.ToArray()).Split("\r\n");
                string[] request = lines[0].Split(' ');
                if (request.Length != 3) return;
                bool tunnel = request[0] == "CONNECT";
                if (!tunnel && request[0] is not ("GET" or "HEAD")) throw new InvalidDataException("Read-only browser proxy.");
                var uri = PublicWeb.ValidateUri(tunnel ? "https://" + request[1] : request[1]);
                var addresses = await Dns.GetHostAddressesAsync(uri.DnsSafeHost, ct);
                if (addresses.Length == 0 || addresses.Any(ip => !PublicWeb.IsPublic(ip))) throw new InvalidDataException("Non-public DNS destination.");
                using var remote = new TcpClient();
                await remote.ConnectAsync(addresses[0], uri.Port, ct);
                var upstream = remote.GetStream();
                if (tunnel) await stream.WriteAsync(Encoding.ASCII.GetBytes("HTTP/1.1 200 Connection Established\r\n\r\n"), ct);
                else
                {
                    string outgoing = request[0] + " " + uri.PathAndQuery + " HTTP/1.1\r\n" + string.Join("\r\n", lines.Skip(1).Where(l => l.Length > 0 && !l.StartsWith("Proxy-", StringComparison.OrdinalIgnoreCase) && !l.StartsWith("Connection:", StringComparison.OrdinalIgnoreCase))) + "\r\nConnection: close\r\n\r\n";
                    await upstream.WriteAsync(Encoding.ASCII.GetBytes(outgoing), ct);
                }
                await Task.WhenAny(CopyAsync(stream, upstream, ct), CopyAsync(upstream, stream, ct));
                deadline.Cancel();
            }
            catch (Exception ex) when (ex is IOException or SocketException or OperationCanceledException or ArgumentException or InvalidDataException or ObjectDisposedException) { }
            finally { _clients.TryRemove(id, out _); }
        }
    }
    private static async Task CopyAsync(Stream from, Stream to, CancellationToken ct)
    {
        byte[] buffer = new byte[16384]; int total = 0;
        try
        {
            while (await from.ReadAsync(buffer, ct) is int read && read > 0)
            {
                total += read; if (total > 8_000_000) break;
                await to.WriteAsync(buffer.AsMemory(0, read), ct);
            }
        }
        catch (Exception ex) when (ex is IOException or OperationCanceledException or ObjectDisposedException) { }
    }
    public void Dispose()
    {
        _stop.Cancel(); _listener.Stop();
        foreach (var client in _clients.Values) client.Dispose();
        _stop.Dispose();
    }
}