LumaTunnel

using System.Net;
using System.Net.Sockets;
using LumaTunnel.Server.Core.Configuration;
using LumaTunnel.Shared.Protocol;

namespace LumaTunnel.Server.Core.Security;

public sealed class TargetPolicy(ServerSettings settings)
{
    private readonly ServerSettings _settings = settings;
    private readonly HashSet<int> _blockedPorts = [.. settings.BlockedPorts];

    public async Task<IPAddress> ResolveAndValidateAsync(string host, int port, CancellationToken cancellationToken = default)
    {
        if (string.IsNullOrWhiteSpace(host) || host.Length > 253 || port is < 1 or > 65535)
            throw new TunnelProtocolException(TunnelErrorCode.InvalidTarget, "Target host or port is invalid.");
        if (_blockedPorts.Contains(port))
            throw new TunnelProtocolException(TunnelErrorCode.TargetBlocked, $"Target port {port} is blocked by server policy.");

        IPAddress[] addresses;
        if (IPAddress.TryParse(host.Trim('[', ']'), out var literal))
        {
            addresses = [literal];
        }
        else
        {
            try
            {
                addresses = await Dns.GetHostAddressesAsync(host, cancellationToken).ConfigureAwait(false);
            }
            catch (SocketException exception)
            {
                throw new TunnelProtocolException(TunnelErrorCode.DnsFailed, $"Unable to resolve target host: {exception.SocketErrorCode}.");
            }
        }

        if (addresses.Length == 0)
            throw new TunnelProtocolException(TunnelErrorCode.DnsFailed, "Target host resolved to no addresses.");

        var allowed = addresses.Where(IsAllowed).ToArray();
        if (allowed.Length != addresses.Length)
            throw new TunnelProtocolException(TunnelErrorCode.TargetBlocked, "Target host resolves to a private or reserved address.");

        return allowed.OrderBy(static ip => ip.AddressFamily == AddressFamily.InterNetwork ? 0 : 1).First();
    }

    public bool IsAllowed(IPAddress address)
    {
        if (_settings.AllowPrivateTargets)
            return !IPAddress.IsLoopback(address);

        if (IPAddress.IsLoopback(address) || address.Equals(IPAddress.Any) || address.Equals(IPAddress.IPv6Any))
            return false;

        if (address.IsIPv4MappedToIPv6)
            address = address.MapToIPv4();

        var bytes = address.GetAddressBytes();
        if (address.AddressFamily == AddressFamily.InterNetwork)
        {
            return !(bytes[0] is 0 or 10 or 127
                     || bytes[0] == 169 && bytes[1] == 254
                     || bytes[0] == 172 && bytes[1] is >= 16 and <= 31
                     || bytes[0] == 192 && bytes[1] == 168
                     || bytes[0] >= 224);
        }

        if (address.AddressFamily == AddressFamily.InterNetworkV6)
        {
            var isUniqueLocal = (bytes[0] & 0xFE) == 0xFC;
            var isLinkLocal = bytes[0] == 0xFE && (bytes[1] & 0xC0) == 0x80;
            var isMulticast = bytes[0] == 0xFF;
            return !isUniqueLocal && !isLinkLocal && !isMulticast && !address.Equals(IPAddress.IPv6Loopback);
        }

        return false;
    }
}