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 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 list = ranges.Where(r => r.Number != 0).ToList();
    var systems = list.GroupBy(r => r.Number).OrderBy(g => g.Key)
      .Select(g => new AsnEntry { Number = g.Key, Country = g.First().Country ?? "", Name = g.First().Name ?? "" }).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(list.Where(r => r.Start.AddressFamily == AddressFamily.InterNetwork), a => BinaryPrimitives.ReadUInt32BigEndian(a.GetAddressBytes()), 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(list.Where(r => r.Start.AddressFamily == AddressFamily.InterNetworkV6), a => BinaryPrimitives.ReadUInt128BigEndian(a.GetAddressBytes()), 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(IEnumerable<AsnRange> ranges, Func<IPAddress, UInt128> number, UInt128 max, Dictionary<uint, int> index)
  {
    var result = new List<(UInt128 Start, int System)>();
    var nextFree = UInt128.Zero;
    var lastSystem = -1;
    var full = false;
    foreach (var range in ranges.Select(r => (Start: number(r.Start), End: number(r.End), System: index[r.Number])).OrderBy(r => r.Start))
    {
      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;
    }
  }
}