SocialPub/PrivaPub/Domain/Media/MediaProxy.cs

using PrivaPub.Models.Federation;
using System.Collections.Concurrent;
using Microsoft.Extensions.Options;

using MongoDB.Entities;

using PrivaPub.Federation.Actors;
using PrivaPub.Infrastructure.Http;
using PrivaPub.Models.Media;

using System.Security.Cryptography;
using System.Text;

namespace PrivaPub.Domain.Media
{
	public enum ProxyOutcome
	{
		Cached,
		TooBig,
		Failed
	}

	public interface IMediaProxy
	{
		string Wrap(string remoteUrl);
		string Verified(string signature, string encodedUrl);
		/// <summary>Whether a remote file's server is suspended or has its media rejected: nothing of it is proxied.</summary>
		bool Refuses(string url);
		(string Path, string ContentType) Cached(string url);
		/// <summary>Downloads a remote file into the cache, once however many ask at the same time.</summary>
		Task<(ProxyOutcome Outcome, string Path, string ContentType)> Download(string url, CancellationToken token);
		Task<(string Path, string ContentType)> Fetch(string signature, string encodedUrl, CancellationToken token);
		Task<HttpResponseMessage> Open(string url, System.Net.Http.Headers.RangeHeaderValue range, CancellationToken token);
		/// <summary>Deletes what the cache holds of a domain and its subdomains; how many files.</summary>
		int Purge(string domain);
		/// <summary>Trims the cache to its size, oldest first.</summary>
		void Trim();
	}

	// Remote media, fetched for clients so that they never contact a remote server (CLAUDE.md, media invariants):
	// - URLs are HMAC-signed, so only what PrivaPub showed is fetched;
	// - a download is shared by everyone asking for the same URL at once, streamed into a temporary file beside its place
	//   and renamed there (a reader never sees half a file), and never held in memory;
	// - what is too big to cache is remembered for an hour (it is streamed instead), what failed for five minutes;
	// - the cache's size is counted as it grows, and trimmed as soon as it passes the cap;
	// - nothing of a suspended server, or of one whose media are rejected, is proxied, and blocking one purges its files.
	public class MediaProxy : IMediaProxy
	{
		static readonly TimeSpan TooBigFor = TimeSpan.FromHours(1);
		static readonly TimeSpan FailedFor = TimeSpan.FromMinutes(5);
		const int Downloads = 8;//remote files fetched at once, however many clients ask

		readonly ILocalActorService _localActors;
		readonly IFederationHttp _http;
		readonly IMediaService _media;
		readonly IOptionsMonitor<MediaOptions> _options;
		readonly Federation.Moderation.IDomainBlocks _domainBlocks;
		// Lazy: GetOrAdd may run its factory more than once when asked at the same time, but only one Lazy is ever evaluated
		readonly ConcurrentDictionary<string, Lazy<Task<(ProxyOutcome, string, string)>>> _inFlight = new();
		readonly ConcurrentDictionary<string, (ProxyOutcome Outcome, DateTime Until)> _refused = new();
		readonly SemaphoreSlim _downloads = new(Downloads);
		readonly object _keyLock = new();
		byte[] _key;
		long _bytes = -1;//what the cache holds, counted once from the disk and then as it changes
		int _trimming;

		public MediaProxy(ILocalActorService localActors, IFederationHttp http, IMediaService media, IOptionsMonitor<MediaOptions> options,
			Federation.Moderation.IDomainBlocks domainBlocks = default)
		{
			_localActors = localActors;
			_http = http;
			_media = media;
			_options = options;
			_domainBlocks = domainBlocks;
		}

		byte[] Key
		{
			get
			{
				if (_key != default)
					return _key;
				lock (_keyLock)
					return _key ??= LoadKey();
			}
		}

