From 4bf7e00248337b05ad0f62f9a89f1df8f1f80d42 Mon Sep 17 00:00:00 2001
From: Maksim Strebkov <257byte@gmail.com>
Date: Tue, 26 Mar 2024 03:05:24 +0300
Subject: [PATCH] Count restaking when validating finalize_unstake

---
 .../Commits/Operations/StakingCommit.cs       | 56 ++++++++++++++-----
 1 file changed, 42 insertions(+), 14 deletions(-)

diff --git a/Tzkt.Sync/Protocols/Handlers/Proto18/Commits/Operations/StakingCommit.cs b/Tzkt.Sync/Protocols/Handlers/Proto18/Commits/Operations/StakingCommit.cs
index 8421f892..3037eecb 100644
--- a/Tzkt.Sync/Protocols/Handlers/Proto18/Commits/Operations/StakingCommit.cs
+++ b/Tzkt.Sync/Protocols/Handlers/Proto18/Commits/Operations/StakingCommit.cs
@@ -271,6 +271,14 @@ public async Task Apply(Block block, JsonElement op, JsonElement content)
                                     x.Cycle >= operation.FirstCycleUnstaked.Value &&
                                     x.Cycle <= operation.LastCycleUnstaked.Value)
                                 .SumAsync(x => x.Amount);
+
+                            requestedAmount -= await Db.AutostakingOps
+                                .Where(x =>
+                                    x.BakerId == baker.Id &&
+                                    x.Action == AutostakingAction.Restake &&
+                                    x.Cycle >= operation.FirstCycleUnstaked.Value &&
+                                    x.Cycle <= operation.LastCycleUnstaked.Value)
+                                .SumAsync(x => x.Amount);
                         }
 
                         if (operation.Amount != requestedAmount)
@@ -413,25 +421,45 @@ public async Task Revert(Block block, StakingOperation operation)
                             .ThenBy(x => x.Id)
                             .ToListAsync();
 
+                        var unstakeOps = stakingOps.Select(x => (x.BakerId, x.Amount.Value))
+                            .Concat(delegationOps.Select(x => (x.PrevDelegateId, x.UnstakedBalance.Value + x.UnstakedRewards.Value)));
+
                         // TODO: review in P
-                        var autostakingOps = await Db.AutostakingOps
-                            .AsNoTracking()
-                            .Where(x =>
-                                x.BakerId == sender.Id &&
-                                x.Action == AutostakingAction.Unstake &&
-                                x.Cycle >= operation.FirstCycleUnstaked.Value &&
-                                x.Cycle <= operation.LastCycleUnstaked.Value)
-                            .OrderBy(x => x.Level)
-                            .ThenBy(x => x.Id)
-                            .ToListAsync();
+                        if (sender is Data.Models.Delegate baker)
+                        {
+                            var autostakingOps = await Db.AutostakingOps
+                                .AsNoTracking()
+                                .Where(x =>
+                                    x.BakerId == baker.Id &&
+                                    x.Action == AutostakingAction.Unstake &&
+                                    x.Cycle >= operation.FirstCycleUnstaked.Value &&
+                                    x.Cycle <= operation.LastCycleUnstaked.Value)
+                                .OrderBy(x => x.Level)
+                                .ThenBy(x => x.Id)
+                                .ToListAsync();
 
-                        var unstakeOps = stakingOps.Select(x => (x.BakerId, x.Amount.Value))
-                            .Concat(delegationOps.Select(x => (x.PrevDelegateId, x.UnstakedBalance.Value + x.UnstakedRewards.Value)))
-                            .Concat(autostakingOps.Select(x => ((int?)x.BakerId, x.Amount)));
+                            var autostakingOps2 = await Db.AutostakingOps
+                                .AsNoTracking()
+                                .Where(x =>
+                                    x.BakerId == baker.Id &&
+                                    x.Action == AutostakingAction.Restake &&
+                                    x.Cycle >= operation.FirstCycleUnstaked.Value &&
+                                    x.Cycle <= operation.LastCycleUnstaked.Value)
+                                .OrderBy(x => x.Level)
+                                .ThenBy(x => x.Id)
+                                .ToListAsync();
+
+                            unstakeOps = unstakeOps
+                                .Concat(autostakingOps.Select(x => ((int?)x.BakerId, x.Amount)))
+                                .Concat(autostakingOps2.Select(x => ((int?)x.BakerId, -x.Amount)));
+                        }
+
+                        if (unstakeOps.Any() && unstakeOps.Sum(x => x.Item2) != operation.Amount)
+                            throw new NotImplementedException("Manual staking is not fully implemented in O2, because it's partially forbidden...");
 
                         foreach (var (prevBakerId, unstakedAmount) in unstakeOps)
                         {
-                            if (sender.UnstakedBalance == 0 && unstakedAmount > 0)
+                            if (sender.UnstakedBalance == 0 && unstakedAmount != 0)
                                 sender.UnstakedBakerId = prevBakerId;
 
                             sender.UnstakedBalance += unstakedAmount;