From e10e6e3a505f75a5ff7789888977683f4a40972e Mon Sep 17 00:00:00 2001
From: Roshan <19766713+rpalakkal@users.noreply.github.com>
Date: Sat, 24 Feb 2024 20:31:54 -0500
Subject: [PATCH] fix: specifying caller for Query.send (#8)

---
 src/AxiomTest.sol         |  9 ++++++---
 src/AxiomVm.sol           | 18 +++++-------------
 test/AverageBalance.t.sol | 13 +++++++++----
 3 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/src/AxiomTest.sol b/src/AxiomTest.sol
index 4057ff7..5d3fbe1 100644
--- a/src/AxiomTest.sol
+++ b/src/AxiomTest.sol
@@ -107,7 +107,8 @@ abstract contract AxiomTest is Test {
             input,
             callbackTarget,
             bytes(""),
-            IAxiomV2Query.AxiomV2FeeData({ maxFeePerGas: 25 gwei, callbackGasLimit: 1_000_000, overrideAxiomQueryFee: 0 })
+            IAxiomV2Query.AxiomV2FeeData({ maxFeePerGas: 25 gwei, callbackGasLimit: 1_000_000, overrideAxiomQueryFee: 0 }),
+            msg.sender
         );
     }
 
@@ -117,12 +118,14 @@ abstract contract AxiomTest is Test {
     /// @param callbackTarget The address of the contract to send a callback to
     /// @param callbackExtraData Extra data to include in the callback
     /// @param feeData The fee data for the query
+    /// @param caller the address of the caller
     function query(
         bytes32 _querySchema,
         bytes memory input,
         address callbackTarget,
         bytes memory callbackExtraData,
-        IAxiomV2Query.AxiomV2FeeData memory feeData
+        IAxiomV2Query.AxiomV2FeeData memory feeData,
+        address caller
     ) internal view returns (Query memory) {
         return Query({
             querySchema: _querySchema,
@@ -132,7 +135,7 @@ abstract contract AxiomTest is Test {
             feeData: feeData,
             axiomVm: axiomVm,
             outputString: "",
-            caller: msg.sender
+            caller: caller
         });
     }
 }
diff --git a/src/AxiomVm.sol b/src/AxiomVm.sol
index 15561b2..d85f557 100644
--- a/src/AxiomVm.sol
+++ b/src/AxiomVm.sol
@@ -79,18 +79,8 @@ library Axiom {
     /// @param self The query to send
     function send(Query memory self) public {
         self.outputString = self.axiomVm.getArgsAndSendQuery(
-            self.querySchema, self.input, self.callbackTarget, self.callbackExtraData, self.feeData
-        );
-    }
-
-    /// @dev Sends a query to Axiom
-    /// @param self The query to send
-    /// @param caller The address of the caller of the original query into Axiom
-    function send(Query memory self, address caller) public {
-        self.outputString = self.axiomVm.getArgsAndSendQuery(
-            self.querySchema, self.input, self.callbackTarget, self.callbackExtraData, self.feeData
+            self.querySchema, self.input, self.callbackTarget, self.callbackExtraData, self.feeData, self.caller
         );
-        self.caller = caller;
     }
 
     /// @dev Pranks a callback from Axiom
@@ -434,6 +424,7 @@ contract AxiomVm is Test {
      * @param callbackTarget the callback contract address
      * @param callbackExtraData extra data to be passed to the callback contract
      * @param feeData the fee data
+     * @param caller the address of the caller
      * @return queryString the query string
      */
     function getArgsAndSendQuery(
@@ -441,12 +432,13 @@ contract AxiomVm is Test {
         bytes memory input,
         address callbackTarget,
         bytes memory callbackExtraData,
-        IAxiomV2Query.AxiomV2FeeData memory feeData
+        IAxiomV2Query.AxiomV2FeeData memory feeData,
+        address caller
     ) public returns (string memory queryString) {
         (QueryArgs memory args, string memory _queryString) =
             sendQueryArgs(querySchema, input, callbackTarget, callbackExtraData, feeData);
         queryString = _queryString;
-        vm.prank(msg.sender);
+        vm.prank(caller);
         IAxiomV2Query(axiomV2QueryAddress).sendQuery{ value: args.value }(
             args.sourceChainId,
             args.dataQueryHash,
diff --git a/test/AverageBalance.t.sol b/test/AverageBalance.t.sol
index 798f273..8a4ceb4 100644
--- a/test/AverageBalance.t.sol
+++ b/test/AverageBalance.t.sol
@@ -62,12 +62,17 @@ contract AverageBalanceTest is AxiomTest {
             overrideAxiomQueryFee: 0
         });
 
-        // create a query into Axiom with custom `callbackExtraData` and `feeData`
-        Query memory q = query(querySchema, abi.encode(input), address(averageBalance), callbackExtraData, feeData);
+        // create a query into Axiom with custom `callbackExtraData`, `feeData`, and `caller`
+        address caller = address(123);
+        vm.deal(caller, 1 ether);
+        Query memory q = query(querySchema, abi.encode(input), address(averageBalance), callbackExtraData, feeData, caller);
 
         // send the query to Axiom
-        address caller = address(123);
-        q.send(caller);
+        vm.expectEmit(false, false, false, false);
+        emit QueryFeeInfoRecorded(0, address(0), 0, 0, 0, 0);
+        vm.expectEmit(true, false, false, false);
+        emit QueryInitiatedOnchain(caller, bytes32(0), 0, bytes32(0), caller, address(averageBalance), hex"");
+        q.send();
 
         // prank fulfillment of the query, returning the Axiom results
         vm.expectEmit();