Local prison per coordinator (WalletWasabi#13313)

* Simplify logic and code * Prison per coordinator * Fix deserialization --------- Co-authored-by: Turbolay <[email protected]>
Kukks · Aug 15, 2024 · 8886ac6 · 8886ac6
1 parent d899b5d
commit 8886ac6
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 94 deletions.
diff --git a/WalletWasabi.Daemon/Global.cs b/WalletWasabi.Daemon/Global.cs
@@ -134,7 +134,8 @@ public Global(string dataDir, string configFilePath, Config config)
 		WalletManager = new WalletManager(config.Network, DataDir, new WalletDirectories(Config.Network, DataDir), walletFactory);
 		TransactionBroadcaster = new TransactionBroadcaster(Network, BitcoinStore, HttpClientFactory, WalletManager);
 
-		CoinPrison = CoinPrison.CreateOrLoadFromFile(DataDir);
+		var prisonForCoordinator = Path.Combine(DataDir, config.GetCoordinatorUri().Host);
+		CoinPrison = CoinPrison.CreateOrLoadFromFile(prisonForCoordinator);
 		WalletManager.WalletStateChanged += WalletManager_WalletStateChanged;
 	}
 

diff --git a/WalletWasabi/WabiSabi/Client/Banning/CoinPrison.cs b/WalletWasabi/WabiSabi/Client/Banning/CoinPrison.cs
@@ -1,118 +1,56 @@
 using Newtonsoft.Json;
 using System.Collections.Generic;
-using System.Diagnostics.CodeAnalysis;
 using System.IO;
 using System.Linq;
+using NBitcoin;
 using WalletWasabi.Blockchain.TransactionOutputs;
 using WalletWasabi.Helpers;
 using WalletWasabi.Logging;
 using WalletWasabi.Wallets;
 
 namespace WalletWasabi.WabiSabi.Client.Banning;
 
