Skip to content

Commit da1f654

Browse files
gator-boigator-boi
committed
evm: add variable width id encoding
1 parent e137dbc commit da1f654

File tree

4 files changed

+194
-28
lines changed

4 files changed

+194
-28
lines changed

‎evm/src/NativeTransfers/NonFungibleNttManager.sol

+25-2
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,25 @@ contract NonFungibleNttManager is INonFungibleNttManager, ManagerBase {
2121

2222
// =============== Immutables ============================================================
2323

24+
// Hard cap on the number of NFTs that can be transferred in a single batch. This is to prevent
25+
// the contract from running out of gas when processing large batches of NFTs.
2426
uint8 constant MAX_BATCH_SIZE = 50;
2527

28+
// The number of bytes each NFT token ID occupies in the payload. All tokenIDs must fit within
29+
// this width.
30+
uint8 immutable tokenIdWidth;
31+
2632
// =============== Setup =================================================================
2733

28-
constructor(address _token, Mode _mode, uint16 _chainId) ManagerBase(_token, _mode, _chainId) {}
34+
constructor(
35+
address _token,
36+
uint8 _tokenIdWidth,
37+
Mode _mode,
38+
uint16 _chainId
39+
) ManagerBase(_token, _mode, _chainId) {
40+
_validateTokenIdWidth(_tokenIdWidth);
41+
tokenIdWidth = _tokenIdWidth;
42+
}
2943

3044
function __NonFungibleNttManager_init() internal onlyInitializing {
3145
// check if the owner is the deployer of this contract
@@ -211,7 +225,7 @@ contract NonFungibleNttManager is INonFungibleNttManager, ManagerBase {
211225
TransceiverStructs.ManagerMessage(
212226
bytes32(uint256(sequence)),
213227
toWormholeFormat(msg.sender),
214-
TransceiverStructs.encodeNonFungibleNativeTokenTransfer(nft)
228+
TransceiverStructs.encodeNonFungibleNativeTokenTransfer(nft, tokenIdWidth)
215229
)
216230
);
217231

@@ -272,4 +286,13 @@ contract NonFungibleNttManager is INonFungibleNttManager, ManagerBase {
272286
revert InvalidPeer(sourceChainId, peerAddress);
273287
}
274288
}
289+
290+
function _validateTokenIdWidth(uint8 _tokenIdWidth) internal pure {
291+
if (
292+
_tokenIdWidth != 1 && _tokenIdWidth != 2 && _tokenIdWidth != 4 && _tokenIdWidth != 8
293+
&& _tokenIdWidth != 16 && _tokenIdWidth != 32
294+
) {
295+
revert InvalidTokenIdWidth(_tokenIdWidth);
296+
}
297+
}
275298
}

‎evm/src/interfaces/INonFungibleNttManager.sol

