From 2ab2db77d61da0dd6e17085aad1a11d626c3100b Mon Sep 17 00:00:00 2001
From: Juan Ignacio Ubeira <juani@balancerlabs.dev>
Date: Wed, 19 Feb 2025 12:19:37 -0300
Subject: [PATCH] Add vault getter to aggregator router (#1311)

---
 pkg/interfaces/contracts/vault/IAggregatorRouter.sol | 4 ++++
 pkg/vault/contracts/AggregatorRouter.sol             | 5 +++++
 pkg/vault/test/foundry/AggregatorsRouter.t.sol       | 5 +++++
 3 files changed, 14 insertions(+)

diff --git a/pkg/interfaces/contracts/vault/IAggregatorRouter.sol b/pkg/interfaces/contracts/vault/IAggregatorRouter.sol
index 501c6f7e5..ca9175d5d 100644
--- a/pkg/interfaces/contracts/vault/IAggregatorRouter.sol
+++ b/pkg/interfaces/contracts/vault/IAggregatorRouter.sol
@@ -4,6 +4,7 @@ pragma solidity ^0.8.24;
 
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 
+import { IVault } from "./IVault.sol";
 import { SwapKind } from "./VaultTypes.sol";
 
 interface IAggregatorRouter {
@@ -13,6 +14,9 @@ interface IAggregatorRouter {
     /// @notice Thrown when the sender does not transfer the correct amount of tokens to the Vault.
     error SwapInsufficientPayment();
 
+    /// @notice Get the address of the Balancer Vault.
+    function getVault() external view returns (IVault);
+
     /**
      * @notice Executes a swap operation specifying an exact input token amount.
      * @param pool Address of the liquidity pool
diff --git a/pkg/vault/contracts/AggregatorRouter.sol b/pkg/vault/contracts/AggregatorRouter.sol
index 3248c51e5..e18b803f8 100644
--- a/pkg/vault/contracts/AggregatorRouter.sol
+++ b/pkg/vault/contracts/AggregatorRouter.sol
@@ -32,6 +32,11 @@ contract AggregatorRouter is IAggregatorRouter, SenderGuard, VaultGuard, Reentra
         // solhint-disable-previous-line no-empty-blocks
     }
 
+    /// @inheritdoc IAggregatorRouter
+    function getVault() public view returns (IVault) {
+        return _vault;
+    }
+
     /***************************************************************************
                                        Swaps
     ***************************************************************************/
diff --git a/pkg/vault/test/foundry/AggregatorsRouter.t.sol b/pkg/vault/test/foundry/AggregatorsRouter.t.sol
index 22c0383cd..b16acb288 100644
--- a/pkg/vault/test/foundry/AggregatorsRouter.t.sol
+++ b/pkg/vault/test/foundry/AggregatorsRouter.t.sol
@@ -74,6 +74,11 @@ contract AggregatorsRouterTest is BaseVaultTest {
         poolArgs = abi.encode(vault, name, symbol);
     }
 
+    function testGetVault() public view {
+        assertNotEq(address(vault), address(0), "Vault not set");
+        assertEq(address(aggregatorsRouter.getVault()), address(vault), "Wrong vault");
+    }
+
     /************************************
                   EXACT IN
     ************************************/