feat: Return Jettons, that were sent by accident

.husky/pre-commit

@@ -1,4 +1,4 @@
 yarn fmt
-yarn tact-fmt
+yarn tact-fmt --write ./
 yarn deduplicate
 yarn lint-staged

sources/contracts/utils/utils.tact

@@ -1,5 +1,9 @@
 asm fun emptyAddress(): Address { b{00} PUSHSLICE }
 
+asm fun sliceWithOneZeroBit(): Slice {
+    b{0} PUSHSLICE
+}
+
 asm fun muldiv(x: Int, y: Int, z: Int): Int { MULDIV }
 
 // CALLCC resets c0 and c1 to default quit-cont (extraordinary continuation),

sources/contracts/vaults/jetton-vault.tact

@@ -7,6 +7,7 @@ import "./ton-vault.tact";
 import "./proofs/check-proof";
 import "./proofs/tep-89-proofer";
 import "../core/liquidity-deposit";
+import "../utils/utils";
 
 // TEP-74 JettonNotify, but with forwardPayload serialized as expected fields
 message(0x7362d09c) JettonNotifyWithActionRequest {
@@ -48,16 +49,56 @@ contract JettonVault(
     receive(msg: JettonNotifyWithActionRequest) {
         if (self.jettonWallet == null) {
             // This function proofs that jetton wallet is valid
-            // This function throws if proof is invalid
-            checkProof(self.jettonMaster, msg.proofType, msg.proof, msg.toCell());
+            // This function throws if proof is invalid and tries to send Jettons back
+            let proofCheckRes = checkProof(self.jettonMaster, msg.proofType, msg.proof, msg.toCell());
+
+            // Proof is not valid, or not supported, let's try to send jettons back
+            if (!proofCheckRes) {
+                message(MessageParameters {
+                    mode: SendRemainingValue | SendIgnoreErrors,
+                    body: SendViaJettonTransfer {
+                        queryId: msg.queryId,
+                        amount: msg.amount,
+                        destination: msg.sender,
+                        responseDestination: msg.sender,
+                        customPayload: null,
+                        forwardTonAmount: 0,
+                        forwardPayload: sliceWithOneZeroBit(),
+                    }.toCell(),
+                    value: 0,
+                    to: sender(),
+                    bounce: true,
+                });
+                commit();
+                require(false, "JettonVault: Proof is invalid");
+            }
             self.jettonWallet = sender();
+
             // There is no sense to bounce, as Jettons won't be sent back
             // However, we can save jettonWallet address to the storage, as the proof succeeded
-
             saveState(self.jettonWallet, self.jettonMaster);
             commit();
         }
-        require(sender() == self.jettonWallet, "JettonVault: Sender must be jetton wallet");
+        // Maybe someone messed up with address, so let's try to send jettons back
+        if (sender() != self.jettonWallet) {
+            message(MessageParameters {
+                mode: SendRemainingValue | SendIgnoreErrors,
+                body: SendViaJettonTransfer {
+                    queryId: msg.queryId,
+                    amount: msg.amount,
+                    destination: msg.sender,
+                    responseDestination: msg.sender,
+                    customPayload: null,
+                    forwardTonAmount: 0,
+                    forwardPayload: sliceWithOneZeroBit(),
+                }.toCell(),
+                value: 0,
+                to: sender(),
+                bounce: true,
+            });
+            commit();
+            require(false, "JettonVault: Sender must be jetton wallet");
+        }
         actionHandler(msg);
     }
 

sources/contracts/vaults/proofs/check-proof.tact

@@ -14,7 +14,7 @@ const PROOF_STATE_INIT: Int = 2;
 const PROOF_STATE_TO_THE_BLOCK: Int = 3;
 const PROOF_JETTON_BURN: Int = 4;
 
-inline fun checkProof(jettonMaster: Address, proofType: Int, proof: Slice, msgCell: Cell) {
+inline fun checkProof(jettonMaster: Address, proofType: Int, proof: Slice, msgCell: Cell): Bool {
     if (proofType == PROOF_TEP89) {
         let prooferStateInit = initOf TEP89Proofer(
             jettonMaster,
@@ -48,14 +48,13 @@ inline fun checkProof(jettonMaster: Address, proofType: Int, proof: Slice, msgCe
             }) == jettonMaster,
             "JettonVault: StateInit proof is invalid",
         );
-        return;
+        return true;
     } else if (proofType == PROOF_STATE_TO_THE_BLOCK) {
         require(false, "JettonVault: State proof is not supported");
-        return;
+        return true;
     } else if (proofType == PROOF_JETTON_BURN) {
         require(false, "JettonVault: Burn proof is not supported");
-        return;
+        return true;
     }
-
-    require(false, "JettonVault: Invalid proof type");
+    return false;
 }

sources/contracts/vaults/ton-vault.tact

@@ -3,6 +3,24 @@ import "../math";
 import "../core/messages.tact";
 import "../core/liquidity-deposit";
 import "../core/amm-pool";
+import "../utils/utils";
+
+message(0xf8a7ea5) ReturnJettonsViaJettonTransfer {
+    queryId: Int as uint64;
+    amount: Int as coins;
+    destination: Address;
+    responseDestination: Address?;
+    customPayload: Cell?;
+    forwardTonAmount: Int as coins;
+    forwardPayload: Slice as remaining;
+}
+
+message(0x7362d09c) UnexpectedJettonNotification {
+    queryId: Int as uint64;
+    amount: Int as coins;
+    sender: Address;
+    forwardPayload: Slice as remaining;
+}
 
 contract TonVault(
     admin: Address,
@@ -44,6 +62,27 @@ contract TonVault(
         });
     }
 
+    // Someone possibly transferred us jettons by accident
+    receive(msg: UnexpectedJettonNotification) {
+        message(MessageParameters {
+            mode: SendRemainingValue,
+            value: 0,
+            body: ReturnJettonsViaJettonTransfer {
+                queryId: msg.queryId,
+                amount: msg.amount,
+                destination: msg.sender,
+                responseDestination: msg.sender,
+                customPayload: null,
+                forwardTonAmount: 1,
+                forwardPayload: sliceWithOneZeroBit(),
+            }.toCell(),
+            to: sender(),
+            bounce: true,
+        });
+        commit();
+        require(false, "TonVault: Jetton transfer must be performed to correct Jetton Vault");
+    }
+
     receive(msg: AddLiquidityPartTon) {
         // TODO: exact tests for this
         nativeReserve(msg.amountIn, ReserveExact | ReserveAddOriginalBalance);

sources/tests/proofs.spec.ts

@@ -134,4 +134,125 @@ describe("Proofs", () => {
         )
         expect(await jettonVaultInstance.getInited()).toBe(false)
     })
+    test("Jettons are returned if proof type is incorrect", async () => {
+        const blockchain = await Blockchain.create()
+        const vaultSetup = await createJettonVault(blockchain)
+
+        const _ = await vaultSetup.deploy()
+        const mockActionPayload = beginCell()
+            .storeStringTail("Random action that does not mean anything")
+            .endCell()
+
+        const sendNotifyWithNoProof = await vaultSetup.treasury.transfer(
+            vaultSetup.vault.address,
+            toNano(0.5),
+            createJettonVaultMessage(
+                // We can use any Jetton Vault opcode here because we don't need an actual operation here
+                LPDepositPartOpcode,
+                mockActionPayload,
+                {
+                    proofType: 0n, // No proof attached
+                },
+            ),
+        )
+
+        const toVaultTx = flattenTransaction(
+            findTransactionRequired(sendNotifyWithNoProof.transactions, {
+                to: vaultSetup.vault.address,
+                op: JettonVault.opcodes.JettonNotifyWithActionRequest,
+                success: true, // Because commit was called
+                exitCode: JettonVault.errors["JettonVault: Proof is invalid"],
+            }),
+        )
+
+        expect(sendNotifyWithNoProof.transactions).toHaveTransaction({
+            from: vaultSetup.vault.address,
+            to: toVaultTx.from,
+            op: JettonVault.opcodes.JettonTransfer,
+            success: true,
+        })
+
+        expect(await vaultSetup.isInited()).toBe(false)
+    })
+
+    test("Jettons are returned if sent to wrong vault", async () => {
+        const blockchain = await Blockchain.create()
+        // Create and set up a correct jetton vault
+        const vaultSetup = await createJettonVault(blockchain)
+        const _ = await vaultSetup.deploy()
+
+        // Create a different jetton (wrong one) for testing
+        const wrongJetton = await createJetton(blockchain)
+
+        // Get the initial balance of the wrong jetton wallet
+        const initialWrongJettonBalance = await wrongJetton.wallet.getJettonBalance()
+
+        // Create a mock payload to use with the transfer
+        const mockPayload = beginCell()
+            .store(
+                storeLPDepositPart({
+                    $$type: "LPDepositPart",
+                    liquidityDepositContract: randomAddress(0), // Mock LP contract address
+                    additionalParams: {
+                        $$type: "AdditionalParams",
+                        minAmountToDeposit: 0n,
+                        lpTimeout: 0n,
+                        payloadOnSuccess: null,
+                        payloadOnFailure: null,
+                    },
+                }),
+            )
+            .endCell()
+
+        // Number of jettons to send to the wrong vault
+        const amountToSend = toNano(0.5)
+
+        // First, we need to initialize the vault with the correct jettons
+        const _initVault = await vaultSetup.treasury.transfer(
+            vaultSetup.vault.address,
+            amountToSend,
+            createJettonVaultMessage(
+                // We can use any Jetton Vault opcode here because we don't need an actual operation here
+                LPDepositPartOpcode,
+                mockPayload,
+                {
+                    proofType: PROOF_TEP89,
+                },
+            ),
+        )
+        expect(await vaultSetup.isInited()).toBeTruthy()
+
+        // Send wrong Jetton to the vault
+        const sendJettonsToWrongVault = await wrongJetton.transfer(
+            vaultSetup.vault.address,
+            amountToSend,
+            createJettonVaultMessage(LPDepositPartOpcode, mockPayload, {
+                proofType: PROOF_TEP89,
+            }),
+        )
+
+        // Verify that the transaction to the vault has occurred but failed due to the wrong jetton
+        const toVaultTx = flattenTransaction(
+            findTransactionRequired(sendJettonsToWrongVault.transactions, {
+                to: vaultSetup.vault.address,
+                op: JettonVault.opcodes.JettonNotifyWithActionRequest,
+                success: true, // Because commit was called
+                exitCode: JettonVault.errors["JettonVault: Sender must be jetton wallet"],
+            }),
+        )
+
+        // Check that the jettons were sent back to the original wallet
+        expect(sendJettonsToWrongVault.transactions).toHaveTransaction({
+            from: vaultSetup.vault.address,
+            to: toVaultTx.from,
+            op: JettonVault.opcodes.JettonTransfer,
+            success: true,
+        })
+
+        expect(await vaultSetup.isInited()).toBeTruthy()
+
+        // Verify that the balance of the wrong jetton wallet is unchanged (jettons returned)
+        const finalWrongJettonBalance = await wrongJetton.wallet.getJettonBalance()
+        expect(finalWrongJettonBalance).toEqual(initialWrongJettonBalance)
+    })
 })

