Merge pull request #333 from EYBlockchain/lydia/ownershipError

fix: ensure whenever we have mappings of addresses, the address is al…
EYBlockchain · Sep 17, 2024 · 7e41ad7 · 7e41ad7
2 parents 5a5e9df + 6f51773
commit 7e41ad7
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 13 deletions.
diff --git a/src/traverse/Binding.ts b/src/traverse/Binding.ts
@@ -275,6 +275,7 @@ export class VariableBinding extends Binding {
   }
 
   updateOwnership(ownerNode: any, msgIsMappingKeyorMappingValue?: string | null) {
+    if (this.isOwned && this.owner.mappingOwnershipType === 'key') return;
     if (
       ownerNode.expression?.name === 'msg' &&
       msgIsMappingKeyorMappingValue === 'value'
@@ -521,19 +522,9 @@ export class VariableBinding extends Binding {
    */
   inferOwnership() {
     if (this.kind !== 'VariableDeclaration') return;
-    let msgSenderEverywhereMappingKey: boolean = false;
-    let msgSenderEverywhereMappingValue: boolean = false;
+    let msgSenderEverywhereMappingKey: boolean;
+    let msgSenderEverywhereMappingValue: boolean;
     this.nullifyingPaths.forEach(path => {
-      const functionDefScope = path.scope.getAncestorOfScopeType(
-        'FunctionDefinition',
-      );
-      if (functionDefScope.callerRestriction === 'match') {
-        this.updateOwnership(functionDefScope.callerRestrictionNode);
-        return;
-      }
-      if (functionDefScope.callerRestriction === 'exclude') {
-        this.updateBlacklist(functionDefScope.callerRestrictionNode);
-      }
       if (this.isMapping && path.getAncestorOfType('IndexAccess') && this.addMappingKey(path).isMsgSender ) {
         // if its unassigned, we assign true
         // if its true, it remains true
@@ -542,6 +533,7 @@ export class VariableBinding extends Binding {
         this.isMapping && path.getAncestorOfType('IndexAccess') &&
         (path.isMsgSender(path.getCorrespondingRhsNode()) || path.isMsgValue(path.getCorrespondingRhsNode()))
       ) {
+        msgSenderEverywhereMappingKey = false;
         msgSenderEverywhereMappingValue ??= true;
       } else {
         // if we find a single non-msg sender mapping key, then msg sender can't be the owner
@@ -562,6 +554,19 @@ export class VariableBinding extends Binding {
         this.nullifyingPaths[0].parentPath.parent.rightHandSide;
       this.updateOwnership(owner, 'value');
     }
+    this.nullifyingPaths.forEach(path => {
+      const functionDefScope = path.scope.getAncestorOfScopeType(
+        'FunctionDefinition',
+      );
+      // we update the ownership of the variable if the function has a restriction on the caller, e.g. require(msg.sender == admin)
+      if (functionDefScope.callerRestriction === 'match') {
+        this.updateOwnership(functionDefScope.callerRestrictionNode);
+        return;
+      }
+      if (functionDefScope.callerRestriction === 'exclude') {
+        this.updateBlacklist(functionDefScope.callerRestrictionNode);
+      }
+    }); 
   }
 
   ownerSetToZeroCheck() {

diff --git a/test/contracts/mapping-3.zol b/test/contracts/mapping-3.zol
@@ -8,17 +8,23 @@ contract MyContract {
   secret uint256 private b;
   address public admin;
 
+   constructor() {
+      admin = msg.sender;
+    }
+
   function assign(secret address param1, secret uint256 param2) public {
     unknown a[param1] += param2;
   }
 
   function assign2(secret uint256 param3) public {
+    require(msg.sender == admin);
     a[msg.sender] = a[msg.sender] - param3;
     known b += param3;
   }
 
-  function assign3(secret uint256 param5) public {
+  function assign3(secret uint256 param5, secret address recipient) public {
     require(msg.sender == admin);
     known b += param5;
+    a[recipient] += param5;
   }
 }