		public string Wrap(string remoteUrl)
		{
			if (string.IsNullOrEmpty(remoteUrl) || remoteUrl.StartsWith(_localActors.BaseAddress + "/", StringComparison.OrdinalIgnoreCase))
				return remoteUrl;
			var encoded = Base64Url(Encoding.UTF8.GetBytes(remoteUrl));
			return $"{_localActors.BaseAddress}/media/proxy/{Sign(remoteUrl)}/{encoded}";
		}

		public string Verified(string signature, string encodedUrl)
		{
			string url;
			try
			{
				url = Encoding.UTF8.GetString(FromBase64Url(encodedUrl));
			}
			catch (FormatException)
			{
				return default;
			}
			return CryptographicOperations.FixedTimeEquals(Encoding.ASCII.GetBytes(signature ?? string.Empty), Encoding.ASCII.GetBytes(Sign(url))) ? url : default;
		}

		public bool Refuses(string url)
		{
			if (_domainBlocks == default || !Uri.TryCreate(url, UriKind.Absolute, out var target))
				return false;
			var block = _domainBlocks.Find(target.Host);
			return block is { Severity: DomainBlockSeverity.Suspend } or { RejectMedia: true };
		}

		public (string Path, string ContentType) Cached(string url)
		{
			var (path, typePath) = CachePaths(url);
			try
			{
				if (!File.Exists(path) || !File.Exists(typePath))
					return default;
				File.SetLastWriteTimeUtc(path, DateTime.UtcNow);
				return (path, File.ReadLines(typePath).FirstOrDefault());
			}
			catch (IOException)
			{
				return default;//trimmed meanwhile
			}
		}

		public Task<HttpResponseMessage> Open(string url, System.Net.Http.Headers.RangeHeaderValue range, CancellationToken token) =>
			_http.OpenMedia(url, range, token);

		public async Task<(string Path, string ContentType)> Fetch(string signature, string encodedUrl, CancellationToken token)
		{
			var url = Verified(signature, encodedUrl);
			if (url == default || Refuses(url))
				return default;
			var (outcome, path, contentType) = await Download(url, token);
			return outcome == ProxyOutcome.Cached ? (path, contentType) : default;
		}

		public async Task<(ProxyOutcome Outcome, string Path, string ContentType)> Download(string url, CancellationToken token)
		{
			if (Cached(url) is { Path: not null } cached)
				return (ProxyOutcome.Cached, cached.Path, cached.ContentType);
			if (_refused.TryGetValue(url, out var refused))
			{
				if (refused.Until > DateTime.UtcNow)
					return (refused.Outcome, default, default);
				_refused.TryRemove(url, out _);
			}
			// shared by everyone asking for it now; the download itself isn't cancelled when one of them leaves
			var shared = _inFlight.GetOrAdd(url, key => new Lazy<Task<(ProxyOutcome, string, string)>>(() => DownloadOnce(key)));
			try
			{
				return await shared.Value.WaitAsync(token);
			}
			finally
			{
				if (shared.Value.IsCompleted)
					_inFlight.TryRemove(new KeyValuePair<string, Lazy<Task<(ProxyOutcome, string, string)>>>(url, shared));
			}
		}

		async Task<(ProxyOutcome, string, string)> DownloadOnce(string url)
		{
			await Task.Yield();
			if (!await _downloads.WaitAsync(TimeSpan.FromSeconds(30)))
				return (ProxyOutcome.Failed, default, default);//too busy: the client tries again
			// cached by a download that ended just before this one began
			if (Cached(url) is { Path: not null } done)
			{
				_downloads.Release();
				return (ProxyOutcome.Cached, done.Path, done.ContentType);
			}
			var (path, typePath) = CachePaths(url);
			var part = $"{path}.{Guid.NewGuid():N}.part";
			try
			{
				Directory.CreateDirectory(System.IO.Path.GetDirectoryName(path)!);
				string contentType, refusal;
				await using (var file = new FileStream(part, FileMode.CreateNew, FileAccess.Write, FileShare.None, 64 * 1024, useAsync: true))
					(contentType, refusal) = await _http.DownloadMedia(url, _options.CurrentValue.MaxProxiedBytes, file, CancellationToken.None);
				if (contentType == default)
				{
					var outcome = refusal == "too-large" ? ProxyOutcome.TooBig : ProxyOutcome.Failed;
					_refused[url] = (outcome, DateTime.UtcNow + (outcome == ProxyOutcome.TooBig ? TooBigFor : FailedFor));
					return (outcome, default, default);
				}
				await File.WriteAllTextAsync(typePath, contentType + "\n" + new Uri(url).Host);
				var size = new FileInfo(part).Length;
				File.Move(part, path, overwrite: true);
				Grew(size);
				return (ProxyOutcome.Cached, path, contentType);
			}
			finally
			{
				_downloads.Release();
				if (File.Exists(part))
					File.Delete(part);
			}
		}

