SocialPub/PrivaPub/Infrastructure/Backup/ProtectiveMerge.cs

using MongoDB.Bson;
using MongoDB.Driver;

namespace PrivaPub.Infrastructure.Backup
{
	// A restore may lose what was made since the backup, but it never undoes a protective act (owner decision 2026-10-07).
	// From the pre-restore backup P, after the backup is imported:
	// - who follows whom is P's: a follower that left stays gone, so no followers-only post reaches it, and one gained is
	//   kept; the persona's follows are P's too;
	// - blocks, mutes, domain blocks (the server's and the personas'), being blocked, reserved names, deletion tombstones,
	//   reports, filters and OAuth applications are the union of both, P's row winning;
	// - deletions win: a root, persona, group, post or remote account deleted since is deleted again (P's row, as its
	//   deletion left it), with what a deletion takes away; P's moderation of a persona (silenced, suspended, banned) stays;
	// - a root keeps P's password, e-mail, ban and policies;
	// - roots, personas and groups made since become tombstones: deleted, their names kept; local posts made since answer
	//   410; media made since go to the trash, as do media trashed since (and they leave the posts restored with them);
	// - every session ends, and every server's circuit is closed again.
	// Run on the imported database, it reads P's files only, so it can run again after a failed attempt.
	public static class ProtectiveMerge
	{
		static readonly (string Collection, string[] Key)[] Unions =
		[
			("ReservedName", ["Name"]),
			("DomainBlock", ["Domain"]),
			("Block", ["AvatarId", "TargetActorURI"]),
			("Mute", ["AvatarId", "TargetActorURI"]),
			("AccountDomainBlock", ["AvatarId", "Domain"]),
			("BlockedBy", ["AvatarId", "ActorURI"]),
			("DeletedObject", ["ObjectURI"]),
			("Report", ["_id"]),
			("PersonaFilter", ["_id"]),
			("openiddict.applications", ["client_id"])
		];

		const string Since = "restore: made after the backup";