-public class CoinPrison : IDisposable
+public class CoinPrison(string filePath, Dictionary<OutPoint, PrisonedCoinRecord> bannedCoins) : IDisposable
 {
-	// Coins with banning time longer than this will be reduced to a random value between 2 and 4 days.
-	private static readonly int MaxDaysToTrustLocalPrison = 4;
-
-	public CoinPrison(string filePath)
+	enum BanningStatus
 	{
-		FilePath = filePath;
+		Banned,
+		BanningExpired,
+		NonBanned
 	}
 
-	private HashSet<PrisonedCoinRecord> BannedCoins { get; set; } = new();
-	private string FilePath { get; }
-	private object Lock { get; set; } = new();
+	// Coins with banning time longer than this will be reduced to a random value between 2 and 4 days.
+	private static readonly TimeSpan MaxDaysToTrustLocalPrison = TimeSpan.FromDays(4);
 
-	public bool TryGetOrRemoveBannedCoin(SmartCoin coin, [NotNullWhen(true)] out DateTimeOffset? bannedUntil)
-	{
-		lock (Lock)
-		{
-			bannedUntil = null;
-			if (BannedCoins.SingleOrDefault(record => record.Outpoint == coin.Outpoint) is { } record)
-			{
-				if (DateTimeOffset.UtcNow < record.BannedUntil)
-				{
-					bannedUntil = record.BannedUntil;
-					return true;
-				}
-				RemoveBannedCoinNoLock(coin);
-			}
-			return false;
-		}
-	}
+	private readonly object _lock = new();
 
 	public void Ban(SmartCoin coin, DateTimeOffset until)
 	{
-		lock (Lock)
+		lock (_lock)
 		{
-			if (BannedCoins.Any(record => record.Outpoint == coin.Outpoint))
+			if (bannedCoins.ContainsKey(coin.Outpoint) || until < DateTimeOffset.UtcNow)
 			{
 				return;
 			}
 
-			until = ReduceBanningTimeIfNeeded(until);
-			BannedCoins.Add(new(coin.Outpoint, until));
-			coin.BannedUntilUtc = until;
+			var effectiveBanningTime = EffectiveBanningTime(until);
+			bannedCoins.Add(coin.Outpoint, new PrisonedCoinRecord(coin.Outpoint, effectiveBanningTime));
+			coin.BannedUntilUtc = effectiveBanningTime;
 			ToFile();
 		}
 	}
 
-	private void RemoveBannedCoinNoLock(SmartCoin coin)
-	{
-		var recordToRemove = BannedCoins.SingleOrDefault(record => coin.Outpoint == record.Outpoint);
-		if (recordToRemove == null)
-		{
-			Logger.LogError($"Tried to remove {nameof(coin)} from {nameof(BannedCoins)}, but {nameof(coin)} was null.");
-			return;
-		}
-		BannedCoins.Remove(recordToRemove);
-		coin.BannedUntilUtc = null;
-		ToFile();
-	}
-
-	private void ToFile()
+	public bool IsBanned(OutPoint outpoint)
 	{
-		if (string.IsNullOrWhiteSpace(FilePath))
+		lock (_lock)
 		{
-			return;
+			return GetBanningStatus(outpoint) == BanningStatus.Banned;
 		}
-
-		IoHelpers.EnsureFileExists(FilePath);
-		string json = JsonConvert.SerializeObject(BannedCoins, Formatting.Indented);
-		File.WriteAllText(FilePath, json);
-	}
-
-	/// <summary>
-	///	Reduces local banning time, which we save to disk, if it's longer than the <see cref="MaxDaysToTrustLocalPrison"/>.
-	///	This is to avoid saving absurd long banning times like 1-2 years.
-	///	With this, the coin will retry to participate in a CJ in every 2-4 days and see if the coin is still banned or not according to the backend.
-	///	Random values are used for the new banning period so we don't leak information to the coordinator when the coins get released from the local prison.
-	/// </summary>
-	/// <param name="bannedUntil">Banning time according to the backend.</param>
-	/// <returns>New banning period we want to save to file on client side.</returns>
-	private static DateTimeOffset ReduceBanningTimeIfNeeded(DateTimeOffset bannedUntil)
-	{
-		var currentDate = DateTimeOffset.UtcNow;
-		if (bannedUntil > currentDate.AddDays(MaxDaysToTrustLocalPrison))
-		{
-			Random random = new();
-			int minHours = ((MaxDaysToTrustLocalPrison * 24) - 1) / 2;
-			int maxHours = (MaxDaysToTrustLocalPrison * 24) - 1;
-			int randomHours = random.Next(minHours, maxHours);
-			int randomMinutes = random.Next(0, 60);
-			int randomSeconds = random.Next(0, 60);
-
-			return currentDate.AddHours(randomHours).AddMinutes(randomMinutes).AddSeconds(randomSeconds);
-		}
-
-		return bannedUntil;
 	}
 
 	public static CoinPrison CreateOrLoadFromFile(string containingDirectory)
 	{
 		string prisonFilePath = Path.Combine(containingDirectory, "PrisonedCoins.json");
-		HashSet<PrisonedCoinRecord> prisonedCoinRecords = new();
 		try
 		{
 			IoHelpers.EnsureFileExists(prisonFilePath);
@@ -121,38 +59,70 @@ public static CoinPrison CreateOrLoadFromFile(string containingDirectory)
 			if (string.IsNullOrWhiteSpace(data))
 			{
 				Logger.LogDebug("Prisoned coins file is empty.");
-				return new(prisonFilePath);
+				return new(prisonFilePath, []);
 			}
-			prisonedCoinRecords = JsonConvert.DeserializeObject<HashSet<PrisonedCoinRecord>>(data)
+			var prisonedCoinRecords = JsonConvert.DeserializeObject<HashSet<PrisonedCoinRecord>>(data)
 				?? throw new InvalidDataException("Prisoned coins file is corrupted.");
+
+			return new(prisonFilePath, prisonedCoinRecords.ToDictionary(x=> x.Outpoint, x=>x));
 		}
 		catch (Exception exc)
 		{
 			Logger.LogError($"There was an error during loading {nameof(CoinPrison)}. Deleting corrupt file.", exc);
 			File.Delete(prisonFilePath);
+			return new(prisonFilePath, []);
 		}
-
-		foreach (var item in prisonedCoinRecords)
-		{
-			item.BannedUntil = ReduceBanningTimeIfNeeded(item.BannedUntil);
-		}
-
-		return new(prisonFilePath) { BannedCoins = prisonedCoinRecords };
 	}
 
 	public void UpdateWallet(Wallet wallet)
 	{
-		foreach (var coin in wallet.Coins)
+		lock (_lock)
 		{
-			if (TryGetOrRemoveBannedCoin(coin, out var bannedUntil))
+			foreach (var coin in wallet.Coins.Where(c => GetBanningStatus(c.Outpoint) == BanningStatus.BanningExpired))
 			{
-				coin.BannedUntilUtc = bannedUntil;
+				if (bannedCoins.Remove(coin.Outpoint))
+				{
+					coin.BannedUntilUtc = null;
+				}
 			}
+
+			ToFile();
 		}
 	}
 
 	public void Dispose()
 	{
-		ToFile();
+		lock (_lock)
+		{
+			ToFile();
+		}
+	}
+
+	private BanningStatus GetBanningStatus(OutPoint outpoint)
+	{
+		return !bannedCoins.TryGetValue(outpoint, out var bannedCoin)
+			? BanningStatus.NonBanned
+			: DateTimeOffset.UtcNow < bannedCoin.BannedUntil
+				? BanningStatus.Banned
+				: BanningStatus.BanningExpired;
+	}
+
+	private void ToFile()
+	{
+		if (string.IsNullOrWhiteSpace(filePath))
+		{
+			return;
+		}
+
+		IoHelpers.EnsureFileExists(filePath);
+		string json = JsonConvert.SerializeObject(bannedCoins.Values, Formatting.Indented);
+		File.WriteAllText(filePath, json);
+	}
+
+	private static DateTimeOffset EffectiveBanningTime(DateTimeOffset bannedUntil)
+	{
+		var maxBanDateTime = DateTimeOffset.UtcNow + MaxDaysToTrustLocalPrison;
+		var maxBanUnixDateTime = long.Min(maxBanDateTime.ToUnixTimeSeconds(), bannedUntil.ToUnixTimeSeconds());
+		return DateTimeOffset.FromUnixTimeSeconds(maxBanUnixDateTime);
 	}
 }
diff --git a/WalletWasabi/WabiSabi/Client/CoinJoin/Manager/CoinJoinManager.cs b/WalletWasabi/WabiSabi/Client/CoinJoin/Manager/CoinJoinManager.cs
@@ -310,7 +310,7 @@ private async Task<IEnumerable<SmartCoin>> SelectCandidateCoinsAsync(IWallet wal
 			throw new CoinJoinClientException(CoinjoinError.NoCoinsEligibleToMix, "No candidate coins available to mix.");
 		}
 
-		var bannedCoins = coinCandidates.Where(x => CoinPrison.TryGetOrRemoveBannedCoin(x, out _)).ToArray();
+		var bannedCoins = coinCandidates.Where(x => CoinPrison.IsBanned(x.Outpoint)).ToArray();
 		var immatureCoins = coinCandidates.Where(x => x.Transaction.IsImmature(bestHeight)).ToArray();
 		var unconfirmedCoins = coinCandidates.Where(x => !x.Confirmed).ToArray();
 		var excludedCoins = coinCandidates.Where(x => x.IsExcludedFromCoinJoin).ToArray();