From 8865e22116dcb8c291caa745de78b4e6241315c8 Mon Sep 17 00:00:00 2001
From: TAMARA LIPOWSKI <data.lipowski@extaccount.com>
Date: Tue, 1 Apr 2025 19:55:38 +0200
Subject: [PATCH] feat: Make EncodingContext.router_address optional

---
 .../evm/strategy_encoder/strategy_encoders.rs |  4 +--
 .../evm/swap_encoder/swap_encoders.rs         | 34 ++++++++++++-------
 src/encoding/models.rs                        |  2 +-
 3 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/src/encoding/evm/strategy_encoder/strategy_encoders.rs b/src/encoding/evm/strategy_encoder/strategy_encoders.rs
index 1dcc3bd..aed32da 100644
--- a/src/encoding/evm/strategy_encoder/strategy_encoders.rs
+++ b/src/encoding/evm/strategy_encoder/strategy_encoders.rs
@@ -207,7 +207,7 @@ impl StrategyEncoder for SplitSwapStrategyEncoder {
                 let encoding_context = EncodingContext {
                     receiver: self.router_address.clone(),
                     exact_out: solution.exact_out,
-                    router_address: self.router_address.clone(),
+                    router_address: Some(self.router_address.clone()),
                     group_token_in: grouped_swap.input_token.clone(),
                     group_token_out: grouped_swap.output_token.clone(),
                 };
@@ -331,7 +331,7 @@ impl StrategyEncoder for ExecutorStrategyEncoder {
             let encoding_context = EncodingContext {
                 receiver: receiver.clone(),
                 exact_out: solution.exact_out,
-                router_address: Bytes::from("0x0000000000000000000000000000000000000000"),
+                router_address: None,
                 group_token_in: grouped_swap.input_token.clone(),
                 group_token_out: grouped_swap.output_token.clone(),
             };
diff --git a/src/encoding/evm/swap_encoder/swap_encoders.rs b/src/encoding/evm/swap_encoder/swap_encoders.rs
index f119681..6a8ef9d 100644
--- a/src/encoding/evm/swap_encoder/swap_encoders.rs
+++ b/src/encoding/evm/swap_encoder/swap_encoders.rs
@@ -229,13 +229,21 @@ impl SwapEncoder for BalancerV2SwapEncoder {
     ) -> Result<Vec<u8>, EncodingError> {
         let token_approvals_manager = ProtocolApprovalsManager::new()?;
         let token = bytes_to_address(&swap.token_in)?;
-        let router_address = bytes_to_address(&encoding_context.router_address)?;
-        let approval_needed = token_approvals_manager.approval_needed(
-            token,
-            router_address,
-            Address::from_str(&self.vault_address)
-                .map_err(|_| EncodingError::FatalError("Invalid vault address".to_string()))?,
-        )?;
+        let approval_needed: bool;
+
+        if let Some(router_address) = encoding_context.router_address {
+            let tycho_router_address = bytes_to_address(
+                &router_address,
+            )?;
+            approval_needed = token_approvals_manager.approval_needed(
+                token,
+                tycho_router_address,
+                Address::from_str(&self.vault_address)
+                    .map_err(|_| EncodingError::FatalError("Invalid vault address".to_string()))?,
+            )?;
+        } else {
+            approval_needed = true;
+        }
 
         let component_id = AlloyBytes::from_str(&swap.component.id)
             .map_err(|_| EncodingError::FatalError("Invalid component ID".to_string()))?;
@@ -350,7 +358,7 @@ mod tests {
         let encoding_context = EncodingContext {
             receiver: Bytes::from("0x0000000000000000000000000000000000000001"),
             exact_out: false,
-            router_address: Bytes::zero(20),
+            router_address: Some(Bytes::zero(20)),
             group_token_in: token_in.clone(),
             group_token_out: token_out.clone(),
         };
@@ -397,7 +405,7 @@ mod tests {
         let encoding_context = EncodingContext {
             receiver: Bytes::from("0x0000000000000000000000000000000000000001"),
             exact_out: false,
-            router_address: Bytes::zero(20),
+            router_address: Some(Bytes::zero(20)),
             group_token_in: token_in.clone(),
             group_token_out: token_out.clone(),
         };
@@ -445,7 +453,7 @@ mod tests {
             // The receiver was generated with `makeAddr("bob") using forge`
             receiver: Bytes::from("0x1d96f2f6bef1202e4ce1ff6dad0c2cb002861d3e"),
             exact_out: false,
-            router_address: Bytes::zero(20),
+            router_address: Some(Bytes::zero(20)),
             group_token_in: token_in.clone(),
             group_token_out: token_out.clone(),
         };
@@ -503,7 +511,7 @@ mod tests {
             receiver: Bytes::from("0x5615deb798bb3e4dfa0139dfa1b3d433cc23b72f"),
             exact_out: false,
             // Same as the executor address
-            router_address: Bytes::from("0x5615deb798bb3e4dfa0139dfa1b3d433cc23b72f"),
+            router_address: Some(Bytes::from("0x5615deb798bb3e4dfa0139dfa1b3d433cc23b72f")),
 
             group_token_in: token_in.clone(),
             group_token_out: token_out.clone(),
@@ -567,7 +575,7 @@ mod tests {
         let encoding_context = EncodingContext {
             receiver: Bytes::from("0x0000000000000000000000000000000000000001"),
             exact_out: false,
-            router_address: Bytes::zero(20),
+            router_address: Some(Bytes::zero(20)),
             group_token_in: group_token_in.clone(),
             // Token out is the same as the group token out
             group_token_out: token_out.clone(),
@@ -606,7 +614,7 @@ mod tests {
         let context = EncodingContext {
             receiver: receiver_address.clone(),
             exact_out: false,
-            router_address: router_address.clone(),
+            router_address: Some(router_address.clone()),
             group_token_in: usde_address.clone(),
             group_token_out: wbtc_address.clone(),
         };
diff --git a/src/encoding/models.rs b/src/encoding/models.rs
index 04c2c44..a8fce6a 100644
--- a/src/encoding/models.rs
+++ b/src/encoding/models.rs
@@ -110,7 +110,7 @@ pub struct Transaction {
 pub struct EncodingContext {
     pub receiver: Bytes,
     pub exact_out: bool,
-    pub router_address: Bytes,
+    pub router_address: Option<Bytes>,
     pub group_token_in: Bytes,
     pub group_token_out: Bytes,
 }