		public static async Task Apply(IMongoDatabase database, string previous, RestoreReport report, CancellationToken token)
		{
			IEnumerable<BsonDocument> P(string collection) => ServerRestore.Read(Path.Combine(previous, collection + ".jsonl.gz"));
			var now = DateTime.UtcNow;

			// who follows whom: P's
			foreach (var collection in new[] { "Follower", "Following" })
			{
				await database.DropCollectionAsync(collection, token);
				await ServerRestore.Insert(database.GetCollection<BsonDocument>(collection), P(collection), token);
			}

			foreach (var (collection, key) in Unions)
				report.ProtectiveKept += await Union(database.GetCollection<BsonDocument>(collection), P(collection), key, token);

			var roots = database.GetCollection<BsonDocument>("RootUser");
			foreach (var root in P("RootUser"))
			{
				var restored = await ById(roots, root["_id"], token);
				if (restored == default)
				{
					// made since: deleted, as RootRemoval leaves a root
					report.RootsTombstoned.Add(root.GetValue("UserName", "").ToString());
					root["UserName"] = $"deleted-{root["_id"]}";
					root["Email"] = BsonNull.Value;
					root["HashedPassword"] = BsonNull.Value;
					root["Policies"] = new BsonArray();
					root["IsBanned"] = false;
					root["IsEmailValidated"] = false;
					root["DeletedAt"] = Deleted(root, "DeletedAt", now);
					await roots.InsertOneAsync(root, cancellationToken: token);
					continue;
				}
				if (IsSet(root, "DeletedAt"))
				{
					if (!IsSet(restored, "DeletedAt"))
						report.RootsDeleted++;
					await roots.ReplaceOneAsync(Id(root["_id"]), root, cancellationToken: token);
					continue;
				}
				await roots.UpdateOneAsync(Id(root["_id"]), Copy(root, "HashedPassword", "Email", "IsEmailValidated", "IsBanned", "Policies",
					"CredentialsChangedAt", "ResetPasswordToken", "ResetPasswordTokenSentAt"), cancellationToken: token);
			}

			var avatars = database.GetCollection<BsonDocument>("Avatar");
			foreach (var avatar in P("Avatar"))
			{
				var restored = await ById(avatars, avatar["_id"], token);
				if (restored == default)
				{
					report.PersonasTombstoned.Add(avatar.GetValue("UserName", "").ToString());
					avatar["DeletionAt"] = Deleted(avatar, "DeletionAt", now);
					await avatars.InsertOneAsync(avatar, cancellationToken: token);
					continue;
				}
				if (IsSet(avatar, "DeletionAt"))
				{
					if (!IsSet(restored, "DeletionAt"))
						report.PersonasDeleted++;
					await avatars.ReplaceOneAsync(Id(avatar["_id"]), avatar, cancellationToken: token);
					// what RootRemoval takes with a persona, which the backup still had
					foreach (var collection in new[] { "PersonaListMember", "PersonaList", "PersonaFilter", "FollowedTag", "ScheduledStatus" })
						await database.GetCollection<BsonDocument>(collection).DeleteManyAsync(new BsonDocument("AvatarId", avatar["_id"].ToString()), token);
					continue;
				}
				var moderated = new[] { "SilencedAt", "SuspendedAt", "BannedAt" }.Where(field => IsSet(avatar, field)).ToArray();
				if (moderated.Length > 0)
					await avatars.UpdateOneAsync(Id(avatar["_id"]), Copy(avatar, moderated), cancellationToken: token);
			}

			var links = database.GetCollection<BsonDocument>("RootToAvatar");
			var linked = (await (await links.FindAsync(FilterDefinition<BsonDocument>.Empty, cancellationToken: token)).ToListAsync(token))
				.Select(l => l.GetValue("AvatarId", BsonNull.Value).ToString()).ToHashSet(StringComparer.Ordinal);
			await ServerRestore.Insert(links, P("RootToAvatar").Where(l => !linked.Contains(l.GetValue("AvatarId", BsonNull.Value).ToString())), token);

			var groups = database.GetCollection<BsonDocument>("Group");
			foreach (var group in P("Group"))
			{
				var restored = await ById(groups, group["_id"], token);
				if (restored == default)
				{
					report.GroupsTombstoned.Add(group.GetValue("UserName", "").ToString());
					group["DeletionAt"] = Deleted(group, "DeletionAt", now);
					await groups.InsertOneAsync(group, cancellationToken: token);
					continue;
				}
				if (!IsSet(group, "DeletionAt"))
					continue;
				if (!IsSet(restored, "DeletionAt"))
					report.GroupsDeleted++;
				await groups.ReplaceOneAsync(Id(group["_id"]), group, cancellationToken: token);
			}

			// remote accounts deleted or gone since: as P has them
			var foreign = database.GetCollection<BsonDocument>("ForeignAvatar");
			foreach (var account in P("ForeignAvatar").Where(a => IsSet(a, "DeletionAt") || IsSet(a, "ForgottenAt")))
				await foreign.ReplaceOneAsync(Id(account["_id"]), account, cancellationToken: token);

			await Posts(database, P("Post"), report, token);
			await Media(database, P("MediaAttachment"), report, token);

			// every session ends: a new stamp for each root's tokens, and the personas' OAuth tokens go
			foreach (var root in await (await roots.FindAsync(FilterDefinition<BsonDocument>.Empty, cancellationToken: token)).ToListAsync(token))
				await roots.UpdateOneAsync(Id(root["_id"]), Builders<BsonDocument>.Update
					.Set("SessionStamp", Guid.NewGuid().ToString("N"))
					.Set("CredentialsChangedAt", now), cancellationToken: token);
			foreach (var collection in new[] { "openiddict.tokens", "openiddict.authorizations" })
				report.SessionsEnded += (int)(await database.GetCollection<BsonDocument>(collection).DeleteManyAsync(FilterDefinition<BsonDocument>.Empty, token)).DeletedCount;

			// what the backup remembered of servers being down is long past
			await database.GetCollection<BsonDocument>("RemoteInstance").UpdateManyAsync(FilterDefinition<BsonDocument>.Empty,
				Builders<BsonDocument>.Update.Set("ConsecutiveFailures", 0).Set("UnavailableUntil", BsonNull.Value), cancellationToken: token);
		}

		// posts deleted since: deleted again, out of timelines and pins; local posts made since: 410
		static async Task Posts(IMongoDatabase database, IEnumerable<BsonDocument> previous, RestoreReport report, CancellationToken token)
		{
			var posts = database.GetCollection<BsonDocument>("Post");
			var gone = database.GetCollection<BsonDocument>("DeletedObject");
			foreach (var batch in previous.Where(p => IsSet(p, "DeletedAt") || IsLocal(p)).Chunk(500))
			{
				var ids = new BsonArray(batch.Select(p => p["_id"]));
				var restored = (await (await posts.FindAsync(new BsonDocument("_id", new BsonDocument("$in", ids)), cancellationToken: token)).ToListAsync(token))
					.ToDictionary(p => p["_id"]);
				foreach (var post in batch)
				{
					if (!restored.TryGetValue(post["_id"], out var mine))
					{
						if (IsLocal(post) && post.GetValue("ObjectURI", BsonNull.Value) is BsonString uri)
						{
							await gone.UpdateOneAsync(new BsonDocument("ObjectURI", uri), Builders<BsonDocument>.Update
								.SetOnInsert("ObjectURI", uri).SetOnInsert("DeletedAt", DateTime.UtcNow), new UpdateOptions { IsUpsert = true }, token);
							report.PostsGone++;
						}
						continue;
					}
					if (!IsSet(post, "DeletedAt") || IsSet(mine, "DeletedAt"))
						continue;
					await posts.ReplaceOneAsync(Id(post["_id"]), post, cancellationToken: token);
					var id = post["_id"].ToString();
					await database.GetCollection<BsonDocument>("TimelineEntry").DeleteManyAsync(
						new BsonDocument("$or", new BsonArray { new BsonDocument("PostId", id), new BsonDocument("ReblogOfPostId", id) }), token);
					await database.GetCollection<BsonDocument>("Pin").DeleteManyAsync(new BsonDocument("PostId", id), token);
					report.PostsDeleted++;
				}
			}
		}

