chore(contracts): Stop counting expired/revoked/replaced Attestation …

…via the balanceOf function (#835)
Consensys · Nov 28, 2024 · afc3d34 · afc3d34
1 parent 3ce6449
commit afc3d34
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 17 deletions.
diff --git a/contracts/src/AttestationRegistry.sol b/contracts/src/AttestationRegistry.sol
@@ -293,20 +293,23 @@ contract AttestationRegistry is RouterManager {
   }
 
   /**
-   * @notice Checks if an address owns a given attestation following ERC-1155
+   * @notice Checks if an address owns a valid attestation following the ERC-1155 interface
    * @param account The address of the token holder
    * @param id ID of the attestation
    * @return The _owner's balance of the attestations on a given attestation ID
+   * @dev Only considers non-revoked, non-replaced and non-expired attestations
    */
   function balanceOf(address account, uint256 id) public view returns (uint256) {
     bytes32 attestationId = generateAttestationId(id);
     Attestation memory attestation = attestations[attestationId];
-    if (attestation.subject.length > 20 && keccak256(attestation.subject) == keccak256(abi.encode(account))) {
-      return 1;
-    }
-    if (attestation.subject.length == 20 && keccak256(attestation.subject) == keccak256(abi.encodePacked(account))) {
-      return 1;
-    }
+
+    if (attestation.attestationId == bytes32(0)) return 0;
+    if (attestation.revoked == true) return 0;
+    if (attestation.expirationDate != 0 && attestation.expirationDate < block.timestamp) return 0;
+
+    if (attestation.subject.length == 32 && abi.decode(attestation.subject, (address)) == account) return 1;
+    if (attestation.subject.length == 20 && address(uint160(bytes20(attestation.subject))) == account) return 1;
+
     return 0;
   }
 

diff --git a/contracts/test/AttestationRegistry.t.sol b/contracts/test/AttestationRegistry.t.sol
@@ -598,45 +598,105 @@ contract AttestationRegistryTest is Test {
     vm.stopPrank();
   }
 
-  function test_balanceOf_notFound(AttestationPayload memory attestationPayload) public {
+  function test_balanceOf_attestationNotFound(AttestationPayload memory attestationPayload) public {
     vm.assume(attestationPayload.subject.length != 0);
     vm.assume(attestationPayload.attestationData.length != 0);
     SchemaRegistryMock schemaRegistryMock = SchemaRegistryMock(router.getSchemaRegistry());
     attestationPayload.schemaId = schemaRegistryMock.getIdFromSchemaString("schemaString");
     schemaRegistryMock.createSchema("name", "description", "context", "schemaString");
 
     vm.startPrank(portal);
-    attestationPayload.subject = abi.encode(address(1234));
+    attestationPayload.subject = abi.encode(address(1));
+    attestationRegistry.attest(attestationPayload, attester);
+
+    uint256 balance = attestationRegistry.balanceOf(address(1), 1234);
+    assertEq(balance, 0);
+  }
+
+  function test_balanceOf_attestationRevoked() public {
+    SchemaRegistryMock schemaRegistryMock = SchemaRegistryMock(router.getSchemaRegistry());
+    schemaRegistryMock.createSchema("name", "description", "context", "schemaString");
+
+    AttestationPayload memory attestationPayload;
+    attestationPayload.schemaId = schemaRegistryMock.getIdFromSchemaString("schemaString");
+    attestationPayload.expirationDate = 0;
+    attestationPayload.subject = abi.encode(address(1));
+    attestationPayload.attestationData = abi.encode("data");
+
+    vm.startPrank(portal);
     attestationRegistry.attest(attestationPayload, attester);
+    vm.startPrank(portal);
+    attestationRegistry.revoke(0x0003000000000000000000000000000000000000000000000000000000000001);
 
     uint256 balance = attestationRegistry.balanceOf(address(1), 1);
     assertEq(balance, 0);
   }
 
-  function test_balanceOf_subjectEncodedAddress(AttestationPayload memory attestationPayload) public {
-    vm.assume(attestationPayload.subject.length != 0);
-    vm.assume(attestationPayload.attestationData.length != 0);
+  function test_balanceOf_attestationExpired() public {
+    // Sets the block timestamp to 1641070800
+    vm.warp(1641070800);
+
     SchemaRegistryMock schemaRegistryMock = SchemaRegistryMock(router.getSchemaRegistry());
+    schemaRegistryMock.createSchema("name", "description", "context", "schemaString");
+
+    AttestationPayload memory attestationPayload;
     attestationPayload.schemaId = schemaRegistryMock.getIdFromSchemaString("schemaString");
+    attestationPayload.expirationDate = block.timestamp > 1 days ? uint64(block.timestamp - 1 days) : 0;
+    attestationPayload.subject = abi.encode(address(1));
+    attestationPayload.attestationData = abi.encode("data");
+
+    vm.startPrank(portal);
+    attestationRegistry.attest(attestationPayload, attester);
+
+    uint256 balance = attestationRegistry.balanceOf(address(1), 1);
+    assertEq(balance, 0);
+  }
+
+  function test_balanceOf_subjectNotFound() public {
+    SchemaRegistryMock schemaRegistryMock = SchemaRegistryMock(router.getSchemaRegistry());
     schemaRegistryMock.createSchema("name", "description", "context", "schemaString");
 
+    AttestationPayload memory attestationPayload;
+    attestationPayload.schemaId = schemaRegistryMock.getIdFromSchemaString("schemaString");
+    attestationPayload.expirationDate = 0;
+    attestationPayload.subject = abi.encode(address(1234));
+    attestationPayload.attestationData = abi.encode("data");
+
     vm.startPrank(portal);
+    attestationRegistry.attest(attestationPayload, attester);
+
+    uint256 balance = attestationRegistry.balanceOf(address(1), 1);
+    assertEq(balance, 0);
+  }
+
+  function test_balanceOf_subjectEncodedAddress() public {
+    SchemaRegistryMock schemaRegistryMock = SchemaRegistryMock(router.getSchemaRegistry());
+    schemaRegistryMock.createSchema("name", "description", "context", "schemaString");
+
+    AttestationPayload memory attestationPayload;
+    attestationPayload.schemaId = schemaRegistryMock.getIdFromSchemaString("schemaString");
+    attestationPayload.expirationDate = 0;
     attestationPayload.subject = abi.encode(address(1));
+    attestationPayload.attestationData = abi.encode("data");
+
+    vm.startPrank(portal);
     attestationRegistry.attest(attestationPayload, attester);
 
     uint256 balance = attestationRegistry.balanceOf(address(1), 1);
     assertEq(balance, 1);
   }
 
-  function test_balanceOf_subjectPackedEncodedAddress(AttestationPayload memory attestationPayload) public {
-    vm.assume(attestationPayload.subject.length != 0);
-    vm.assume(attestationPayload.attestationData.length != 0);
+  function test_balanceOf_subjectPackedEncodedAddress() public {
     SchemaRegistryMock schemaRegistryMock = SchemaRegistryMock(router.getSchemaRegistry());
-    attestationPayload.schemaId = schemaRegistryMock.getIdFromSchemaString("schemaString");
     schemaRegistryMock.createSchema("name", "description", "context", "schemaString");
 
-    vm.startPrank(portal);
+    AttestationPayload memory attestationPayload;
+    attestationPayload.schemaId = schemaRegistryMock.getIdFromSchemaString("schemaString");
+    attestationPayload.expirationDate = 0;
     attestationPayload.subject = abi.encodePacked(address(1));
+    attestationPayload.attestationData = abi.encode("data");
+
+    vm.startPrank(portal);
     attestationRegistry.attest(attestationPayload, attester);
 
     uint256 balance = attestationRegistry.balanceOf(address(1), 1);