		void Grew(long bytes)
		{
			if (Interlocked.Read(ref _bytes) < 0)
				Interlocked.CompareExchange(ref _bytes, Measure(), -1);
			if (Interlocked.Add(ref _bytes, bytes) > _options.CurrentValue.ProxyCacheBytes)
				_ = Task.Run(Trim);
		}

		public void Trim()
		{
			if (Interlocked.Exchange(ref _trimming, 1) == 1)
				return;
			try
			{
				Interlocked.Exchange(ref _bytes, TrimDirectory(_media.ProxyRoot, _options.CurrentValue.ProxyCacheBytes));
			}
			finally
			{
				Interlocked.Exchange(ref _trimming, 0);
			}
		}

		public int Purge(string domain)
		{
			var directory = new DirectoryInfo(_media.ProxyRoot);
			if (string.IsNullOrEmpty(domain) || !directory.Exists)
				return 0;
			var purged = 0;
			foreach (var type in directory.EnumerateFiles("*.type", SearchOption.AllDirectories))
			{
				try
				{
					var host = File.ReadLines(type.FullName).Skip(1).FirstOrDefault();
					if (host == default || !(host.Equals(domain, StringComparison.OrdinalIgnoreCase) || host.EndsWith("." + domain, StringComparison.OrdinalIgnoreCase)))
						continue;
					var file = new FileInfo(type.FullName[..^".type".Length]);
					if (file.Exists)
						file.Delete();
					type.Delete();
					purged++;
				}
				catch (IOException)
				{
				}
			}
			Interlocked.Exchange(ref _bytes, -1);
			return purged;
		}

		long Measure()
		{
			var directory = new DirectoryInfo(_media.ProxyRoot);
			return directory.Exists ? directory.EnumerateFiles("*", SearchOption.AllDirectories).Where(f => f.Extension is not (".type" or ".part")).Sum(f => f.Length) : 0;
		}

		// deletes the oldest files until the cache is within its cap; what it holds afterwards
		public static long TrimDirectory(string root, long cap)
		{
			var directory = new DirectoryInfo(root);
			if (!directory.Exists)
				return 0;
			var files = directory.EnumerateFiles("*", SearchOption.AllDirectories).Where(f => f.Extension is not (".type" or ".part")).OrderBy(f => f.LastWriteTimeUtc).ToList();
			var total = files.Sum(f => f.Length);
			foreach (var file in files)
			{
				if (total <= cap)
					break;
				try
				{
					total -= file.Length;
					file.Delete();
					var type = new FileInfo(file.FullName + ".type");
					if (type.Exists)
						type.Delete();
				}
				catch (IOException)
				{
				}
			}
			return total;
		}

		(string Path, string TypePath) CachePaths(string url)
		{
			var name = Convert.ToHexStringLower(SHA256.HashData(Encoding.UTF8.GetBytes(url)));
			var path = System.IO.Path.Combine(_media.ProxyRoot, name[..2], name);
			return (path, path + ".type");
		}

		string Sign(string url) => Base64Url(HMACSHA256.HashData(Key, Encoding.UTF8.GetBytes(url))[..16]);

