From 28a20c8a68e8a397e57b9b03326c578850b8b2a9 Mon Sep 17 00:00:00 2001
From: Credence <hello@zerocredence.com>
Date: Mon, 15 Jul 2024 18:47:21 +0100
Subject: [PATCH] contracts: disallow join after battle ends

---
 contracts/src/systems/combat/contracts.cairo | 52 ++++++++++++++------
 1 file changed, 36 insertions(+), 16 deletions(-)

diff --git a/contracts/src/systems/combat/contracts.cairo b/contracts/src/systems/combat/contracts.cairo
index 300857143..d22448f4b 100644
--- a/contracts/src/systems/combat/contracts.cairo
+++ b/contracts/src/systems/combat/contracts.cairo
@@ -577,8 +577,40 @@ mod combat_systems {
             get!(world, attacking_army_id, EntityOwner).assert_caller_owner(world);
 
             let mut defending_army: Army = get!(world, defending_army_id, Army);
+            if defending_army.battle_id.is_non_zero() {
+                // defending army appears to be in battle
+                // so we want to update the defending army's battle status
+                // to see if the battle has ended. if it has ended, then the 
+                // army will be removed from the battle
+                let mut defending_army_battle: Battle = get!(
+                    world, defending_army.battle_id, Battle
+                );
+                InternalCombatImpl::update_battle_and_army(
+                    world, ref defending_army_battle, ref defending_army
+                );
+                set!(world, (defending_army_battle))
+            }
+            // ensure defending army is not in battle
             defending_army.assert_not_in_battle();
 
+            let troop_config = TroopConfigImpl::get(world);
+            let attacking_army_health: Health = get!(world, attacking_army_id, Health);
+            let defending_army_health: Health = get!(world, defending_army_id, Health);
+            // ensure health invariant checks pass
+            assert!(
+                attacking_army_health.current == attacking_army.troops.full_health(troop_config),
+                "attacking army health sanity check fail"
+            );
+            assert!(
+                defending_army_health.current == defending_army.troops.full_health(troop_config),
+                "defending army health sanity check fail"
+            );
+
+            // ensure both armies are alive
+            attacking_army_health.assert_alive("your army");
+            defending_army_health.assert_alive("the army you are attacking");
+
+            // ensure both armies are in the same location
             let attacking_army_position: Position = get!(world, attacking_army_id, Position);
             let defending_army_position: Position = get!(world, defending_army_id, Position);
             attacking_army_position.assert_same_location(defending_army_position.into());
@@ -609,22 +641,6 @@ mod combat_systems {
             }
 
             // create battle 
-            let troop_config = TroopConfigImpl::get(world);
-            let attacking_army_health: Health = get!(world, attacking_army_id, Health);
-            // health sanity check 
-            assert!(
-                attacking_army_health.current == attacking_army.troops.full_health(troop_config),
-                "attacking army health sanity check fail"
-            );
-            attacking_army_health.assert_alive("Army");
-
-            let defending_army_health: Health = get!(world, defending_army_id, Health);
-            // health sanity check 
-            assert!(
-                defending_army_health.current == defending_army.troops.full_health(troop_config),
-                "defending army health sanity check fail"
-            );
-
             let mut battle: Battle = Default::default();
             battle.entity_id = battle_id;
             battle.attack_army = attacking_army.into();
@@ -668,6 +684,10 @@ mod combat_systems {
             let mut battle: Battle = get!(world, battle_id, Battle);
             battle.update_state();
 
+            if battle.has_ended() {
+                panic!("Battle has ended");
+            }
+
             // ensure caller army is not in battle
             let mut caller_army: Army = get!(world, army_id, Army);
             caller_army.assert_not_in_battle();