MeWeb/src/MeWeb.Core/Knowledge/AsnTable.cs

using System.Buffers.Binary;
using System.IO.Compression;
using System.Net;
using System.Net.Sockets;
using System.Text;

namespace MeWeb.Core.Knowledge;

public class AsnEntry
{
  public uint Number { get; set; }
  public string Country { get; set; }
  public string Name { get; set; }
}

public class AsnRange
{
  public IPAddress Start { get; set; }
  public IPAddress End { get; set; }
  public uint Number { get; set; }
  public string Country { get; set; }
  public string Name { get; set; }
}

public sealed class AsnTable
{
  const uint Magic = 0x314B574D;

  AsnEntry[] systems;
  uint[] v4Start;
  int[] v4System;
  UInt128[] v6Start;
  int[] v6System;

  public DateTimeOffset SourceDate { get; private set; }

  public int Count => systems.Length;

  public AsnEntry Find(IPAddress address)
  {
    if (address is null) return null;
    if (address.IsIPv4MappedToIPv6) address = address.MapToIPv4();
    var index = address.AddressFamily == AddressFamily.InterNetwork
      ? Search(v4Start, BinaryPrimitives.ReadUInt32BigEndian(address.GetAddressBytes()), v4System)
      : Search(v6Start, BinaryPrimitives.ReadUInt128BigEndian(address.GetAddressBytes()), v6System);
    return index < 0 ? null : systems[index];
  }

  public AsnEntry Find(uint number)
  {
    var low = 0;
    var high = systems.Length - 1;
    while (low <= high)
    {
      var middle = (low + high) / 2;
      var found = systems[middle].Number;
      if (found == number) return systems[middle];
      if (found < number) low = middle + 1;
      else high = middle - 1;
    }
    return null;
  }

  public static AsnTable Read(Stream compressed)
  {
    using var brotli = new BrotliStream(compressed, CompressionMode.Decompress);
    using var reader = new BinaryReader(brotli, Encoding.UTF8);
    if (reader.ReadUInt32() != Magic) throw new InvalidDataException("not a me-web? provider table");
    var table = new AsnTable { SourceDate = DateTimeOffset.FromUnixTimeSeconds(reader.ReadInt64()) };
    table.systems = new AsnEntry[reader.Read7BitEncodedInt()];
    for (var i = 0; i < table.systems.Length; i++)
      table.systems[i] = new AsnEntry { Number = (uint)reader.Read7BitEncodedInt64(), Country = reader.ReadString(), Name = reader.ReadString() };

    var count = reader.Read7BitEncodedInt();
    table.v4Start = new uint[count];
    table.v4System = new int[count];
    var previous4 = 0u;
    for (var i = 0; i < count; i++)
    {
      previous4 += (uint)reader.Read7BitEncodedInt64();
      table.v4Start[i] = previous4;
      table.v4System[i] = reader.Read7BitEncodedInt() - 1;
    }

    count = reader.Read7BitEncodedInt();
    table.v6Start = new UInt128[count];
    table.v6System = new int[count];
    var previous6 = UInt128.Zero;
    for (var i = 0; i < count; i++)
    {
      previous6 += ReadBig(reader);
      table.v6Start[i] = previous6;
      table.v6System[i] = reader.Read7BitEncodedInt() - 1;
    }
    return table;
  }

