From 902463415ef58be5ae17497c005118e6b53c7d97 Mon Sep 17 00:00:00 2001
From: Erwan Or <erwanor@penumbralabs.xyz>
Date: Wed, 10 Apr 2024 22:01:21 -0400
Subject: [PATCH] dex: DiD guard against bad position transitions

---
 .../dex/src/component/position_manager.rs     | 40 ++++++++++++++++++-
 1 file changed, 39 insertions(+), 1 deletion(-)

diff --git a/crates/core/component/dex/src/component/position_manager.rs b/crates/core/component/dex/src/component/position_manager.rs
index 6d6ebdc7c8..dd513c89e0 100644
--- a/crates/core/component/dex/src/component/position_manager.rs
+++ b/crates/core/component/dex/src/component/position_manager.rs
@@ -402,7 +402,8 @@ pub(crate) trait Inner: StateWrite {
 
         let id = new_state.id();
 
-        }
+        // Assert `update_position` state transitions invariants:
+        Self::guard_invalid_transitions(&prev_state, &new_state, &id)?;
 
         // Update the DEX engine indices:
         self.update_position_by_price_index(&prev_state, &new_state, &id)?;
@@ -416,7 +417,44 @@ pub(crate) trait Inner: StateWrite {
         Ok(())
     }
 
+    fn guard_invalid_transitions(
+        prev_state: &Option<Position>,
+        new_state: &Position,
+        id: &position::Id,
+    ) -> Result<()> {
+        use position::State::*;
+
+        if let Some(prev_lp) = prev_state {
+            tracing::debug!(?id, prev = ?prev_lp.state, new = ?new_state.state, "evaluating state transition");
+            match (prev_lp.state, new_state.state) {
+                (Opened, Opened) => {}
+                (Opened, Closed) => {}
+                (Closed, Closed) => { /* no-op but allowed */ }
+                (Closed, Withdrawn { sequence }) => {
+                    ensure!(
+                        sequence == 0,
+                        "withdrawn positions must have their sequence start at zero (found: {})",
+                        sequence
+                    );
+                }
+                (Withdrawn { sequence: old_seq }, Withdrawn { sequence: new_seq }) => {
+                    let expected_seq = old_seq.saturating_add(1);
+                    ensure!(
+                        new_seq == expected_seq,
+                        "withdrawn must increase 1-by-1 (old: {}, new: {}, expected: {})",
+                        old_seq,
+                        new_seq,
+                        expected_seq
+                    );
+                }
+                _ => bail!("invalid transition"),
             }
+        } else {
+            ensure!(
+                matches!(new_state.state, Opened),
+                "fresh positions MUST start in the `Opened` state (found: {:?})",
+                new_state.state
+            );
         }
 
         Ok(())