		// the oldest key, so that two first uses at once agree on one
		static byte[] LoadKey()
		{
			var secret = DB.Default.Find<MediaSecret>().Sort(m => m.ID, Order.Ascending).ExecuteFirstAsync().GetAwaiter().GetResult();
			if (secret == default)
			{
				DB.Default.SaveAsync(new MediaSecret { Key = Convert.ToBase64String(RandomNumberGenerator.GetBytes(32)) }).GetAwaiter().GetResult();
				secret = DB.Default.Find<MediaSecret>().Sort(m => m.ID, Order.Ascending).ExecuteFirstAsync().GetAwaiter().GetResult();
			}
			return Convert.FromBase64String(secret.Key);
		}

		static string Base64Url(byte[] bytes) => Convert.ToBase64String(bytes).TrimEnd('=').Replace('+', '-').Replace('/', '_');

		static byte[] FromBase64Url(string value)
		{
			var padded = value.Replace('-', '+').Replace('_', '/');
			return Convert.FromBase64String(padded + new string('=', (4 - padded.Length % 4) % 4));
		}
	}

	public class MediaJanitor : BackgroundService
	{
		static readonly TimeSpan FirstPass = TimeSpan.FromMinutes(5);
		static readonly TimeSpan Interval = TimeSpan.FromHours(1);
		static readonly TimeSpan UnattachedLifetime = TimeSpan.FromDays(1);
		// a trashed file waits this long before it is deleted, already out of what /media/files serves
		public static readonly TimeSpan TrashGrace = TimeSpan.FromDays(1);

		readonly IMediaService _media;
		readonly IOptionsMonitor<MediaOptions> _options;
		readonly ILogger<MediaJanitor> _logger;

		public MediaJanitor(IMediaService media, IOptionsMonitor<MediaOptions> options, ILogger<MediaJanitor> logger)
		{
			_media = media;
			_options = options;
			_logger = logger;
		}

		protected override async Task ExecuteAsync(CancellationToken stoppingToken)
		{
			var wait = FirstPass;
			while (!stoppingToken.IsCancellationRequested)
			{
				try
				{
					await Task.Delay(wait, stoppingToken);
					wait = Interval;
					await Sweep(stoppingToken);
				}
				catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
				{
					return;
				}
				catch (Exception ex)
				{
					_logger.LogWarning(ex, "{Service} pass failed", nameof(MediaJanitor));
				}
			}
		}

		// one pass:
		// - uploads never posted for a day (and not waiting for a scheduled post, nor a profile picture) go to the trash;
		// - trashed files still served (a crash between the mark and the move) are moved out;
		// - trashed files past their grace are deleted, with their rows, unless a backup or a restore is running (one may be
		//   reading the trash);
		// - the proxy cache is trimmed to its size, oldest first
		public async Task Sweep(CancellationToken token)
		{
			var cutoff = DateTime.UtcNow - UnattachedLifetime;
			var unattached = await _media.Trash(m => m.PostId == null && m.ScheduledStatusId == null && m.ProfileOfAvatarId == null && m.CreatedAt < cutoff,
				"never posted", token);

			var trashed = await DB.Default.Find<MediaAttachment>().Match(m => m.TrashedAt != null).Limit(2000).ExecuteAsync(token);
			var purgeBefore = DateTime.UtcNow - TrashGrace;
			var maintenance = await Infrastructure.Backup.MaintenanceLock.Current(token) != default;
			var purged = 0;
			foreach (var row in trashed)
			{
				if (row.TrashedAt < purgeBefore && !maintenance)
				{
					await _media.Purge(row, token);
					purged++;
				}
				else
					_media.Hide(row);
			}
			if (unattached > 0 || purged > 0)
				_logger.LogInformation("{Service}: {Unattached} uploads never posted trashed, {Purged} trashed files deleted",
					nameof(MediaJanitor), unattached, purged);
			TrimProxyCache();
		}

		void TrimProxyCache() => MediaProxy.TrimDirectory(_media.ProxyRoot, _options.CurrentValue.ProxyCacheBytes);
	}
}