  public static void Write(Stream output, DateTimeOffset sourceDate, IEnumerable<AsnRange> ranges)
  {
    var known = new Dictionary<uint, AsnEntry>();
    var v4Ranges = new List<(UInt128 Start, UInt128 End, uint Number)>();
    var v6Ranges = new List<(UInt128 Start, UInt128 End, uint Number)>();
    foreach (var range in ranges)
    {
      if (range.Number == 0) continue;
      if (!known.ContainsKey(range.Number)) known[range.Number] = new AsnEntry { Number = range.Number, Country = range.Country ?? "", Name = range.Name ?? "" };
      if (range.Start.AddressFamily == AddressFamily.InterNetwork)
        v4Ranges.Add((BinaryPrimitives.ReadUInt32BigEndian(range.Start.GetAddressBytes()), BinaryPrimitives.ReadUInt32BigEndian(range.End.GetAddressBytes()), range.Number));
      else
        v6Ranges.Add((BinaryPrimitives.ReadUInt128BigEndian(range.Start.GetAddressBytes()), BinaryPrimitives.ReadUInt128BigEndian(range.End.GetAddressBytes()), range.Number));
    }
    var systems = known.Values.OrderBy(s => s.Number).ToList();
    var index = systems.Select((s, i) => (s.Number, i)).ToDictionary(p => p.Number, p => p.i);

    using var brotli = new BrotliStream(output, CompressionLevel.SmallestSize, leaveOpen: true);
    using var writer = new BinaryWriter(brotli, Encoding.UTF8);
    writer.Write(Magic);
    writer.Write(sourceDate.ToUnixTimeSeconds());
    writer.Write7BitEncodedInt(systems.Count);
    foreach (var system in systems)
    {
      writer.Write7BitEncodedInt64(system.Number);
      writer.Write(system.Country);
      writer.Write(system.Name);
    }

    var v4 = Flatten(v4Ranges, uint.MaxValue, index);
    writer.Write7BitEncodedInt(v4.Count);
    var previous4 = 0u;
    foreach (var (start, system) in v4)
    {
      writer.Write7BitEncodedInt64((uint)start - previous4);
      writer.Write7BitEncodedInt(system + 1);
      previous4 = (uint)start;
    }

    var v6 = Flatten(v6Ranges, UInt128.MaxValue, index);
    writer.Write7BitEncodedInt(v6.Count);
    var previous6 = UInt128.Zero;
    foreach (var (start, system) in v6)
    {
      WriteBig(writer, start - previous6);
      writer.Write7BitEncodedInt(system + 1);
      previous6 = start;
    }
  }

  static List<(UInt128 Start, int System)> Flatten(List<(UInt128 Start, UInt128 End, uint Number)> ranges, UInt128 max, Dictionary<uint, int> index)
  {
    ranges.Sort((a, b) => a.Start.CompareTo(b.Start));
    var result = new List<(UInt128 Start, int System)>();
    var nextFree = UInt128.Zero;
    var lastSystem = -1;
    var full = false;
    foreach (var (rangeStart, rangeEnd, number) in ranges)
    {
      var range = (Start: rangeStart, End: rangeEnd, System: index[number]);
      if (full || range.End < nextFree) continue;
      var start = range.Start < nextFree ? nextFree : range.Start;
      if (start > nextFree && lastSystem != -1)
      {
        result.Add((nextFree, -1));
        lastSystem = -1;
      }
      if (range.System != lastSystem) result.Add((start, range.System));
      lastSystem = range.System;
      full = range.End == max;
      nextFree = full ? max : range.End + 1;
    }
    if (lastSystem != -1 && !full) result.Add((nextFree, -1));
    return result;
  }

  static int Search<T>(T[] starts, T value, int[] systems) where T : IComparable<T>
  {
    var position = Array.BinarySearch(starts, value);
    if (position < 0) position = ~position - 1;
    return position < 0 ? -1 : systems[position];
  }

  static void WriteBig(BinaryWriter writer, UInt128 value)
  {
    do
    {
      var low = (byte)(value & 0x7F);
      value >>= 7;
      writer.Write(value == UInt128.Zero ? low : (byte)(low | 0x80));
    }
    while (value != UInt128.Zero);
  }

  static UInt128 ReadBig(BinaryReader reader)
  {
    var value = UInt128.Zero;
    var shift = 0;
    while (true)
    {
      var b = reader.ReadByte();
      value |= (UInt128)(b & 0x7F) << shift;
      if ((b & 0x80) == 0) return value;
      shift += 7;
    }
  }
}