Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 31 additions & 3 deletions src/paymasters/BondTreasuryPaymaster.sol
Original file line number Diff line number Diff line change
Expand Up @@ -9,17 +9,24 @@ import {WhitelistPaymaster} from "./WhitelistPaymaster.sol";

/// @notice ZkSync paymaster plus ERC-20 bond treasury for whitelisted contracts (e.g. `IBondTreasury` implementers).
/// @dev Extends `WhitelistPaymaster` (user + destination whitelist, gas sponsorship). Adds bond token balance,
/// `consumeSponsoredBond` with `QuotaControl`, and ERC-20 withdrawal. Constructor seeds `address(this)` as a
/// whitelisted destination so management txs can be sponsored once those EOAs are user-whitelisted.
/// `consumeSponsoredBond` with `QuotaControl`, per-user token allowance, and ERC-20 withdrawal.
/// Constructor seeds `address(this)` as a whitelisted destination so management txs can be sponsored
/// once those EOAs are user-whitelisted.
contract BondTreasuryPaymaster is WhitelistPaymaster, QuotaControl {
using SafeERC20 for IERC20;

IERC20 public immutable bondToken;

/// @notice Per-user remaining token allowance for sponsored bonds.
mapping(address => uint256) public userBondAllowance;

event TokensWithdrawn(address indexed token, address indexed to, uint256 amount);
event UserBondAllowanceSet(address indexed user, uint256 newAllowance);
event UserBondAllowanceIncreased(address indexed user, uint256 addedAmount, uint256 newAllowance);

error CallerNotWhitelistedContract();
error InsufficientBondBalance();
error UserBondAllowanceExceeded();

constructor(
address admin,
Expand Down Expand Up @@ -60,19 +67,40 @@ contract BondTreasuryPaymaster is WhitelistPaymaster, QuotaControl {
}
}

/// @notice Validate whitelist + consume quota for a sponsored bond.
/// @notice Validate whitelist + consume quota + decrement per-user allowance for a sponsored bond.
/// @dev Callable only by `isWhitelistedContract`. Caller pulls via `transferFrom`.
function consumeSponsoredBond(address user, uint256 amount) external {
if (!isWhitelistedContract[msg.sender]) revert CallerNotWhitelistedContract();
if (!isWhitelistedUser[user]) revert UserIsNotWhitelisted();
if (bondToken.balanceOf(address(this)) < amount) revert InsufficientBondBalance();

uint256 remaining = userBondAllowance[user];
if (remaining < amount) revert UserBondAllowanceExceeded();
unchecked {
userBondAllowance[user] = remaining - amount;
}

_checkedResetClaimed();
_checkedUpdateClaimed(amount);

bondToken.forceApprove(msg.sender, amount);
}

/// @notice Set the total bond allowance for a user (overwrites previous value).
function setUserBondAllowance(address user, uint256 allowance) external {
_checkRole(WHITELIST_ADMIN_ROLE);
userBondAllowance[user] = allowance;
emit UserBondAllowanceSet(user, allowance);
}

/// @notice Add to a user's existing bond allowance.
function increaseUserBondAllowance(address user, uint256 amount) external {
_checkRole(WHITELIST_ADMIN_ROLE);
uint256 newAllowance = userBondAllowance[user] + amount;
userBondAllowance[user] = newAllowance;
emit UserBondAllowanceIncreased(user, amount, newAllowance);
}

/// @notice Withdraw ERC-20 tokens (e.g. excess bond token) from this contract.
function withdrawTokens(address token, address to, uint256 amount) external {
_checkRole(WITHDRAWER_ROLE);
Expand Down
44 changes: 32 additions & 12 deletions src/swarms/FleetIdentityUpgradeable.sol
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,9 @@

pragma solidity ^0.8.24;

import {ERC721EnumerableUpgradeable} from
"@openzeppelin/contracts-upgradeable/token/ERC721/extensions/ERC721EnumerableUpgradeable.sol";
import {
ERC721EnumerableUpgradeable
} from "@openzeppelin/contracts-upgradeable/token/ERC721/extensions/ERC721EnumerableUpgradeable.sol";
import {Ownable2StepUpgradeable} from "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
Expand Down Expand Up @@ -198,14 +199,24 @@ contract FleetIdentityUpgradeable is
mapping(uint16 => uint32[]) internal _countryAdminAreas;
mapping(uint32 => uint256) internal _countryAdminAreaIndex;

// ──────────────────────────────────────────────
// V2 Storage: Bond payer tracking
// ──────────────────────────────────────────────

/// @notice UUID -> address that paid the ownership bond.
/// @dev For self-funded claims this is msg.sender; for sponsored claims
/// this is the treasury. On burn the bond refunds to this address.
/// Zero (pre-upgrade tokens) falls back to uuidOwner.
mapping(bytes16 => address) public uuidOwnershipBondPayer;

// ──────────────────────────────────────────────
// Storage Gap (for future upgrades)
// ──────────────────────────────────────────────

/// @dev Reserved storage slots for future upgrades.
/// When adding new storage in V2+, reduce this gap accordingly.
// solhint-disable-next-line var-name-mixedcase
uint256[50] private __gap;
uint256[49] private __gap;

// ──────────────────────────────────────────────
// Events
Expand Down Expand Up @@ -433,10 +444,14 @@ contract FleetIdentityUpgradeable is

// Use snapshot for accurate refund
uint256 ownershipBond = uuidOwnershipBondPaid[uuid];
// Refund to whoever paid the bond (treasury for sponsored, owner for self-funded).
// Zero (pre-upgrade tokens) falls back to the UUID owner.
address bondPayer = uuidOwnershipBondPayer[uuid];
if (bondPayer == address(0)) bondPayer = owner;

_burn(tokenId);
_clearUuidOwnership(uuid);
_refundBond(owner, ownershipBond);
_refundBond(bondPayer, ownershipBond);

emit FleetBurned(tokenHolder, tokenId, region, 0, ownershipBond);
} else {
Expand Down Expand Up @@ -486,6 +501,7 @@ contract FleetIdentityUpgradeable is
_mint(msg.sender, tokenId);

uuidOwnershipBondPaid[uuid] = _baseBond;
uuidOwnershipBondPayer[uuid] = msg.sender;
_pullBond(msg.sender, _baseBond);

emit UuidClaimed(msg.sender, uuid, operatorOf(uuid));
Expand Down Expand Up @@ -525,6 +541,7 @@ contract FleetIdentityUpgradeable is

tokenId = uint256(uint128(uuid));
uuidOwnershipBondPaid[uuid] = _baseBond;
uuidOwnershipBondPayer[uuid] = treasury;
_mint(msg.sender, tokenId);

// Bond transfer uses the trusted _bondToken — cannot be faked by the treasury.
Expand Down Expand Up @@ -660,11 +677,7 @@ contract FleetIdentityUpgradeable is
}

/// @notice Builds a bundle containing ONLY country-level fleets.
function buildCountryOnlyBundle(uint16 countryCode)
external
view
returns (bytes16[] memory uuids, uint256 count)
{
function buildCountryOnlyBundle(uint16 countryCode) external view returns (bytes16[] memory uuids, uint256 count) {
if (countryCode == 0 || countryCode > MAX_COUNTRY_CODE) revert InvalidCountryCode();

uint32 countryKey = uint32(countryCode);
Expand Down Expand Up @@ -761,6 +774,7 @@ contract FleetIdentityUpgradeable is
delete uuidOperator[uuid];
delete uuidTotalTierBonds[uuid];
delete uuidOwnershipBondPaid[uuid];
delete uuidOwnershipBondPayer[uuid];
}

function _decrementUuidCount(bytes16 uuid) internal returns (uint256 newCount) {
Expand Down Expand Up @@ -832,6 +846,7 @@ contract FleetIdentityUpgradeable is
tokenId = _mintFleetToken(uuid, region, targetTier);
tokenTier0Bond[tokenId] = tier0Bond;
uuidOwnershipBondPaid[uuid] = _baseBond;
uuidOwnershipBondPayer[uuid] = msg.sender;

_pullBond(msg.sender, _baseBond + targetTierBond);

Expand Down Expand Up @@ -975,8 +990,7 @@ contract FleetIdentityUpgradeable is
uint32 adminKey = makeAdminRegion(countryCode, adminCode);
uint32 candidateRegion = isCountry ? countryKey : adminKey;

(, uint256 count, uint256 highestTier, uint256 lowestTier) =
_buildHighestBondedUuidBundle(countryKey, adminKey);
(, uint256 count, uint256 highestTier, uint256 lowestTier) = _buildHighestBondedUuidBundle(countryKey, adminKey);

for (uint256 tier = lowestTier; tier <= highestTier; ++tier) {
bool tierHasCapacity = _regionTierMembers[candidateRegion][tier].length < TIER_CAPACITY;
Expand Down Expand Up @@ -1114,7 +1128,13 @@ contract FleetIdentityUpgradeable is

uint32 region = tokenRegion(tokenId);
if (region == OWNED_REGION_KEY && from != address(0) && to != address(0)) {
uuidOwner[tokenUuid(tokenId)] = to;
bytes16 uuid = tokenUuid(tokenId);
uuidOwner[uuid] = to;
// Transfer bond-refund right only if it was the previous owner (self-funded).
// Sponsored bonds (treasury != from) stay with the original payer.
if (uuidOwnershipBondPayer[uuid] == from) {
uuidOwnershipBondPayer[uuid] = to;
}
}

return from;
Expand Down
Loading
Loading