XFESystem

using System.Text;

namespace XFESystem.Network;

/// <summary>Bounded RFC 1035 A/CNAME parser. Compression cycles, truncation and unrelated answers are rejected.</summary>
public static class DnsMessage
{
    public static string NormalizeName(string value)
    {
        string name = value.TrimEnd('.').ToLowerInvariant();
        if (name.Length is 0 or > 253 || name.Split('.').Any(label => label.Length is 0 or > 63 ||
            label.Any(c => c is not (>= 'a' and <= 'z') and not (>= '0' and <= '9') and not '-')))
            throw new ArgumentException("DNS requires a bounded ASCII host name.", nameof(value));
        return name;
    }
    public static int WriteQuery(Span<byte> output, ushort id, string host)
    {
        string name = NormalizeName(host);
        if (output.Length < name.Length + 18) throw new ArgumentException("DNS output is too small.");
        output.Clear(); NetworkPackets.Write16(output, 0, id); NetworkPackets.Write16(output, 2, 0x0100);
        NetworkPackets.Write16(output, 4, 1); int at = 12;
        foreach (string label in name.Split('.'))
        { output[at++] = (byte)label.Length; at += Encoding.ASCII.GetBytes(label, output[at..]); }
        output[at++] = 0; NetworkPackets.Write16(output, at, 1); NetworkPackets.Write16(output, at + 2, 1);
        return at + 4;
    }
    static string Name(ReadOnlySpan<byte> packet, ref int offset)
    {
        int at = offset, resumed = -1, steps = 0; var name = new StringBuilder(64);
        while (true)
        {
            if (++steps > 128 || at >= packet.Length) throw new InvalidDataException("Invalid DNS name or compression cycle.");
            byte length = packet[at++];
            if ((length & 0xc0) == 0xc0)
            {
                if (at >= packet.Length) throw new InvalidDataException("Truncated DNS pointer.");
                int pointer = ((length & 63) << 8) | packet[at++];
                if (pointer >= packet.Length) throw new InvalidDataException("DNS pointer is outside the packet.");
                if (resumed < 0) resumed = at;
                at = pointer; continue;
            }
            if ((length & 0xc0) != 0 || length > 63 || at + length > packet.Length)
                throw new InvalidDataException("Invalid DNS label.");
            if (length == 0) { offset = resumed < 0 ? at : resumed; return name.ToString().ToLowerInvariant(); }
            if (name.Length != 0) name.Append('.');
            for (int i = 0; i < length; i++)
            {
                byte c = packet[at++];
                if (c is not (>= (byte)'a' and <= (byte)'z') and not (>= (byte)'A' and <= (byte)'Z') and
                    not (>= (byte)'0' and <= (byte)'9') and not (byte)'-')
                    throw new InvalidDataException("Invalid DNS host character.");
                name.Append((char)c);
            }
            if (name.Length > 253) throw new InvalidDataException("DNS host name is too long.");
        }
    }
    public static uint ReadAddress(ReadOnlySpan<byte> packet, ushort id, string host)
    {
        if (packet.Length < 12 || NetworkPackets.Read16(packet, 0) != id) throw new InvalidDataException("DNS transaction mismatch.");
        ushort flags = NetworkPackets.Read16(packet, 2);
        if ((flags & 0xf80f) != 0x8000 || (flags & 0x0200) != 0 || NetworkPackets.Read16(packet, 4) != 1)
            throw new InvalidDataException("DNS response is an error, truncated, or has an unsupported opcode.");
        string expected = NormalizeName(host); int at = 12;
        if (Name(packet, ref at) != expected || at + 4 > packet.Length || NetworkPackets.Read16(packet, at) != 1 || NetworkPackets.Read16(packet, at + 2) != 1)
            throw new InvalidDataException("DNS question mismatch.");
        at += 4; int count = NetworkPackets.Read16(packet, 6);
        if (count > 64) throw new InvalidDataException("Too many DNS answers.");
        var aliases = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
        var addresses = new Dictionary<string, uint>(StringComparer.OrdinalIgnoreCase);
        for (int i = 0; i < count; i++)
        {
            string owner = Name(packet, ref at);
            if (at + 10 > packet.Length) throw new InvalidDataException("Truncated DNS answer.");
            ushort type = NetworkPackets.Read16(packet, at), family = NetworkPackets.Read16(packet, at + 2), length = NetworkPackets.Read16(packet, at + 8);
            at += 10; int end = at + length;
            if (end > packet.Length) throw new InvalidDataException("DNS answer exceeds packet.");
            if (family == 1 && type == 1 && length == 4) addresses[owner] = NetworkPackets.Read32(packet, at);
            if (family == 1 && type == 5)
            { int position = at; string alias = Name(packet, ref position); if (position != end) throw new InvalidDataException("Invalid CNAME length."); aliases[owner] = alias; }
            at = end;
        }
        for (int step = 0; step < 8; step++)
        {
            if (addresses.TryGetValue(expected, out uint address) && address != 0 && (address >> 28) < 14) return address;
            if (!aliases.TryGetValue(expected, out string? next)) break;
            expected = next;
        }
        throw new InvalidDataException("DNS has no A address for the requested host/CNAME chain.");
    }
}