131 smart contract nft internal security audit

smart-contracts/assembly/contracts/NFT/NFT-internals.ts

@@ -3,6 +3,7 @@ import {
   bytesToU64,
   u64ToBytes,
   Args,
+  SafeMath,
 } from '@massalabs/as-types';
 import { Storage, Context, validateAddress } from '@massalabs/massa-as-sdk';
 
@@ -58,21 +59,28 @@ export function _constructor(args: Args): void {
  */
 export function _increment(): u64 {
   const currentID = bytesToU64(Storage.get(counterKey));
-  const newID = currentID + 1;
+  const newID = SafeMath.add(currentID, 1);
   Storage.set(counterKey, u64ToBytes(newID));
   return newID;
 }
 
 export function _updateBalanceOf(address: string, increment: boolean): void {
   const balanceKey = stringToBytes('balanceOf_' + address);
-  const number = increment ? 1 : -1;
+  const number = 1;
 
   if (Storage.has(balanceKey)) {
     const balance = bytesToU64(Storage.get(balanceKey));
-    const newBalance = balance + number;
+    let newBalance: u64;
+
+    if (increment) {
+      newBalance = SafeMath.add(balance, number);
+    } else {
+      newBalance = SafeMath.sub(balance, number);
+    }
+
     Storage.set(balanceKey, u64ToBytes(newBalance));
   } else {
-    assert(number == 1, 'Balance cannot be negative');
+    assert(increment, 'Balance cannot be negative');
     Storage.set(balanceKey, u64ToBytes(1));
   }
 }
@@ -97,8 +105,9 @@ export function _onlyOwner(): bool {
  * @returns true if the caller is token's owner
  */
 export function _isTokenOwner(address: string, tokenId: u64): bool {
-  // as we need to compare two byteArrays, we need to compare the pointers
-  // we transform our byte array to their pointers and we compare them
+  // To compare two byte arrays, we compare their contents.
+  // The byte arrays are transformed into pointers, and `memory.compare` is used to
+  // compare the values (bytes) referenced by these pointers.
   const left = nft1_ownerOf(u64ToBytes(tokenId));
   return (
     memory.compare(

smart-contracts/assembly/contracts/NFT/__tests__/NFT.spec.ts

@@ -31,6 +31,7 @@ import {
   nft1_setApprovalForAll,
   nft1_isApprovedForAll,
 } from '../NFT';
+import { _increment, _updateBalanceOf } from '../NFT-internals';
 
 const callerAddress = 'A12UBnqTHDQALpocVBnkPNy7y5CndUJQTLutaVDDFgMJcq5kQiKq';
 
@@ -209,3 +210,60 @@ function approveAddress(tokenId: u64, address: string): void {
   const args = new Args().add(tokenId).add(address).serialize();
   nft1_approve(args);
 }
+describe('NFT internal', () => {
+  describe('_increment function', () => {
+    throws('should throw an error when overflow occurs', () => {
+      // Set the counter to its maximum value
+      const maxU64: u64 = u64.MAX_VALUE;
+      Storage.set(counterKey, u64ToBytes(maxU64));
+
+      // Expect an error to be thrown when trying to increment beyond the maximum value
+      _increment();
+    });
+
+    test('should increment the counter', () => {
+      // Set the counter to 0
+      Storage.set(counterKey, u64ToBytes(0));
+
+      // Expect the counter to be incremented
+      expect(_increment()).toBe(1);
+      const result = Storage.get(counterKey);
+      expect(result).toStrictEqual(u64ToBytes(1));
+    });
+  });
+
+  describe('_updateBalanceOf function', () => {
+    it('should throw an error when overflow occurs', () => {
+      const maxU64 = u64.MAX_VALUE;
+      const balanceKey = stringToBytes('balanceOf_testAddress');
+      Storage.set(balanceKey, u64ToBytes(maxU64));
+
+      expect(() => _updateBalanceOf('testAddress', true)).toThrow();
+    });
+
+    it('should throw an error when underflow occurs', () => {
+      const balanceKey = stringToBytes('balanceOf_testAddress');
+      Storage.set(balanceKey, u64ToBytes(0));
+
+      expect(() => _updateBalanceOf('testAddress', false)).toThrow();
+    });
+
+    it('should correctly increment the balance', () => {
+      const balanceKey = stringToBytes('balanceOf_testAddress');
+      Storage.set(balanceKey, u64ToBytes(10));
+
+      _updateBalanceOf('testAddress', true);
+      const newBalance = bytesToU64(Storage.get(balanceKey));
+      expect(newBalance).toBe(11);
+    });
+
+    it('should correctly decrement the balance', () => {
+      const balanceKey = stringToBytes('balanceOf_testAddress');
+      Storage.set(balanceKey, u64ToBytes(10));
+
+      _updateBalanceOf('testAddress', false);
+      const newBalance = bytesToU64(Storage.get(balanceKey));
+      expect(newBalance).toBe(9);
+    });
+  });
+});