feature(target_chains/ton): add upgrade standard (#1939)

* add execute_upgrade_contract

* add upgrade feature

* precommit

target_chains/ton/contracts/contracts/Main.fc

@@ -13,14 +13,16 @@
 
     ;; * A 32-bit (big-endian) unsigned integer `op`, identifying the `operation` to be performed, or the `method` of the smart contract to be invoked.
     int op = in_msg_body~load_uint(32);
-    ;; * A 64-bit (big-endian) unsigned integer `query_id`, used in all query-response internal messages to indicate that a response is related to a query (the `query_id` of a response must be equal to the `query_id` of the corresponding query). If `op` is not a query-response method (e.g., it invokes a method that is not expected to send an answer), then `query_id` may be omitted.
-    int query_id = in_msg_body~load_uint(64);
+    cell data = in_msg_body~load_ref();
+    slice data_slice = data.begin_parse();
 
     ;; * The remainder of the message body is specific for each supported value of `op`.
     if (op == OP_UPDATE_GUARDIAN_SET) {
-        update_guardian_set(in_msg_body);
+        update_guardian_set(data_slice);
     } elseif (op == OP_EXECUTE_GOVERNANCE_ACTION) {
-        execute_governance_action(in_msg_body);
+        execute_governance_action(data_slice);
+    } elseif (op == OP_UPGRADE_CONTRACT) {
+        execute_upgrade_contract(data);
     } else {
         throw(0xffff); ;; Throw exception for unknown op
     }

target_chains/ton/contracts/contracts/Pyth.fc

@@ -254,8 +254,25 @@ int apply_decimal_expo(int value, int expo) {
     return result;
 }
 
-() execute_upgrade_contract(slice payload) impure {
-    ;; TODO: Implement
+() execute_upgrade_contract(cell new_code) impure {
+    load_data();
+    int hash_code = cell_hash(new_code);
+    throw_unless(ERROR_INVALID_CODE_HASH, upgrade_code_hash == hash_code);
+
+    ;; Set the new code
+    set_code(new_code);
+
+    ;; Set the code continuation to the new code
+    set_c3(new_code.begin_parse().bless());
+
+    ;; Throw an exception to end the current execution
+    ;; The contract will be restarted with the new code
+    throw(0);
+}
+
+() execute_authorize_upgrade_contract(slice payload) impure {
+    int code_hash = payload~load_uint(256);
+    upgrade_code_hash = code_hash;
 }
 
 () execute_authorize_governance_data_source_transfer(slice payload) impure {
@@ -317,8 +334,8 @@ int apply_decimal_expo(int value, int expo) {
 }
 
 () execute_governance_payload(int action, slice payload) impure {
-    if (action == UPGRADE_CONTRACT) {
-        execute_upgrade_contract(payload);
+    if (action == AUTHORIZE_UPGRADE_CONTRACT) {
+        execute_authorize_upgrade_contract(payload);
     } elseif (action == AUTHORIZE_GOVERNANCE_DATA_SOURCE_TRANSFER) {
         execute_authorize_governance_data_source_transfer(payload);
     } elseif (action == SET_DATA_SOURCES) {

target_chains/ton/contracts/contracts/common/errors.fc

@@ -40,6 +40,7 @@ const int ERROR_OLD_GOVERNANCE_MESSAGE = 1033;
 const int ERROR_INVALID_GOVERNANCE_TARGET = 1034;
 const int ERROR_INVALID_GOVERNANCE_MAGIC = 1035;
 const int ERROR_INVALID_GOVERNANCE_MODULE = 1036;
+const int ERROR_INVALID_CODE_HASH = 1037;
 
 ;; Common
-const int ERROR_INSUFFICIENT_GAS = 1037;
+const int ERROR_INSUFFICIENT_GAS = 1038;

target_chains/ton/contracts/contracts/common/governance_actions.fc

@@ -1,4 +1,4 @@
-const int UPGRADE_CONTRACT = 0;
+const int AUTHORIZE_UPGRADE_CONTRACT = 0;
 const int AUTHORIZE_GOVERNANCE_DATA_SOURCE_TRANSFER = 1;
 const int SET_DATA_SOURCES = 2;
 const int SET_FEE = 3;

target_chains/ton/contracts/contracts/common/op.fc

@@ -1,3 +1,4 @@
 const int OP_UPDATE_GUARDIAN_SET = 1;
 const int OP_UPDATE_PRICE_FEEDS = 2;
 const int OP_EXECUTE_GOVERNANCE_ACTION = 3;
+const int OP_UPGRADE_CONTRACT = 4;

target_chains/ton/contracts/contracts/common/storage.fc

@@ -11,6 +11,7 @@ global int single_update_fee;
 global cell data_sources; ;; Dictionary of DataSource tuples, keyed by u8
 global int num_data_sources;
 global cell is_valid_data_source; ;; Dictionary of int (0 as false, -1 as true), keyed by DataSource cell_hash
+global int upgrade_code_hash; ;; 256-bit unsigned integer
 
 
 ;; Wormhole
@@ -54,6 +55,7 @@ global int governance_data_source_index; ;; u32
         .store_ref(governance_data_source)
         .store_uint(last_executed_governance_sequence, 64)
         .store_uint(governance_data_source_index, 32)
+        .store_uint(upgrade_code_hash, 256)
         .end_cell();
 
     begin_cell()
@@ -94,6 +96,7 @@ global int governance_data_source_index; ;; u32
     governance_data_source = governance_slice~load_ref();
     last_executed_governance_sequence = governance_slice~load_uint(64);
     governance_data_source_index = governance_slice~load_uint(32);
+    upgrade_code_hash = governance_slice~load_uint(256);
 
     ds.end_parse();
 }

target_chains/ton/contracts/contracts/tests/PythTest.fc

@@ -20,12 +20,15 @@
 
     int op = in_msg_body~load_uint(32);
     cell data = in_msg_body~load_ref();
+    slice data_slice = data.begin_parse();
     if (op == OP_UPDATE_GUARDIAN_SET) {
-        update_guardian_set(data.begin_parse());
+        update_guardian_set(data_slice);
     } elseif (op == OP_UPDATE_PRICE_FEEDS) {
-        update_price_feeds(msg_value, data.begin_parse());
+        update_price_feeds(msg_value, data_slice);
     } elseif (op == OP_EXECUTE_GOVERNANCE_ACTION) {
-        execute_governance_action(data.begin_parse());
+        execute_governance_action(data_slice);
+    } elseif (op == OP_UPGRADE_CONTRACT) {
+        execute_upgrade_contract(data);
     } else {
         throw(0xffff); ;; Throw exception for unknown op
     }

target_chains/ton/contracts/contracts/tests/PythTestUpgraded.fc

@@ -0,0 +1,78 @@
+{-
+  This test contract is an upgraded version of PythTest.fc. This is used to test the upgrade functionality of the Pyth contract.
+-}
+
+#include "../imports/stdlib.fc";
+#include "../Pyth.fc";
+#include "../Wormhole.fc";
+#include "../common/op.fc";
+
+() recv_internal(int balance, int msg_value, cell in_msg_full, slice in_msg_body) impure {
+    if (in_msg_body.slice_empty?()) {
+        return ();
+    }
+
+    int op = in_msg_body~load_uint(32);
+    cell data = in_msg_body~load_ref();
+    slice data_slice = data.begin_parse();
+    if (op == OP_UPDATE_GUARDIAN_SET) {
+        update_guardian_set(data_slice);
+    } elseif (op == OP_UPDATE_PRICE_FEEDS) {
+        update_price_feeds(msg_value, data_slice);
+    } elseif (op == OP_EXECUTE_GOVERNANCE_ACTION) {
+        execute_governance_action(data_slice);
+    } elseif (op == OP_UPGRADE_CONTRACT) {
+        execute_upgrade_contract(data);
+    } else {
+        throw(0xffff); ;; Throw exception for unknown op
+    }
+}
+
+(int, int, int, int) test_get_price_unsafe(int price_feed_id) method_id {
+    return get_price_unsafe(price_feed_id);
+}
+
+(int, int, int, int) test_get_price_no_older_than(int time_period, int price_feed_id) method_id {
+    return get_price_no_older_than(time_period, price_feed_id);
+}
+
+(int, int, int, int) test_get_ema_price_unsafe(int price_feed_id) method_id {
+    return get_ema_price_unsafe(price_feed_id);
+}
+
+(int, int, int, int) test_get_ema_price_no_older_than(int time_period, int price_feed_id) method_id {
+    return get_ema_price_no_older_than(time_period, price_feed_id);
+}
+
+(int) test_get_update_fee(slice in_msg_body) method_id {
+    return get_update_fee(in_msg_body);
+}
+
+(int) test_get_single_update_fee() method_id {
+    return get_single_update_fee();
+}
+
+(int) test_get_chain_id() method_id {
+    return get_chain_id();
+}
+
+(int) test_get_last_executed_governance_sequence() method_id {
+    return get_last_executed_governance_sequence();
+}
+
+(int) test_get_governance_data_source_index() method_id {
+    return get_governance_data_source_index();
+}
+
+(cell) test_get_governance_data_source() method_id {
+    return get_governance_data_source();
+}
+
+(int) test_get_is_valid_data_source(cell data_source) method_id {
+    return get_is_valid_data_source(data_source);
+}
+
+;; Add a new function to demonstrate the upgrade
+(int) test_new_function() method_id {
+    return 1;
+}

target_chains/ton/contracts/contracts/tests/WormholeTest.fc

@@ -19,8 +19,9 @@
 
     int op = in_msg_body~load_uint(32);
     cell data = in_msg_body~load_ref();
+    slice data_slice = data.begin_parse();
     if (op == OP_UPDATE_GUARDIAN_SET) {
-        update_guardian_set(data.begin_parse());
+        update_guardian_set(data_slice);
     } else {
         throw(0xffff); ;; Throw exception for unknown op
     }

target_chains/ton/contracts/package.json

@@ -18,6 +18,7 @@
     "@ton/ton": "^13.11.2",
     "@types/jest": "^29.5.12",
     "@types/node": "^20.14.10",
+    "@wormhole-foundation/sdk-definitions": "^0.10.7",
     "jest": "^29.7.0",
     "prettier": "^3.3.2",
     "ts-jest": "^29.2.0",

target_chains/ton/contracts/tests/PythTest.spec.ts

@@ -12,10 +12,17 @@ import {
   TEST_GUARDIAN_ADDRESS1,
   PYTH_AUTHORIZE_GOVERNANCE_DATA_SOURCE_TRANSFER,
   PYTH_REQUEST_GOVERNANCE_DATA_SOURCE_TRANSFER,
+  TEST_GUARDIAN_ADDRESS2,
 } from "./utils/pyth";
 import { GUARDIAN_SET_0, MAINNET_UPGRADE_VAAS } from "./utils/wormhole";
 import { DataSource } from "@pythnetwork/xc-admin-common";
-import { parseDataSource } from "./utils";
+import { parseDataSource, createAuthorizeUpgradePayload } from "./utils";
+import {
+  UniversalAddress,
+  createVAA,
+  serialize,
+} from "@wormhole-foundation/sdk-definitions";
+import { mocks } from "@wormhole-foundation/sdk-definitions/testing";
 
 const TIME_PERIOD = 60;
 const PRICE = new Price({
@@ -49,6 +56,11 @@ const TEST_GOVERNANCE_DATA_SOURCES: DataSource[] = [
     emitterAddress:
       "000000000000000000000000000000000000000000000000000000000000002b",
   },
+  {
+    emitterChain: 1,
+    emitterAddress:
+      "0000000000000000000000000000000000000000000000000000000000000000",
+  },
 ];
 
 describe("PythTest", () => {
@@ -343,7 +355,7 @@ describe("PythTest", () => {
       from: deployer.address,
       to: pythTest.address,
       success: false,
-      exitCode: 1037, // ERROR_INSUFFICIENT_GAS
+      exitCode: 1038, // ERROR_INSUFFICIENT_GAS
     });
   });
 
@@ -719,4 +731,144 @@ describe("PythTest", () => {
       exitCode: 1033, // ERROR_OLD_GOVERNANCE_MESSAGE
     });
   });
+
+  it("should successfully upgrade the contract", async () => {
+    // Compile the upgraded contract
+    const upgradedCode = await compile("PythTestUpgraded");
+    const upgradedCodeHash = upgradedCode.hash();
+
+    // Create the authorize upgrade payload
+    const authorizeUpgradePayload =
+      createAuthorizeUpgradePayload(upgradedCodeHash);
+
+    const authorizeUpgradeVaa = createVAA("Uint8Array", {
+      guardianSet: 0,
+      timestamp: 0,
+      nonce: 0,
+      emitterChain: "Solana",
+      emitterAddress: new UniversalAddress(new Uint8Array(32)),
+      sequence: 1n,
+      consistencyLevel: 0,
+      signatures: [],
+      payload: authorizeUpgradePayload,
+    });
+
+    const guardianSet = mocks.devnetGuardianSet();
+    guardianSet.setSignatures(authorizeUpgradeVaa);
+
+    await deployContract(
+      BTC_PRICE_FEED_ID,
+      TIME_PERIOD,
+      PRICE,
+      EMA_PRICE,
+      SINGLE_UPDATE_FEE,
+      DATA_SOURCES,
+      0,
+      [TEST_GUARDIAN_ADDRESS2],
+      1,
+      1,
+      "0000000000000000000000000000000000000000000000000000000000000000",
+      TEST_GOVERNANCE_DATA_SOURCES[2]
+    );
+
+    // Execute the upgrade
+    const sendExecuteGovernanceActionResult =
+      await pythTest.sendExecuteGovernanceAction(
+        deployer.getSender(),
+        Buffer.from(serialize(authorizeUpgradeVaa))
+      );
+
+    expect(sendExecuteGovernanceActionResult.transactions).toHaveTransaction({
+      from: deployer.address,
+      to: pythTest.address,
+      success: true,
+    });
+
+    // Execute the upgrade
+    const sendUpgradeContractResult = await pythTest.sendUpgradeContract(
+      deployer.getSender(),
+      upgradedCode
+    );
+
+    expect(sendUpgradeContractResult.transactions).toHaveTransaction({
+      from: deployer.address,
+      to: pythTest.address,
+      success: true,
+    });
+
+    // Verify that the contract has been upgraded by calling a new method
+    const newMethodResult = await pythTest.getNewFunction();
+    expect(newMethodResult).toBe(1);
+  });
+
+  it("should fail to upgrade the contract with modified code", async () => {
+    // Compile the upgraded contract
+    const upgradedCode = await compile("PythTestUpgraded");
+    const upgradedCodeHash = upgradedCode.hash();
+
+    // Create the authorize upgrade payload
+    const authorizeUpgradePayload =
+      createAuthorizeUpgradePayload(upgradedCodeHash);
+
+    const authorizeUpgradeVaa = createVAA("Uint8Array", {
+      guardianSet: 0,
+      timestamp: 0,
+      nonce: 0,
+      emitterChain: "Solana",
+      emitterAddress: new UniversalAddress(new Uint8Array(32)),
+      sequence: 1n,
+      consistencyLevel: 0,
+      signatures: [],
+      payload: authorizeUpgradePayload,
+    });
+
+    const guardianSet = mocks.devnetGuardianSet();
+    guardianSet.setSignatures(authorizeUpgradeVaa);
+
+    await deployContract(
+      BTC_PRICE_FEED_ID,
+      TIME_PERIOD,
+      PRICE,
+      EMA_PRICE,
+      SINGLE_UPDATE_FEE,
+      DATA_SOURCES,
+      0,
+      [TEST_GUARDIAN_ADDRESS2],
+      1,
+      1,
+      "0000000000000000000000000000000000000000000000000000000000000000",
+      TEST_GOVERNANCE_DATA_SOURCES[2]
+    );
+
+    // Execute the upgrade authorization
+    const sendExecuteGovernanceActionResult =
+      await pythTest.sendExecuteGovernanceAction(
+        deployer.getSender(),
+        Buffer.from(serialize(authorizeUpgradeVaa))
+      );
+
+    expect(sendExecuteGovernanceActionResult.transactions).toHaveTransaction({
+      from: deployer.address,
+      to: pythTest.address,
+      success: true,
+    });
+
+    // Attempt to execute the upgrade with a different code
+    const wormholeTestCode = await compile("WormholeTest");
+    const sendUpgradeContractResult = await pythTest.sendUpgradeContract(
+      deployer.getSender(),
+      wormholeTestCode
+    );
+
+    // Expect the transaction to fail
+    expect(sendUpgradeContractResult.transactions).toHaveTransaction({
+      from: deployer.address,
+      to: pythTest.address,
+      success: false,
+      exitCode: 1037, // ERROR_INVALID_CODE_HASH
+    });
+
+    // Verify that the contract has not been upgraded by attempting to call the new method
+    await expect(pythTest.getNewFunction()).rejects.toThrow();
+  });
 });

target_chains/ton/contracts/tests/utils.ts

@@ -1,5 +1,11 @@
 import { DataSource } from "@pythnetwork/xc-admin-common";
 import { Cell, Transaction, beginCell } from "@ton/core";
+import { Buffer } from "buffer";
+
+const GOVERNANCE_MAGIC = 0x5054474d;
+const GOVERNANCE_MODULE = 1;
+const AUTHORIZE_UPGRADE_CONTRACT_ACTION = 0;
+const TARGET_CHAIN_ID = 1;
 
 export function createCellChain(buffer: Buffer): Cell {
   let chunks = bufferToChunks(buffer, 127);
@@ -67,3 +73,13 @@ export function printTxGasStats(name: string, transaction: Transaction) {
   console.log(`${name} gas cost: ${txComputed.gasFees}`);
   return txComputed.gasFees;
 }
+
+export function createAuthorizeUpgradePayload(newCodeHash: Buffer): Buffer {
+  const payload = Buffer.alloc(8);
+  payload.writeUInt32BE(GOVERNANCE_MAGIC, 0);
+  payload.writeUInt8(GOVERNANCE_MODULE, 4);
+  payload.writeUInt8(AUTHORIZE_UPGRADE_CONTRACT_ACTION, 5);
+  payload.writeUInt16BE(TARGET_CHAIN_ID, 6);
+
+  return Buffer.concat([payload, newCodeHash]);
+}

target_chains/ton/contracts/tests/utils/pyth.ts

@@ -40,6 +40,9 @@ export const BTC_PRICE_FEED_ID =
 export const TEST_GUARDIAN_ADDRESS1 =
   "0x686b9ea8e3237110eaaba1f1b7467559a3273819";
 
+export const TEST_GUARDIAN_ADDRESS2 =
+  "0xbefa429d57cd18b7f8a4d91a2da9ab4af05d0fbe";
+
 // A Pyth governance instruction to authorize governance data source transfer signed by the test guardian #1.
 // From: target_chains/starknet/contracts/tests/data.cairo::pyth_auth_transfer()
 export const PYTH_AUTHORIZE_GOVERNANCE_DATA_SOURCE_TRANSFER =

target_chains/ton/contracts/wrappers/PythTest.ts

@@ -156,8 +156,9 @@ export class PythTest implements Contract {
               .endCell()
           : beginCell().endCell()
       ) // governance_data_source
-      .storeUint(0, 64) // last_executed_governance_sequence
-      .storeUint(0, 32) // governance_data_source_index
+      .storeUint(0, 64) // last_executed_governance_sequence, set to 0 for initial state
+      .storeUint(0, 32) // governance_data_source_index, set to 0 for initial state
+      .storeUint(0, 256) // upgrade_code_hash, set to 0 for initial state
       .endCell();
 
     // Create the main cell with references to grouped data
@@ -350,6 +351,23 @@ export class PythTest implements Contract {
     });
   }
 
+  async sendUpgradeContract(
+    provider: ContractProvider,
+    via: Sender,
+    newCode: Cell
+  ) {
+    const messageBody = beginCell()
+      .storeUint(4, 32) // OP_UPGRADE_CONTRACT
+      .storeRef(newCode)
+      .endCell();
+
+    await provider.internal(via, {
+      value: toNano("0.1"),
+      sendMode: SendMode.PAY_GAS_SEPARATELY,
+      body: messageBody,
+    });
+  }
+
   async getIsValidDataSource(
     provider: ContractProvider,
     dataSource: DataSource
@@ -365,4 +383,9 @@ export class PythTest implements Contract {
     ]);
     return result.stack.readBoolean();
   }
+
+  async getNewFunction(provider: ContractProvider) {
+    const result = await provider.get("test_new_function", []);
+    return result.stack.readNumber();
+  }
 }

target_chains/ton/contracts/wrappers/PythTestUpgraded.compile.ts

@@ -0,0 +1,6 @@
+import { CompilerConfig } from "@ton/blueprint";
+
+export const compile: CompilerConfig = {
+  lang: "func",
+  targets: ["contracts/tests/PythTestUpgraded.fc"],
+};

target_chains/ton/contracts/wrappers/WormholeTest.ts

@@ -85,6 +85,7 @@ export class WormholeTest implements Contract {
       .storeRef(beginCell()) // governance_data_source, empty for initial state
       .storeUint(0, 64) // last_executed_governance_sequence, set to 0 for initial state
       .storeUint(0, 32) // governance_data_source_index, set to 0 for initial state
+      .storeUint(0, 256) // upgrade_code_hash, set to 0 for initial state
       .endCell();
 
     // Create the main cell with references to grouped data