		// media trashed since, and media made since: in the trash (the janitor deletes their files after its grace), and out
		// of the posts restored; media of posts deleted now: in the trash too
		static async Task Media(IMongoDatabase database, IEnumerable<BsonDocument> previous, RestoreReport report, CancellationToken token)
		{
			var media = database.GetCollection<BsonDocument>("MediaAttachment");
			var posts = database.GetCollection<BsonDocument>("Post");
			var now = DateTime.UtcNow;
			foreach (var batch in previous.Chunk(1000))
			{
				var ids = new BsonArray(batch.Select(m => m["_id"]));
				var restored = (await (await media.FindAsync(new BsonDocument("_id", new BsonDocument("$in", ids)), cancellationToken: token)).ToListAsync(token))
					.ToDictionary(m => m["_id"]);
				var made = new List<BsonDocument>();
				foreach (var row in batch)
				{
					if (!restored.TryGetValue(row["_id"], out var mine))
					{
						if (!IsSet(row, "TrashedAt"))
						{
							row["TrashedAt"] = now;
							row["TrashReason"] = Since;
						}
						made.Add(row);
						continue;
					}
					if (!IsSet(row, "TrashedAt") || IsSet(mine, "TrashedAt"))
						continue;
					await media.UpdateOneAsync(Id(row["_id"]), Copy(row, "TrashedAt", "TrashReason"), cancellationToken: token);
					if (mine.GetValue("PostId", BsonNull.Value) is BsonString postId)
						await posts.UpdateOneAsync(Id(ObjectId.TryParse(postId.AsString, out var post) ? post : postId), Builders<BsonDocument>.Update.PullFilter<BsonDocument>("Media",
							new BsonDocument("AttachmentId", row["_id"].ToString())), cancellationToken: token);
					report.MediaTrashed++;
				}
				await ServerRestore.Insert(media, made, token);
				report.MediaTrashed += made.Count(m => m["TrashReason"] == Since);
			}

			var deleted = await (await posts.FindAsync(new BsonDocument("DeletedAt", new BsonDocument("$ne", BsonNull.Value)),
				new FindOptions<BsonDocument> { Projection = new BsonDocument("_id", 1) }, token)).ToListAsync(token);
			foreach (var batch in deleted.Select(p => (BsonValue)p["_id"].ToString()).Chunk(1000))
				report.MediaTrashed += (int)(await media.UpdateManyAsync(
					new BsonDocument { { "PostId", new BsonDocument("$in", new BsonArray(batch)) }, { "TrashedAt", BsonNull.Value } },
					Builders<BsonDocument>.Update.Set("TrashedAt", now).Set("TrashReason", "restore: its post was deleted"), cancellationToken: token)).ModifiedCount;
		}

		// P's rows added, each replacing the restored row with the same key: how many P had
		static async Task<int> Union(IMongoCollection<BsonDocument> collection, IEnumerable<BsonDocument> previous, string[] key, CancellationToken token)
		{
			var count = 0;
			foreach (var batch in previous.Chunk(500))
			{
				var writes = new List<WriteModel<BsonDocument>>();
				foreach (var row in batch)
				{
					var match = new BsonDocument(key.Select(field => new BsonElement(field, row.GetValue(field, BsonNull.Value))));
					writes.Add(new DeleteManyModel<BsonDocument>(match));
					writes.Add(new InsertOneModel<BsonDocument>(row));
				}
				await collection.BulkWriteAsync(writes, new BulkWriteOptions { IsOrdered = true, BypassDocumentValidation = true }, token);
				count += batch.Length;
			}
			return count;
		}

		static bool IsLocal(BsonDocument post) =>
			post.GetValue("GroupUserId", BsonNull.Value) is BsonString && !post.GetValue("IsFederatedCopy", false).ToBoolean();

		static bool IsSet(BsonDocument row, string field) => row.TryGetValue(field, out var value) && !value.IsBsonNull;

		static BsonValue Deleted(BsonDocument row, string field, DateTime now) => IsSet(row, field) ? row[field] : now;

		static FilterDefinition<BsonDocument> Id(BsonValue id) => new BsonDocument("_id", id);

		static async Task<BsonDocument> ById(IMongoCollection<BsonDocument> collection, BsonValue id, CancellationToken token) =>
			await (await collection.FindAsync(Id(id), cancellationToken: token)).FirstOrDefaultAsync(token);

		static UpdateDefinition<BsonDocument> Copy(BsonDocument from, params string[] fields) =>
			Builders<BsonDocument>.Update.Combine(fields.Select(field => Builders<BsonDocument>.Update.Set(field, from.GetValue(field, BsonNull.Value))));
	}
}