MachinePay confidential x402 protocol spec: How BITE Phase 2 enables confidential x402 transfers

This document provides a detailed specification and a Solidity implementation for MachinePay encrypted x402 token standard to be used with x402 protocol. The spec incorporates ERC-3009 (Transfer With Authorization) capabilities.

This design leverages the synergy between BITE Phase 1 (input encryption) and BITE Phase 2 (contract-triggered decryption) to achieve comprehensive confidentiality.

1. Overview

The x402 token implementation aims to provide confidentiality for both user balances (state) and transaction details (inputs).

  • Transaction Confidentiality (BITE Phase 1): Users submit transactions using the BITE Phase 1 format. The entire transaction payload (e.g., the transfer function call including the amount and recipient) is encrypted by the client. The transaction is addressed to BITE_MAGIC_NUMBER. The BITE protocol decrypts the payload securely within the EVM execution environment. This ensures that inputs (like the transfer amount) are hidden on the public ledger.
  • State Confidentiality (BITE Phase 2): Balances are stored encrypted on-chain. BITE Phase 2 Contract-Triggered Transactions (CTXs) are used to securely decrypt these balances, perform the transfer logic using the plaintext inputs provided by Phase 1, and re-encrypt the results.

2. Architecture: Dual Encryption Strategy

To allow the contract to manage the ledger while enabling users to view their own balances, a dual encryption strategy is employed:

  1. Threshold Encryption (T_Key): Stored balances are encrypted using the BITE network threshold key. The smart contract can decrypt these via BITE Phase 2 CTXs.
  2. User Encryption (U_Key): Balances are simultaneously encrypted using the individual user’s public key (e.g., ECIES). Users can decrypt this locally.

3. BITE Infrastructure and Libraries

This design relies on the BITE infrastructure:

  1. BITE Phase 1 Protocol: Handles the encryption, transport (via BITE_MAGIC_NUMBER), and decryption of CALLDATA.
  2. BITE Phase 2 Precompile: Provides the decryptAndExecute mechanism.
  3. BITE Encryption Libraries: Standard library contracts provided by BITE (not precompiles) are used for re-encrypting data when state changes:
    • IThresholdEncryptor: Encrypts data using the BITE T_Key.
    • IUserEncryptor: Encrypts data using a specified U_Key.

4. Confidential Transfer Flow

Transfers are asynchronous and fully confidential.

Preparation (Off-chain / Client-side)

  1. The user prepares a standard transaction: transfer(recipient, amount).
  2. The client encrypts this payload using the BITE Phase 1 protocol (RLP encoding, AES encryption, addressing to BITE_MAGIC_NUMBER).

Block N: Initiation (transfer or transferWithAuthorization)

  1. The BITE protocol decrypts the Phase 1 payload.
  2. The EVM executes the intended function (e.g., transfer). The recipient and amount are visible to the contract logic in plaintext.
  3. Verification: The contract verifies locks and key registrations. For ERC-3009, the EIP-712 signature over the plaintext amount is verified (as the user signed plaintext, and the contract receives plaintext).
  4. CTX Scheduling: The contract calls BITE.decryptAndExecute. Since the amount is already known (via P1), it is passed efficiently as context.
    • encryptedArguments: [Sender T_Key Balance, Receiver T_Key Balance].
    • plaintextArguments: [Context ID, Sender Address, Receiver Address, Amount].

Block N+1: Execution (onDecrypt)

  1. The CTX executes onDecrypt.
  2. Validation: The contract validates the Context ID and the CTX origin (msg.sender).
  3. Data Retrieval: The contract receives the decrypted balances (decryptedArguments) and retrieves the Amount from plaintextArguments.
  4. Logic: The contract checks Sender Balance >= Amount.
  5. Re-encryption: New balances are calculated and re-encrypted using the IThresholdEncryptor and IUserEncryptor BITE library contracts.
  6. State Update: Storage is updated

MachinePay Encrypted x402 Token β€” Architecture Diagram

