Repository: openmined/pysyft
Issue: Fix Store Mutation by Guest

## Description
This might be fixed already but I just want to make sure the old PR is resolved and merged.

[Fixed override store keys bug #6313](https://github.com/OpenMined/PySyft/pull/6313)


Changed files:
packages/syft/tests/syft/security/mutate_objects_test.py
packages/syft/src/syft/core/node/common/action/przs_action.py
packages/syft/src/syft/core/node/common/action/run_class_method_action.py
packages/syft/src/syft/core/node/common/node_manager/redis_store.py
packages/syft/src/syft/core/node/common/action/save_object_action.py
packages/syft/src/syft/core/node/common/node_service/upload_service/upload_service_messages.py
packages/syft/src/syft/core/node/common/node_manager/dict_store.py
packages/syft/tests/syft/core/pointer/pointer_test.py
packages/syft/src/syft/core/store/store_interface.py
packages/syft/src/syft/core/node/common/action/run_class_method_smpc_action.py
packages/syft/src/syft/core/node/common/action/function_or_constructor_action.py

diff --git a/packages/syft/src/syft/core/node/common/action/function_or_constructor_action.py b/packages/syft/src/syft/core/node/common/action/function_or_constructor_action.py
index aa91cda87..d27ddbbdf 100644
--- a/packages/syft/src/syft/core/node/common/action/function_or_constructor_action.py
+++ b/packages/syft/src/syft/core/node/common/action/function_or_constructor_action.py
@@ -82,6 +82,10 @@ class RunFunctionOrConstructorAction(ImmediateActionWithoutReply):
 
     def execute_action(self, node: AbstractNode, verify_key: VerifyKey) -> None:
         method = node.lib_ast(self.path)
+
+        # If if there's another object with the same ID.
+        node.store.check_collision(self.id_at_location)
+
         result_read_permissions: Union[None, Dict[VerifyKey, Optional[UID]]] = None
         result_write_permissions: Union[None, Dict[VerifyKey, Optional[UID]]] = {
             verify_key: None
diff --git a/packages/syft/src/syft/core/node/common/action/przs_action.py b/packages/syft/src/syft/core/node/common/action/przs_action.py
index 932e2458b..e04364cb4 100644
--- a/packages/syft/src/syft/core/node/common/action/przs_action.py
+++ b/packages/syft/src/syft/core/node/common/action/przs_action.py
@@ -76,6 +76,8 @@ class PRZSAction(ImmediateActionWithoutReply):
         return f"PRZSAction ({arg_names}, {kwargs_names})"
 
     def execute_action(self, node: AbstractNode, verify_key: VerifyKey) -> None:
+        # If if there's another object with the same ID.
+        node.store.check_collision(self.id_at_location)
 
         (
             upcasted_args,
diff --git a/packages/syft/src/syft/core/node/common/action/run_class_method_action.py b/packages/syft/src/syft/core/node/common/action/run_class_method_action.py
index b2884a116..42becc5ba 100644
--- a/packages/syft/src/syft/core/node/common/action/run_class_method_action.py
+++ b/packages/syft/src/syft/core/node/common/action/run_class_method_action.py
@@ -97,6 +97,9 @@ class RunClassMethodAction(ImmediateActionWithoutReply):
     def execute_action(self, node: AbstractNode, verify_key: VerifyKey) -> None:
         method = node.lib_ast(self.path)
 
+        # If if there's another object with the same ID.
+        node.store.check_collision(self.id_at_location)
+
         mutating_internal = False
         if (
             self.path.startswith("torch.Tensor")
diff --git a/packages/syft/src/syft/core/node/common/action/run_class_method_smpc_action.py b/packages/syft/src/syft/core/node/common/action/run_class_method_smpc_action.py
index 6276dd783..801274aa5 100644
--- a/packages/syft/src/syft/core/node/common/action/run_class_method_smpc_action.py
+++ b/packages/syft/src/syft/core/node/common/action/run_class_method_smpc_action.py
@@ -95,6 +95,9 @@ class RunClassMethodSMPCAction(ImmediateActionWithoutReply):
         # relative
         from ....tensor.smpc import context
 
+        # If if there's another object with the same ID.
+        node.store.check_collision(self.id_at_location)
+
         resolved_self = retrieve_object(node, self._self.id_at_location, self.path)
         result_read_permissions = resolved_self.read_permissions
 
diff --git a/packages/syft/src/syft/core/node/common/action/save_object_action.py b/packages/syft/src/syft/core/node/common/action/save_object_action.py
index 1c41eea70..bb0de354e 100644
--- a/packages/syft/src/syft/core/node/common/action/save_object_action.py
+++ b/packages/syft/src/syft/core/node/common/action/save_object_action.py
@@ -42,11 +42,9 @@ class SaveObjectAction(ImmediateActionWithoutReply):
         return f"SaveObjectAction {obj_str}"
 
     def execute_action(self, node: AbstractNode, verify_key: VerifyKey) -> None:
-        # Check if it uses an id previously registered.
-        old_obj = node.store.get_or_none(key=self.obj.id, proxy_only=True)
 
-        if old_obj:
-            raise Exception("You're not allowed to perform this operation.")
+        # If if there's another object with the same ID.
+        node.store.check_collision(self.obj.id)
 
         self.obj.read_permissions = {
             node.verify_key: node.id,
diff --git a/packages/syft/src/syft/core/node/common/node_manager/dict_store.py b/packages/syft/src/syft/core/node/common/node_manager/dict_store.py
index b3aeda11c..d74757a42 100644
--- a/packages/syft/src/syft/core/node/common/node_manager/dict_store.py
+++ b/packages/syft/src/syft/core/node/common/node_manager/dict_store.py
@@ -40,6 +40,13 @@ class DictStore(ObjectStore):
         except KeyError as e:  # noqa: F841
             return None
 
+    def check_collision(self, key: UID) -> None:
+        # Check ID collision with pointer's result.
+        if self.get_or_none(key=key, proxy_only=True):
+            raise Exception(
+                "You're not allowed to perform this operation using this ID."
+            )
+
     def get_objects_of_type(self, obj_type: type) -> Iterable[StorableObject]:
         # raise NotImplementedError("get_objects_of_type")
         # return [obj for obj in self.values() if isinstance(obj.data, obj_type)]
diff --git a/packages/syft/src/syft/core/node/common/node_manager/redis_store.py b/packages/syft/src/syft/core/node/common/node_manager/redis_store.py
index f0c155fed..ed86ef8d9 100644
--- a/packages/syft/src/syft/core/node/common/node_manager/redis_store.py
+++ b/packages/syft/src/syft/core/node/common/node_manager/redis_store.py
@@ -49,6 +49,13 @@ class RedisStore(ObjectStore):
         except KeyError:
             return None
 
+    def check_collision(self, key: UID) -> None:
+        # Check ID collision with pointer's result.
+        if self.get_or_none(key=key, proxy_only=True):
+            raise Exception(
+                "You're not allowed to perform this operation using this ID."
+            )
+
     def get_objects_of_type(self, obj_type: type) -> Iterable[StorableObject]:
         # TODO: remove this kind of operation which pulls all the data out in one go
         # raise NotImplementedError("get_objects_of_type")
diff --git a/packages/syft/src/syft/core/node/common/node_service/upload_service/upload_service_messages.py b/packages/syft/src/syft/core/node/common/node_service/upload_service/upload_service_messages.py
index bb4b664a6..c44d41c18 100644
--- a/packages/syft/src/syft/core/node/common/node_service/upload_service/upload_service_messages.py
+++ b/packages/syft/src/syft/core/node/common/node_service/upload_service/upload_service_messages.py
@@ -64,10 +64,8 @@ class UploadDataMessage(SyftMessage, DomainMessageRegistry):
         # Then we need to check if this UID was registered previously.
         if UID.is_valid_uuid(key.split("/")[-1]):
             id_at_location = UID.from_string(key.split("/")[-1])  # Get Object ID.
-            old_obj = node.store.get_or_none(key=id_at_location, proxy_only=True)
-
-            if old_obj:
-                raise Exception("You're not allowed to perform this operation.")
+            # If if there's another object with the same ID.
+            node.store.check_collision(id_at_location)
 
         s3_client = get_s3_client(settings=node.settings)
         result = s3_client.create_multipart_upload(Bucket=node.id.no_dash, Key=key)
diff --git a/packages/syft/src/syft/core/store/store_interface.py b/packages/syft/src/syft/core/store/store_interface.py
index afd84361d..5bc57197e 100644
--- a/packages/syft/src/syft/core/store/store_interface.py
+++ b/packages/syft/src/syft/core/store/store_interface.py
@@ -142,6 +142,9 @@ class ObjectStore(ABC):
     ) -> Optional[StorableObject]:
         traceback_and_raise(NotImplementedError)
 
+    def check_collision(self, key: UID) -> None:
+        traceback_and_raise(NotImplementedError)
+
     def get(self, key: StoreKey, proxy_only: bool = False) -> StorableObject:
         traceback_and_raise(NotImplementedError)
 
diff --git a/packages/syft/tests/syft/core/pointer/pointer_test.py b/packages/syft/tests/syft/core/pointer/pointer_test.py
index 704f16194..c3acd93d2 100644
--- a/packages/syft/tests/syft/core/pointer/pointer_test.py
+++ b/packages/syft/tests/syft/core/pointer/pointer_test.py
@@ -91,7 +91,10 @@ def test_overwrite_tensor(
     with pytest.raises(Exception) as exc_info:
         client.send_immediate_msg_without_reply(msg=obj_msg)
 
-    assert str(exc_info.value) == "You're not allowed to perform this operation."
+    assert (
+        str(exc_info.value)
+        == "You're not allowed to perform this operation using this ID."
+    )
 
     # Check if tensor remains the same
     assert tensor_ptr.get_copy() == ten
diff --git a/packages/syft/tests/syft/security/mutate_objects_test.py b/packages/syft/tests/syft/security/mutate_objects_test.py
index 465d2bdf8..853c0248a 100644
--- a/packages/syft/tests/syft/security/mutate_objects_test.py
+++ b/packages/syft/tests/syft/security/mutate_objects_test.py
@@ -2,13 +2,13 @@
 from copy import copy
 
 # third party
+import pytest
 import torch as th
 
 # syft absolute
 import syft as sy
 from syft.core.common.uid import UID
-
-# from syft.core.node.common.action.run_class_method_action import RunClassMethodAction
+from syft.core.node.common.action.run_class_method_action import RunClassMethodAction
 
 
 def test_store_object_mutation(
@@ -75,40 +75,47 @@ def test_store_object_mutation(
     assert all(new_result == (y + y)) is True
 
 
-# TODO: Fix
-# def test_store_overwrite_key(
-#     client: sy.VirtualMachineClient, root_client: sy.VirtualMachineClient
-# ) -> None:
-#     """Check if someone can overwrite any address."""
-
-#     # root creates two tensors one visible one not
-#     x = th.tensor([1, 2, 3])
-#     x_ptr = x.send(root_client, pointable=True, tags=["visible"])
-
-#     y = th.tensor([3, 6, 9])
-#     y_ptr = y.send(root_client, pointable=False, tags=["invisible"])
-
-#     # guest gets a pointer to the visible one
-#     guest_x = client.store[x_ptr.id_at_location]
-#     guest_x.add_(guest_x)
-
-#     # guest constructs a pointer to the guessed hidden object
-#     guest_y = copy(guest_x)
-#     target_uid = UID.from_string(y_ptr.id_at_location.no_dash)
-#     guest_y.id_at_location = target_uid
-#     guest_y.add_(guest_y)
-
-#     # guest should not be able to overwrite a destination with RunClassMethodAction
-#     cmd = RunClassMethodAction(
-#         path="torch.Tensor.add_",
-#         _self=guest_y,
-#         args=[guest_y],
-#         kwargs={},
-#         id_at_location=target_uid,  # normal destination changed
-#         address=client.address,
-#     )
-#     client.send_immediate_msg_without_reply(msg=cmd)
-
-#     # y should not have changed
-#     y_result = y_ptr.get(delete_obj=False)
-#     assert all(y_result == y) is True
+def test_store_overwrite_key(
+    client: sy.VirtualMachineClient, root_client: sy.VirtualMachineClient
+) -> None:
+    """Check if someone can overwrite any address."""
+
+    # root creates two tensors one visible one not
+    x = th.tensor([1, 2, 3])
+    x_ptr = x.send(root_client, pointable=True, tags=["visible"])
+
+    y = th.tensor([3, 6, 9])
+    y_ptr = y.send(root_client, pointable=False, tags=["invisible"])
+
+    # guest gets a pointer to the visible one
+    guest_x = client.store[x_ptr.id_at_location]
+    guest_x.add_(guest_x)
+
+    # guest constructs a pointer to the guessed hidden object
+    guest_y = copy(guest_x)
+    target_uid = sy.common.UID.from_string(y_ptr.id_at_location.no_dash)
+    guest_y.id_at_location = target_uid
+    guest_y.add_(guest_y)
+
+    # guest should not be able to overwrite a destination with RunClassMethodAction
+    cmd = RunClassMethodAction(
+        path="torch.Tensor.add_",
+        _self=guest_y,
+        args=[guest_y],
+        kwargs={},
+        id_at_location=target_uid,  # normal destination changed
+        address=client.address,
+    )
+
+    # ID collision exception
+    with pytest.raises(Exception) as exc_info:
+        client.send_immediate_msg_without_reply(msg=cmd)
+
+    assert (
+        str(exc_info.value)
+        == "You're not allowed to perform this operation using this ID."
+    )
+
+    # y should not have changed
+    y_result = y_ptr.get(delete_obj=False)
+    assert all(y_result == y) is True