sources/tests/ton-vault.spec.ts

@@ -0,0 +1,46 @@
+import {Blockchain} from "@ton/sandbox"
+import {createJetton, createTonVault} from "../utils/environment"
+import {beginCell} from "@ton/core"
+import {findTransactionRequired, flattenTransaction} from "@ton/test-utils"
+import {randomInt} from "node:crypto"
+import {TonVault} from "../output/DEX_TonVault"
+
+describe("TON Vault", () => {
+    test("Jettons are returned if sent to TON Vault", async () => {
+        const blockchain = await Blockchain.create()
+        const vaultSetup = await createTonVault(blockchain)
+
+        const _ = await vaultSetup.deploy()
+        const mockActionPayload = beginCell().storeStringTail("Random payload").endCell()
+
+        const jetton = await createJetton(blockchain)
+        const initialBalance = await jetton.wallet.getJettonBalance()
+        const numberOfJettons = BigInt(randomInt(0, 100000000000))
+        const sendResult = await jetton.transfer(
+            vaultSetup.vault.address,
+            numberOfJettons,
+            mockActionPayload,
+        )
+
+        const toVaultTx = flattenTransaction(
+            findTransactionRequired(sendResult.transactions, {
+                to: vaultSetup.vault.address,
+                op: TonVault.opcodes.UnexpectedJettonNotification,
+                success: true, // Because commit was called
+                exitCode:
+                    TonVault.errors[
+                        "TonVault: Jetton transfer must be performed to correct Jetton Vault"
+                    ],
+            }),
+        )
+
+        expect(sendResult.transactions).toHaveTransaction({
+            from: vaultSetup.vault.address,
+            to: toVaultTx.from,
+            op: TonVault.opcodes.ReturnJettonsViaJettonTransfer,
+            success: true,
+        })
+        const finalJettonBalance = await jetton.wallet.getJettonBalance()
+        expect(finalJettonBalance).toEqual(initialBalance)
+    })
+})

sources/utils/environment.ts

@@ -75,6 +75,8 @@ export type VaultInterface<T> = {
     }
     treasury: T
     deploy: () => Promise<SandboxSendResult>
+    // TON Vault is always inited, no need to init explicitly
+    isInited: () => Promise<boolean>
     addLiquidity: (
         liquidityDepositContractAddress: Address,
         amount: bigint,
@@ -111,6 +113,10 @@ export const createJettonVault: Create<VaultInterface<JettonTreasury>> = async (
         )
     }
 
+    const isInited = async () => {
+        return await vault.getInited()
+    }
+
     const addLiquidity = async (
         liquidityDepositContractAddress: Address,
         amount: bigint,
@@ -162,6 +168,7 @@ export const createJettonVault: Create<VaultInterface<JettonTreasury>> = async (
     return {
         vault,
         treasury: jetton,
+        isInited,
         deploy,
         addLiquidity,
         sendSwapRequest,
@@ -242,6 +249,9 @@ export const createTonVault: Create<VaultInterface<TonTreasury>> = async (
     return {
         deploy,
         vault,
+        isInited: async () => {
+            return true
+        },
         treasury: wallet,
         addLiquidity,
         sendSwapRequest,