(Using BITE Phase 1 + Phase 2)

                   β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
                   β”‚            User / Client                β”‚
                   β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
                                  |
                                  | 1. User prepares:
                                  |    transfer(recipient, amount)
                                  |
                                  v
                   β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
                   β”‚     BITE Phase 1 Encryption (Client)   β”‚
                   β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€
                   β”‚ - RLP encode payload                   β”‚
                   β”‚ - Encrypt with BITE P1 (AES)           β”‚
                   β”‚ - Route to BITE_MAGIC_NUMBER           β”‚
                   β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
                                  |
                                  | Encrypted CALLDATA
                                  v
    β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
    β”‚                           BLOCK N                                  β”‚
    β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
                                  |
                                  v
                    β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
                    β”‚      BITE Protocol (Execution)          β”‚
                    β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€
                    β”‚ Phase 1: Decrypt CALLDATA inside EVM   β”‚
                    β”‚ β†’ contract sees plaintext (addr, amt)  β”‚
                    β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
                                  |
                                  v
              β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
              β”‚   MachinePay x402 Token Smart Contract (P1)      β”‚
              β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€
              β”‚ - Reads plaintext amount & recipient             β”‚
              β”‚ - Verifies authorization (ERC-3009)             β”‚
              β”‚ - Schedules Phase 2 CTX: decryptAndExecute       β”‚
              β”‚   encryptedArguments: [encBal_S, encBal_R]       β”‚
              β”‚   plaintextArguments: [ctxId, S, R, amount]      β”‚
              β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
                                  |
                                  | CTX scheduled β†’ next block
                                  v
    β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
    β”‚                           BLOCK N+1                                β”‚
    β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
                                  |
                                  v
              β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
              β”‚   BITE Phase 2 CTX (decryptAndExecute)           β”‚
              β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€
              β”‚ - Decrypts balances with T_Key                   β”‚
              β”‚ - Passes plaintextArguments to contract          β”‚
              β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
                                  |
                                  v
     β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
     β”‚     MachinePay x402 Token Smart Contract (onDecrypt handler)   β”‚
     β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€
     β”‚ 1. Validate CTX origin + Context ID                             β”‚
     β”‚ 2. Receive: decrypted S_Balance, R_Balance                      β”‚
     β”‚ 3. Logic: check S_Balance β‰₯ amount                              β”‚
     β”‚ 4. Compute new balances                                         β”‚
     β”‚ 5. Re-encrypt balances using:                                   β”‚
     β”‚      - IThresholdEncryptor (T_Key encrypted)                    β”‚
     β”‚      - IUserEncryptor (U_Key encrypted per user)                β”‚
     β”‚ 6. Update state:                                                β”‚
     β”‚      storage[S] = (encT[S], encU[S])                            β”‚
     β”‚      storage[R] = (encT[R], encU[R])                            β”‚
     β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
                                  |
                                  v
                     β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
                     β”‚        Final Confidential State          β”‚
                     β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€
                     β”‚ - Inputs hidden via Phase 1             β”‚
                     β”‚ - Balances hidden via Phase 2            β”‚
                     β”‚ - Users decrypt U_Key data locally      β”‚
                     β”‚ - Contract decrypts T_Key via CTX       β”‚
                     β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜

Solidity Implementation

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.24;

// Imports for EIP-3009 (EIP-712 Signatures) and ERC20 interface
// Assuming OpenZeppelin contracts are available in the environment.
import {EIP712} from "@openzeppelin/contracts/utils/cryptography/EIP712.sol";
import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";

// -------------------- BITE Interfaces and Libraries --------------------

/**
 * @notice Interface for BITE Phase 2 Precompile (decryptAndExecute).
 */
interface IBitePhase2 {
    function decryptAndExecute(
        bytes[] calldata encryptedArguments,
        bytes[] calldata plaintextArguments
    ) external;
}

/**
 * @notice BITE Library Contract Interface: Encrypts data using the BITE network threshold key (T_Key).
 */
interface IThresholdEncryptor {
    function encrypt(bytes calldata data) external view returns (bytes memory);
}

/**
 * @notice BITE Library Contract Interface: Encrypts data using a specified public key (U_Key, ECIES).
 */
