SocialPub/PrivaPub/Infrastructure/Geo/GeoUpdater.cs

using MaxMind.Db;

using Microsoft.Extensions.Options;

using PrivaPub.Infrastructure.Statistics;

using System.IO.Compression;
using System.Net;

namespace PrivaPub.Infrastructure.Geo
{
	// Keeps the DB-IP Lite databases current by itself, so a deploy needs no timer and no root: once a day it checks
	// whether each database was built this month and, if not, fetches this month's (or, early in the month, last month's),
	// checks it opens as the right kind of database and swaps it in. This, CdnUpdater and SMTP are the server's only traffic
	// that is not federation, which is why each has its own client instead of IFederationHttp.
	public sealed class GeoUpdater : BackgroundService
	{
		public const string ClientName = "geo";
		static readonly string[] Kinds = { "city", "asn" };
		static readonly TimeSpan FirstWait = TimeSpan.FromSeconds(30);
		static readonly TimeSpan Interval = TimeSpan.FromHours(24);

		readonly IHttpClientFactory _http;
		readonly IOptionsMonitor<StatisticsOptions> _options;
		readonly IGeoLocator _locator;
		readonly ILogger<GeoUpdater> _logger;

		public GeoUpdater(IHttpClientFactory http, IOptionsMonitor<StatisticsOptions> options, IGeoLocator locator, ILogger<GeoUpdater> logger)
		{
			_http = http;
			_options = options;
			_locator = locator;
			_logger = logger;
		}

		public static string FileOf(string kind) => $"dbip-{kind}-lite.mmdb";

		protected override async Task ExecuteAsync(CancellationToken stoppingToken)
		{
			try
			{
				await Task.Delay(FirstWait, stoppingToken);
				while (!stoppingToken.IsCancellationRequested)
				{
					if (_options.CurrentValue.Geo.AutoUpdate)
						await Update(DateTime.UtcNow, stoppingToken);
					await Task.Delay(Interval + TimeSpan.FromMinutes(Random.Shared.Next(60)), stoppingToken);
				}
			}
			catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
			{
			}
		}

		// The number of databases replaced. A failure keeps the database already there and is tried again the next day.
		public async Task<int> Update(DateTime now, CancellationToken token)
		{
			var options = _options.CurrentValue;
			var replaced = 0;
			try
			{
				Directory.CreateDirectory(options.GeoDirectory);
			}
			catch (Exception ex) when (ex is IOException or UnauthorizedAccessException)
			{
				_logger.LogWarning(ex, "Cannot create {Directory}; servers are not located", options.GeoDirectory);
				return 0;
			}
			foreach (var kind in Kinds)
			{
				var path = Path.Combine(options.GeoDirectory, FileOf(kind));
				var built = BuiltMonth(path);
				foreach (var month in new[] { Month(now), Month(now).AddMonths(-1) })
				{
					if (built >= month)
						break;
					try
					{
						if (await Install(kind, month, path, options, token))
						{
							replaced++;
							_logger.LogInformation("Installed DB-IP Lite {Kind} {Month:yyyy-MM}", kind, month);
							break;
						}
					}
					catch (Exception ex) when (ex is HttpRequestException or IOException or InvalidDatabaseException or InvalidDataException
						or TaskCanceledException && !token.IsCancellationRequested)
					{
						_logger.LogWarning(ex, "Could not install DB-IP Lite {Kind} {Month:yyyy-MM}", kind, month);
						break;
					}
				}
			}
			if (replaced > 0 && _locator is DbIpLocator locator)
				locator.Reload();
			return replaced;
		}

		static DateTime Month(DateTime at) => new(at.Year, at.Month, 1, 0, 0, 0, DateTimeKind.Utc);

		static DateTime BuiltMonth(string path)
		{
			if (!File.Exists(path))
				return DateTime.MinValue;
			try
			{
				using var reader = new Reader(path, FileAccessMode.MemoryMapped);
				return Month(reader.Metadata.BuildDate);
			}
			catch (Exception ex) when (ex is InvalidDatabaseException or IOException)
			{
				return DateTime.MinValue;
			}
		}

		// False when that month is not published (yet); throws when it is published but cannot be used.
		async Task<bool> Install(string kind, DateTime month, string path, StatisticsOptions options, CancellationToken token)
		{
			var url = $"{options.Geo.DownloadBase.TrimEnd('/')}/dbip-{kind}-lite-{month:yyyy-MM}.mmdb.gz";
			var packed = path + ".gz.part";
			var unpacked = path + ".new";
			try
			{
				using (var response = await _http.CreateClient(ClientName).GetAsync(url, HttpCompletionOption.ResponseHeadersRead, token))
				{
					if (response.StatusCode == HttpStatusCode.NotFound)
						return false;
					response.EnsureSuccessStatusCode();
					await using var source = await response.Content.ReadAsStreamAsync(token);
					await using var target = File.Create(packed);
					await CopyCapped(source, target, options.Geo.MaxDownloadBytes, token);
				}
				await using (var source = new GZipStream(File.OpenRead(packed), CompressionMode.Decompress))
				await using (var target = File.Create(unpacked))
					await CopyCapped(source, target, options.Geo.MaxDatabaseBytes, token);
				using (var reader = new Reader(unpacked, FileAccessMode.MemoryMapped))
					if (!reader.Metadata.DatabaseType.Contains(kind, StringComparison.OrdinalIgnoreCase))
						throw new InvalidDatabaseException($"{url} is a {reader.Metadata.DatabaseType} database, not {kind}");
				File.Move(unpacked, path, overwrite: true);
				return true;
			}
			finally
			{
				File.Delete(packed);
				File.Delete(unpacked);
			}
		}

		static async Task CopyCapped(Stream source, Stream target, long limit, CancellationToken token)
		{
			var buffer = new byte[81920];
			long total = 0;
			int read;
			while ((read = await source.ReadAsync(buffer, token)) > 0)
			{
				total += read;
				if (total > limit)
					throw new InvalidDataException($"more than {limit} bytes");
				await target.WriteAsync(buffer.AsMemory(0, read), token);
			}
		}
	}
}