+2
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,8 @@ interface INonFungibleNttManager is IManagerBase {
6767
/// @param thisChain The current chain.
6868
error InvalidTargetChain(uint16 targetChain, uint16 thisChain);
6969

70+
error InvalidTokenIdWidth(uint8 tokenIdWidth);
71+
7072
/// @notice Sets the corresponding peer.
7173
/// @dev The NonFungiblenttManager that executes the message sets the source NonFungibleNttManager
7274
/// as the peer.

‎evm/src/libraries/TransceiverStructs.sol

+93-13
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ library TransceiverStructs {
1919
/// @param prefix The prefix that was found in the encoded message.
2020
error IncorrectPrefix(bytes4 prefix);
2121
error UnorderedInstructions();
22+
error TokenIdTooLarge(uint256 tokenId, uint256 max);
2223

2324
/// @dev Prefix for all NativeTokenTransfer payloads
2425
/// This is 0x99'N''T''T'
@@ -147,8 +148,18 @@ library TransceiverStructs {
147148
encoded.checkLength(offset);
148149
}
149150

150-
/// @dev Native Token Transfer payload.
151-
/// TODO: Document wire format.
151+
/// @dev Non-Fungible Native Token Transfer payload.
152+
/// The wire format is as follows:
153+
/// - NON_FUNGIBLE_NTT_PREFIX - 4 bytes
154+
/// - to - 32 bytes
155+
/// - toChain - 2 bytes
156+
/// - batchSize - 2 bytes
157+
/// - for each encoded tokenId:
158+
/// - tokenIdWidth - 1 byte
159+
/// - tokenId - `tokenIdWidth` bytes
160+
/// - payloadLength - 2 bytes
161+
/// - payload - `payloadLength` bytes
162+
152163
struct NonFungibleNativeTokenTransfer {
153164
/// @notice Address of the recipient.
154165
bytes32 to;
@@ -160,19 +171,18 @@ library TransceiverStructs {
160171
bytes payload;
161172
}
162173

163-
function encodeNonFungibleNativeTokenTransfer(NonFungibleNativeTokenTransfer memory nft)
164-
public
165-
pure
166-
returns (bytes memory encoded)
167-
{
174+
function encodeNonFungibleNativeTokenTransfer(
175+
NonFungibleNativeTokenTransfer memory nft,
176+
uint8 tokenIdWidth
177+
) public pure returns (bytes memory encoded) {
168178
uint16 batchSize = uint16(nft.tokenIds.length);
169179
uint16 payloadLen = uint16(nft.payload.length);
170180

171181
bytes memory encodedTokenIds = abi.encodePacked(batchSize);
172182
for (uint256 i = 0; i < batchSize; ++i) {
173-
// For now encode each token ID as 32 bytes long.
174-
// TODO: Optimize this to encode only the necessary bytes.
175-
encodedTokenIds = abi.encodePacked(encodedTokenIds, uint8(32), nft.tokenIds[i]);
183+
encodedTokenIds = abi.encodePacked(
184+
encodedTokenIds, tokenIdWidth, encodeTokenId(nft.tokenIds[i], tokenIdWidth)
185+
);
176186
}
177187

178188
return abi.encodePacked(
@@ -200,9 +210,9 @@ library TransceiverStructs {
200210

201211
uint256[] memory tokenIds = new uint256[](batchSize);
202212
for (uint256 i = 0; i < batchSize; ++i) {
203-
uint8 tokenIdLength;
204-
(tokenIdLength, offset) = encoded.asUint8Unchecked(offset);
205-
(tokenIds[i], offset) = encoded.asUint256Unchecked(offset);
213+
uint256 tokenId;
214+
(offset, tokenId) = parseTokenId(encoded, offset);
215+
tokenIds[i] = tokenId;
206216
}
207217
nonFungibleNtt.tokenIds = tokenIds;
208218

@@ -214,6 +224,76 @@ library TransceiverStructs {
214224
encoded.checkLength(offset);
215225
}
216226

227+
function encodeTokenId(
228+
uint256 tokenId,
229+
uint8 tokenIdWidth
230+
) public pure returns (bytes memory) {
231+
if (tokenIdWidth == 1) {
232+
if (tokenId > type(uint8).max) {
233+
revert TokenIdTooLarge(tokenId, type(uint8).max);
234+
} else {
235+
return abi.encodePacked(uint8(tokenId));
236+
}
237+
} else if (tokenIdWidth == 2) {
238+
if (tokenId > type(uint16).max) {
239+
revert TokenIdTooLarge(tokenId, type(uint16).max);
240+
} else {
241+
return abi.encodePacked(uint16(tokenId));
242+
}
243+
} else if (tokenIdWidth == 4) {
244+
if (tokenId > type(uint32).max) {
245+
revert TokenIdTooLarge(tokenId, type(uint32).max);
246+
} else {
247+
return abi.encodePacked(uint32(tokenId));
248+
}
249+
} else if (tokenIdWidth == 8) {
250+
if (tokenId > type(uint64).max) {
251+
revert TokenIdTooLarge(tokenId, type(uint64).max);
252+
} else {
253+
return abi.encodePacked(uint64(tokenId));
254+
}
255+
} else if (tokenIdWidth == 16) {
256+
if (tokenId > type(uint128).max) {
257+
revert TokenIdTooLarge(tokenId, type(uint128).max);
258+
} else {
259+
return abi.encodePacked(uint128(tokenId));
260+
}
261+
} else {
262+
return abi.encodePacked(uint256(tokenId));
263+
}
264+
}
265+
266+
function parseTokenId(
267+
bytes memory encoded,
268+
uint256 offset
269+
) public pure returns (uint256 nextOffset, uint256 tokenId) {
270+
uint8 tokenIdWidth;
271+
(tokenIdWidth, nextOffset) = encoded.asUint8Unchecked(offset);
272+
if (tokenIdWidth == 1) {
273+
uint8 tokenIdValue;
274+
(tokenIdValue, nextOffset) = encoded.asUint8Unchecked(nextOffset);
275+
tokenId = tokenIdValue;
276+
} else if (tokenIdWidth == 2) {
277+
uint16 tokenIdValue;
278+
(tokenIdValue, nextOffset) = encoded.asUint16Unchecked(nextOffset);
279+
tokenId = tokenIdValue;
280+
} else if (tokenIdWidth == 4) {
281+
uint32 tokenIdValue;
282+
(tokenIdValue, nextOffset) = encoded.asUint32Unchecked(nextOffset);
283+
tokenId = tokenIdValue;
284+
} else if (tokenIdWidth == 8) {
285+
uint64 tokenIdValue;
286+
(tokenIdValue, nextOffset) = encoded.asUint64Unchecked(nextOffset);
287+
tokenId = tokenIdValue;
288+
} else if (tokenIdWidth == 16) {
289+
uint128 tokenIdValue;
290+
(tokenIdValue, nextOffset) = encoded.asUint128Unchecked(nextOffset);
291+
tokenId = tokenIdValue;
292+
} else {
293+
(tokenId, nextOffset) = encoded.asUint256Unchecked(nextOffset);
294+
}
295+
}
296+
217297
/// @dev Message emitted by Transceiver implementations.
218298
/// Each message includes an Transceiver-specified 4-byte prefix.
219299
/// The wire format is as follows:

‎evm/test/NonFungibleNttManager.t.sol

+74-13
Original file line numberDiff line numberDiff line change
@@ -24,11 +24,14 @@ import "./libraries/NttManagerHelpers.sol";
2424
import {Utils} from "./libraries/Utils.sol";
2525
import "openzeppelin-contracts/contracts/proxy/ERC1967/ERC1967Proxy.sol";
2626
import "../src/libraries/external/OwnableUpgradeable.sol";
27+
import "wormhole-solidity-sdk/libraries/BytesParsing.sol";
2728

2829
import "./mocks/MockTransceivers.sol";
2930
import "../src/mocks/DummyNft.sol";
3031

3132
contract TestNonFungibleNttManager is Test {
33+
using BytesParsing for bytes;
34+
3235
uint16 constant chainIdOne = 2;
3336
uint16 constant chainIdTwo = 6;
3437
uint16 constant chainIdThree = 10;
@@ -40,6 +43,7 @@ contract TestNonFungibleNttManager is Test {
4043
address relayer = 0x7B1bD7a6b4E61c2a123AC6BC2cbfC614437D0470;
4144
uint8 consistencyLevel = 1;
4245
uint256 baseGasLimit = 500000;
46+
uint8 tokenIdWidth = 2;
4347

4448
address owner = makeAddr("owner");
4549

@@ -60,7 +64,7 @@ contract TestNonFungibleNttManager is Test {
6064
bool shouldInitialize
6165
) internal returns (INonFungibleNttManager) {
6266
NonFungibleNttManager implementation =
63-
new NonFungibleNttManager(address(nft), _mode, _chainId);
67+
new NonFungibleNttManager(address(nft), tokenIdWidth, _mode, _chainId);
6468

6569
NonFungibleNttManager proxy =
6670
NonFungibleNttManager(address(new ERC1967Proxy(address(implementation), "")));
@@ -227,12 +231,60 @@ contract TestNonFungibleNttManager is Test {
227231
managerOne.setPeer(chainId, newPeer);
228232
}
229233

230-
// ============================ Transfer Tests ======================================
231-
232-
function test_lockAndMint(uint256 nftCount, uint256 startId) public {
234+
// ============================ Serde Tests ======================================
235+
236+
function test_serde(
237+
uint8 tokenIdWidth,
238+
bytes32 to,
239+
uint16 toChain,
240+
bytes memory payload,
241+
uint256 nftCount,
242+
uint256 startId
243+
) public {
244+
// Narrow the search.
245+
tokenIdWidth = uint8(bound(tokenIdWidth, 1, 32));
246+
// Ugly, but necessary.
247+
vm.assume(
248+
tokenIdWidth == 1 || tokenIdWidth == 2 || tokenIdWidth == 4 || tokenIdWidth == 8
249+
|| tokenIdWidth == 16 || tokenIdWidth == 32
250+
);
251+
vm.assume(to != bytes32(0));
252+
vm.assume(toChain != 0);
233253
nftCount = bound(nftCount, 1, managerOne.getMaxBatchSize());
234254
startId = bound(startId, 0, type(uint256).max - nftCount);
235255

256+
TransceiverStructs.NonFungibleNativeTokenTransfer memory nftTransfer = TransceiverStructs
257+
.NonFungibleNativeTokenTransfer({
258+
to: to,
259+
toChain: toChain,
260+
payload: payload,
261+
tokenIds: _createBatchTokenIds(nftCount, startId)
262+
});
263+
264+
bytes memory encoded =
265+
TransceiverStructs.encodeNonFungibleNativeTokenTransfer(nftTransfer, 32);
266+
267+
TransceiverStructs.NonFungibleNativeTokenTransfer memory out =
268+
TransceiverStructs.parseNonFungibleNativeTokenTransfer(encoded);
269+
270+
assertEq(out.to, to, "To address should be the same");
271+
assertEq(out.toChain, toChain, "To chain should be the same");
272+
assertEq(out.payload, payload, "Payload should be the same");
273+
assertEq(
274+
out.tokenIds.length, nftTransfer.tokenIds.length, "TokenIds length should be the same"
275+
);
276+
277+
for (uint256 i = 0; i < nftCount; i++) {
278+
assertEq(out.tokenIds[i], nftTransfer.tokenIds[i], "TokenId should be the same");
279+
}
280+
}
281+
282+
// ============================ Transfer Tests ======================================
283+
284+
function test_lockAndMint(uint16 nftCount, uint16 startId) public {
285+
nftCount = uint16(bound(nftCount, 1, managerOne.getMaxBatchSize()));
286+
startId = uint16(bound(startId, 0, type(uint16).max - nftCount));
287+
236288
address recipient = makeAddr("recipient");
237289
uint256[] memory tokenIds = _mintNftBatch(nftOne, recipient, nftCount, startId);
238290

@@ -253,9 +305,9 @@ contract TestNonFungibleNttManager is Test {
253305
assertTrue(_isBatchOwner(nftOne, tokenIds, address(managerOne)), "Manager should own NFTs");
254306
}
255307

256-
function test_burnAndUnlock(uint256 nftCount, uint256 startId) public {
257-
nftCount = bound(nftCount, 1, managerTwo.getMaxBatchSize());
258-
startId = bound(startId, 0, type(uint256).max - nftCount);
308+
function test_burnAndUnlock(uint16 nftCount, uint16 startId) public {
309+
nftCount = uint16(bound(nftCount, 1, managerOne.getMaxBatchSize()));
310+
startId = uint16(bound(startId, 0, type(uint16).max - nftCount));
259311

260312
// Mint nftOne to managerOne to "lock" them.
261313
{
@@ -287,9 +339,9 @@ contract TestNonFungibleNttManager is Test {
287339
assertTrue(_isBatchOwner(nftOne, tokenIds, recipient), "Recipient should own NFTs");
288340
}
289341

290-
function test_burnAndMint(uint256 nftCount, uint256 startId) public {
291-
nftCount = bound(nftCount, 1, managerOne.getMaxBatchSize());
292-
startId = bound(startId, 0, type(uint256).max - nftCount);
342+
function test_burnAndMint(uint16 nftCount, uint16 startId) public {
343+
nftCount = uint16(bound(nftCount, 1, managerOne.getMaxBatchSize()));
344+
startId = uint16(bound(startId, 0, type(uint16).max - nftCount));
293345

294346
address recipient = makeAddr("recipient");
295347
uint256[] memory tokenIds = _mintNftBatch(nftOne, recipient, nftCount, startId);
@@ -483,9 +535,7 @@ contract TestNonFungibleNttManager is Test {
483535

484536
vm.expectRevert(
485537
abi.encodeWithSelector(
486-
INonFungibleNttManager.InvalidTargetChain.selector,
487-
chainIdThree,
488-
chainIdTwo
538+
INonFungibleNttManager.InvalidTargetChain.selector, chainIdThree, chainIdTwo
489539
)
490540
);
491541
transceiverTwo.receiveMessage(encodedVm);
@@ -605,6 +655,17 @@ contract TestNonFungibleNttManager is Test {
605655
return arr;
606656
}
607657

658+
function _createBatchTokenIds(
659+
uint256 len,
660+
uint256 start
661+
) internal pure returns (uint256[] memory) {
662+
uint256[] memory arr = new uint256[](len);
663+
for (uint256 i = 0; i < len; i++) {
664+
arr[i] = start + i;
665+
}
666+
return arr;
667+
}
668+
608669
function _getWormholeMessage(
609670
Vm.Log[] memory logs,
610671
uint16 emitterChain

0 commit comments

Comments
 (0)