interface IUserEncryptor {
    function encryptWithKey(bytes calldata publicKey, bytes calldata data) external view returns (bytes memory);
}

// -------------------- Contract Implementation --------------------

/**
 * @title EncryptedToken_x402
 * @notice Confidential ERC-20 token with ERC-3009 capabilities using BITE Phase 1 and Phase 2.
 */
contract EncryptedToken_x402 is EIP712, IERC20 {
    // -------------------- Configuration & Constants --------------------
    
    // BITE Phase 2 Precompile Address
    address constant BITE_PHASE2_ADDR = 0x0000000000000000000000000000000000000100;
    IBitePhase2 constant BITE = IBitePhase2(BITE_PHASE2_ADDR);

    // References to the BITE Encryption Library Contracts (Set during deployment via constructor)
    IThresholdEncryptor public immutable ThresholdEncryptor;
    IUserEncryptor public immutable UserEncryptor;

    string public constant name = "Encrypted Token x402";
    string public constant symbol = "x402";
    uint8 public constant decimals = 18;

    // EIP-3009 Typehash. Standard format is used because arguments are plaintext 
    // when signed by the user and when processed by the contract (thanks to BITE P1).
    bytes32 public constant TRANSFER_WITH_AUTHORIZATION_TYPEHASH = keccak256(
        "TransferWithAuthorization(address from,address to,uint256 value,uint256 validAfter,uint256 validBefore,bytes32 nonce)"
    );

    // -------------------- Storage --------------------

    // 1. Encrypted with BITE Threshold Key (T_Key) - Used for contract logic
    mapping(address => bytes) private thresholdBalances;
    
    // 2. Encrypted with User's Public Key (U_Key) - Used for local viewing
    mapping(address => bytes) private userBalances;

    // User Public Key Registry
    mapping(address => bytes) public publicKeys;

    // CTX management
    struct CTXState {
        bool active;
        address expectedCaller; // msg.sender that initiated the CTX (User or Relayer)
    }
    // Mapping from a unique context ID to the CTX state
    mapping(bytes32 => CTXState) private pendingCTXs;

    // ERC-3009 Authorization Nonces
    mapping(address => mapping(bytes32 => bool)) public authorizationUsed;

    // -------------------- Events --------------------
    // The standard ERC20 Transfer event MUST NOT be emitted, as it would leak the amount.
    // event Transfer(address indexed from, address indexed to, uint256 value);
    
    // Confidential Events (No amounts included)
    event PublicKeyRegistered(address indexed user, bytes publicKey);
    // Emitted in Block N
    event TransferInitiated(bytes32 indexed contextId, address indexed from, address indexed to);
    // Emitted in Block N+1
    event TransferExecuted(bytes32 indexed contextId, address indexed from, address indexed to);
    event TransferFailed(bytes32 indexed contextId, string reason);


    // -------------------- Initialization --------------------

    /**
     * @param _thresholdEncryptorAddr Address of the deployed BITE Threshold Encryptor library contract.
     * @param _userEncryptorAddr Address of the deployed BITE User Encryptor library contract.
     */
    constructor(address _thresholdEncryptorAddr, address _userEncryptorAddr) EIP712("Encrypted Token x402", "1") {
        ThresholdEncryptor = IThresholdEncryptor(_thresholdEncryptorAddr);
        UserEncryptor = IUserEncryptor(_userEncryptorAddr);
        // Note: Minting/Burning mechanics require specialized CTX flows and are omitted here.
    }

    // -------------------- User Key Management --------------------

    /**
     * @notice Register the user's public key for ECIES encryption. Required before interacting with the token.
     */
    function registerPublicKey(bytes memory publicKey) external {
        require(publicKey.length > 0, "Empty public key");
        publicKeys[msg.sender] = publicKey;
        emit PublicKeyRegistered(msg.sender, publicKey);
    }

    // -------------------- Public Views (ERC20 Compatibility) --------------------

    // Standard ERC20 functions return 0 or revert for confidentiality compatibility.
    function totalSupply() external pure override returns (uint256) { return 0; }
    function balanceOf(address /*account*/) public pure override returns (uint256) { return 0; }
    function allowance(address /*owner*/, address /*spender*/) external pure override returns (uint256) { return 0; }

    /**
     * @notice Returns the balance encrypted with the user's public key (U_Key) for local decryption.
     */
    function getUserEncryptedBalance(address account) external view returns (bytes memory) {
        return userBalances[account];
    }

    // -------------------- Transfer Initiation (Block N) --------------------
    // These functions MUST be called via BITE Phase 1 transactions to ensure confidentiality.

    /**
     * @notice Initiates a confidential transfer. (Standard ERC20)
     * @dev BITE P1 ensures 'to' and 'amount' are confidential during transport.
     * @param to The recipient address.
     * @param amount The amount to transfer (plaintext within the BITE execution context).
     */
    function transfer(address to, uint256 amount) 
        external override
        returns (bool) 
    {
        // msg.sender is the 'from' address and the 'initiator'
        _scheduleTransferCTX(msg.sender, to, amount, msg.sender);
        return true; // Optimistic return, execution is deferred to Block N+1.
    }

    /**
     * @notice EIP-3009 transfer.
     * @dev BITE P1 ensures inputs are confidential during transport.
     */
    function transferWithAuthorization(
        address from,
        address to,
        uint256 amount, // Plaintext (decrypted by P1)
        uint256 validAfter,
        uint256 validBefore,
        bytes32 nonce,
        uint8 v,
        bytes32 r,
        bytes32 s
    ) external {
        require(block.timestamp >= validAfter, "Authorization not yet valid");
        require(block.timestamp <= validBefore, "Authorization expired");
        require(!authorizationUsed[from][nonce], "Authorization already used");

        // Verify EIP-712 Signature over plaintext arguments
        bytes32 digest = _hashTypedDataV4(
            keccak256(
                abi.encode(
                    TRANSFER_WITH_AUTHORIZATION_TYPEHASH,
                    from,
                    to,
                    amount, // Corresponds to 'value' in the typehash definition
                    validAfter,
                    validBefore,
                    nonce
                )
            )
        );

        address signer = ECDSA.recover(digest, v, r, s);
        require(signer == from, "Invalid signature");

        authorizationUsed[from][nonce] = true;

        // The initiator (expected msg.sender of the CTX) will be the relayer (msg.sender here)
        _scheduleTransferCTX(from, to, amount, msg.sender);
    }

    // Note: Implementing encrypted approve/transferFrom requires complex multi-CTX flows and is out of scope for this MVP.
    function approve(address /*spender*/, uint256 /*amount*/) external override returns (bool) { revert("Encrypted approve not supported"); }
    function transferFrom(address /*from*/, address /*to*/, uint256 /*amount*/) external override returns (bool) { revert("Encrypted transferFrom not supported"); }

    function _scheduleTransferCTX(address from, address to, uint256 amount, address initiator) private {
        require(to != address(0), "Transfer to zero address");
        require(amount > 0, "Amount must be greater than zero");

        // Check public key registration (required for re-encryption in Block N+1)
        require(publicKeys[from].length > 0, "Sender key not registered");
        require(publicKeys[to].length > 0, "Recipient key not registered");



        // Prepare arguments for CTX decryption.
        // We only need to decrypt the balances. The amount is already known via P1 decryption.
        bytes[] memory encArgs = new bytes[](2);
        // Arg 0: Sender's T_Key encrypted balance
        encArgs[0] = thresholdBalances[from];
        // Arg 1: Receiver's T_Key encrypted balance
        encArgs[1] = thresholdBalances[to];

        // Generate unique context ID (ensures uniqueness of the CTX request)
        bytes32 contextId = keccak256(abi.encodePacked(block.number, from, to, amount, initiator));

        // Prepare plaintext arguments (Context and Amount)
        bytes[] memory plainArgs = new bytes[](4);
        plainArgs[0] = abi.encode(from);
        plainArgs[1] = abi.encode(to);
        // Arg 2: Amount is passed in plaintext to the CTX as context
        plainArgs[2] = abi.encode(amount); 
        plainArgs[3] = abi.encode(contextId);

        // 5. Manage CTX state
        pendingCTXs[contextId] = CTXState({
            active: true,
            // Per BITE P2 spec, the CTX msg.sender in Block N+1 will match the initiator in Block N
            expectedCaller: initiator 
        });

        emit TransferInitiated(contextId, from, to);

        // 6. Schedule CTX for Block N+1
        BITE.decryptAndExecute(encArgs, plainArgs);
    }

    // -------------------- CTX Callback (Block N+1) --------------------

    /**
     * @notice Callback executed by the CTX in Block N+1.
     * @param decryptedArguments [decryptedSenderBalance, decryptedReceiverBalance]
     * @param plaintextArguments [senderAddress, receiverAddress, amount, contextId]
     */
    function onDecrypt(
        bytes[] calldata decryptedArguments,
        bytes[] calldata plaintextArguments
    ) external {
        // 1. Decode Context and Amount
        require(plaintextArguments.length == 4, "Bad plaintext len");
        (address from) = abi.decode(plaintextArguments[0], (address));
        (address to) = abi.decode(plaintextArguments[1], (address));
        (uint256 amount) = abi.decode(plaintextArguments[2], (uint256));
        (bytes32 contextId) = abi.decode(plaintextArguments[3], (bytes32));

        // 2. Validate CTX State
        CTXState storage pending = pendingCTXs[contextId];
        require(pending.active, "No pending CTX for this context");
        // Security Check: Ensure the caller matches the expected initiator (BITE P2 security model).
        require(msg.sender == pending.expectedCaller, "Unexpected caller (not CTX origin)");


        // 3. Decode Decrypted Balances
        require(decryptedArguments.length == 2, "Bad decrypted len");
        
        // If the encrypted input was empty (balance implicitly 0), the decrypted bytes will be empty.
        uint256 balanceFrom = _decodeOrZero(decryptedArguments[0]);
        uint256 balanceTo = _decodeOrZero(decryptedArguments[1]);

        // 4. Check Sufficiency
        if (balanceFrom < amount) {
            _cleanup(from, to, contextId);
            emit TransferFailed(contextId, "Insufficient funds");
            return;
        }

        // 5. Calculate New Balances
        // Arithmetic is safe due to the check above and using Solidity 0.8+.
        uint256 newBalanceFrom = balanceFrom - amount;
        uint256 newBalanceTo = balanceTo + amount;

        // 6. Re-encryption and Storage Update (using BITE Libraries)
        // Fetch keys (already checked in Block N)
        bytes memory pkFrom = publicKeys[from];
        bytes memory pkTo = publicKeys[to];

        _updateEncryptedBalances(from, newBalanceFrom, pkFrom);
        _updateEncryptedBalances(to, newBalanceTo, pkTo);

        // 7. Cleanup
        _cleanup(from, to, contextId);
        emit TransferExecuted(contextId, from, to);
    }

    // -------------------- Helpers --------------------

    function _cleanup(address from, address to, bytes32 contextId) internal {
        // Clear storage
        delete pendingCTXs[contextId];
    }

    /**
     * @notice Helper to decode bytes to uint256, returning 0 if bytes are empty.
     */
    function _decodeOrZero(bytes memory data) internal pure returns (uint256) {
        if (data.length == 0) {
            return 0;
        }
        // Assumes the decrypted data is an abi-encoded uint256
        return abi.decode(data, (uint256));
    }

    /**
     * @notice Helper to perform dual encryption using BITE library contracts and update storage.
     */
    function _updateEncryptedBalances(address account, uint256 balance, bytes memory publicKey) private {
        bytes memory balanceBytes = abi.encode(balance);
        
        // 1. Threshold Encryption (T_Key) - using BITE library contract
        thresholdBalances[account] = ThresholdEncryptor.encrypt(balanceBytes);

        // 2. User Key Encryption (U_Key) - using BITE library contract
        userBalances[account] = UserEncryptor.encryptWithKey(publicKey, balanceBytes);
    }
}
2 Likes