Repository: OpenMined/PySyft
Pull request: Add a `protocol` kwarg to `encrypt()` method to select between MPC and Paillier HE
Fixes [#3212](https://github.com/OpenMined/PySyft/issues/3212)

> Modify .encrypt() to add a kwarg and be used not only for Paillier encryption but only MPC (default should be MPC, and it should call for mpc .fix_prec(...).share(...) under the hood)

Problem statement:
Enhance encryption API
Two features here:

- Make `crypto_provider` kwargs mandatory in `.share()`and check that it's not a client_worker
- Modify `.encrypt()` to add a kwarg and be used not only for Paillier encryption but only MPC (default should be MPC,  and it should call for mpc .fix_prec(...).share(...)  under the hood)

Hints from issue discussion:
I can try the first part.

Commit message: Cover edge case in test
?

Patch to review:
diff --git a/syft/frameworks/torch/tensors/interpreters/native.py b/syft/frameworks/torch/tensors/interpreters/native.py
index d69acac7a5f..9230c1034b0 100644
--- a/syft/frameworks/torch/tensors/interpreters/native.py
+++ b/syft/frameworks/torch/tensors/interpreters/native.py
@@ -992,19 +992,57 @@ def torch_type(self):
         else:
             return self.child.torch_type()
 
-    def encrypt(self, public_key):
-        """This method will encrypt each value in the tensor using Paillier
-        homomorphic encryption.
+    def encrypt(self, encryption_method="mpc", **kwargs):
+        """
+        This method will encrypt each value in the tensor using Multi Party
+        Computation (default) or Paillier Homomorphic Encryption
 
         Args:
-            *public_key a public key created using
-                syft.frameworks.torch.he.paillier.keygen()
+            encryption_method (str): Currently supports 'mpc' for Multi Party
+                Computation and 'paillier' for Paillier Homomorphic Encryption
+            **kwargs:
+                With Respect to MPC accepts:
+                    workers (list): Parties involved in the sharing of the Tensor
+                    crypto_provider (syft.VirtualWorker): Worker responsible for the
+                        generation of the random numbers for encryption
+                    args_fix_prec (dict): To be parsed as a kwarg for the .fix_prec() method
+
+                With Respect to Paillier accepts:
+                    public_key (phe.paillier.PaillierPublicKey): Can be obtained using
+                        ```public_key, private_key = sy.frameworks.torch.he.paillier.keygen()```
+        Returns:
+            An encrypted version of the Tensor following the protocol specified
+
+        Raises:
+            NotImplementedError: If protocols other than the ones mentioned above are queried
+
         """
+        if encryption_method.lower() == "mpc":
+            args_fix_prec = kwargs.pop("args_fix_prec")
+            workers = kwargs.pop("workers")
+            crypto_provider = kwargs.pop("crypto_provider")
+
+            if args_fix_prec:
+                x_shared = self.fix_prec(**args_fix_prec).share(
+                    *workers, crypto_provider=crypto_provider, **kwargs
+                )
+            else:  # If no args for .fix_prec()
+                x_shared = self.fix_prec().share(
+                    *workers, crypto_provider=crypto_provider, **kwargs
+                )
+            return x_shared
+
+        elif encryption_method.lower() == "paillier":
+            x = self.copy()
+            x_encrypted = PaillierTensor().on(x)  # Instantiate the class
+            x_encrypted.child.encrypt_(kwargs["public_key"])  # Perform blacHomomorphic Encryption
+            return x_encrypted
 
-        x = self.copy()
-        x2 = PaillierTensor().on(x)
-        x2.child.encrypt_(public_key)
-        return x2
+        else:
+            raise NotImplementedError(
+                "Currently the .encrypt() method only supports Paillier Homomorphic "
+                "Encryption and Secure Multi-Party Computation"
+            )
 
     def decrypt(self, private_key):
         """This method will decrypt each value in the tensor, returning a normal
diff --git a/test/torch/tensors/test_native.py b/test/torch/tensors/test_native.py
index bb5c12d00ac..f935824f57d 100644
--- a/test/torch/tensors/test_native.py
+++ b/test/torch/tensors/test_native.py
@@ -205,3 +205,24 @@ def forward(self, x):
 
     ## Forward on the remote model
     pred = model_net(tensor_remote)
+
+
+def test_encrypt():
+    hook = syft.TorchHook(torch)
+
+    bob = syft.VirtualWorker(hook, id="bob")
+    alice = syft.VirtualWorker(hook, id="alice")
+    james = syft.VirtualWorker(hook, id="james")
+
+    x = torch.randint(10, (1, 5), dtype=torch.float32)
+    x_encrypted = x.encrypt(workers=[bob, alice], crypto_provider=james, args_fix_prec={"base": 10})
+    assert torch.all(torch.eq(x_encrypted.get().float_prec(), x))
+
+    x = torch.randint(10, (1, 5), dtype=torch.float32)
+    x_encrypted = x.encrypt(workers=[bob, alice], crypto_provider=james)
+    assert torch.all(torch.eq(x_encrypted.get().float_prec(), x))
+
+    x = torch.randint(10, (1, 5), dtype=torch.float32)
+    public, private = syft.frameworks.torch.he.paillier.keygen()
+    x_encrypted = x.encrypt(encryption_method="paillier", public_key=public)
+    assert torch.all(torch.eq(x_encrypted.decrypt(private), x))
diff --git a/test/torch/tensors/test_paillier.py b/test/torch/tensors/test_paillier.py
index 70837643475..bc6d2deb925 100644
--- a/test/torch/tensors/test_paillier.py
+++ b/test/torch/tensors/test_paillier.py
@@ -11,7 +11,7 @@ def test_encrypt_and_decrypt():
     pub, pri = sy.keygen()
 
     x_tensor = torch.Tensor([1, 2, 3])
-    x = x_tensor.encrypt(pub)
+    x = x_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     y = x.decrypt(pri)
 
@@ -26,7 +26,7 @@ def test_encrypted_encrypted_add():
     pub, pri = sy.keygen()
 
     x_tensor = torch.Tensor([1, 2, 3])
-    x = x_tensor.encrypt(pub)
+    x = x_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     y = (x + x).decrypt(pri)
 
@@ -41,7 +41,7 @@ def test_encrypted_decrypted_add():
     pub, pri = sy.keygen()
 
     x_tensor = torch.Tensor([1, 2, 3])
-    x = x_tensor.encrypt(pub)
+    x = x_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     y = (x + x_tensor).decrypt(pri)
 
@@ -56,7 +56,7 @@ def test_decrypted_encrypted_add():
     pub, pri = sy.keygen()
 
     x_tensor = torch.Tensor([1, 2, 3])
-    x = x_tensor.encrypt(pub)
+    x = x_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     y = (x_tensor + x).decrypt(pri)
 
@@ -68,10 +68,10 @@ def test_encrypted_encrypted_sub():
     pub, pri = sy.keygen()
 
     x_tensor = torch.Tensor([1, 2, 3])
-    x = x_tensor.encrypt(pub)
+    x = x_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     y_tensor = torch.Tensor([2, 2, 2])
-    y = y_tensor.encrypt(pub)
+    y = y_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     z = (x - y).decrypt(pri)
 
@@ -83,10 +83,10 @@ def test_encrypted_decrypted_sub():
     pub, pri = sy.keygen()
 
     x_tensor = torch.Tensor([1, 2, 3])
-    x = x_tensor.encrypt(pub)
+    x = x_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     y_tensor = torch.Tensor([2, 2, 2])
-    y = y_tensor.encrypt(pub)
+    y = y_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     z = (x - y_tensor).decrypt(pri)
 
@@ -98,10 +98,10 @@ def test_decrypted_encrypted_sub():
     pub, pri = sy.keygen()
 
     x_tensor = torch.Tensor([1, 2, 3])
-    x = x_tensor.encrypt(pub)
+    x = x_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     y_tensor = torch.Tensor([2, 2, 2])
-    y = y_tensor.encrypt(pub)
+    y = y_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     z = (x_tensor - y).decrypt(pri)
 
@@ -113,10 +113,10 @@ def test_encrypted_decrypted_mul():
     pub, pri = sy.keygen()
 
     x_tensor = torch.Tensor([1, 2, 3])
-    x = x_tensor.encrypt(pub)
+    x = x_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     y_tensor = torch.Tensor([2, 2, 2])
-    y = y_tensor.encrypt(pub)
+    y = y_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     z = (x * y_tensor).decrypt(pri)
 
@@ -128,10 +128,10 @@ def test_decrypted_encrypted_mul():
     pub, pri = sy.keygen()
 
     x_tensor = torch.Tensor([1, 2, 3])
-    x = x_tensor.encrypt(pub)
+    x = x_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     y_tensor = torch.Tensor([2, 2, 2])
-    y = y_tensor.encrypt(pub)
+    y = y_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     z = (x_tensor * y).decrypt(pri)
 
@@ -143,10 +143,10 @@ def test_encrypted_decrypted_matmul():
     pub, pri = sy.keygen()
 
     x_tensor = torch.tensor([1, 2, 3])
-    x = x_tensor.encrypt(pub)
+    x = x_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     y_tensor = torch.tensor([2, 2, 2])
-    y = y_tensor.encrypt(pub)
+    y = y_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     z = (x.mm(y_tensor)).decrypt(pri)
 
@@ -158,10 +158,10 @@ def test_decrypted_encrypted_matmul():
     pub, pri = sy.keygen()
 
     x_tensor = torch.Tensor([[1, 2, 3]])
-    x = x_tensor.encrypt(pub)
+    x = x_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     y_tensor = torch.Tensor([[2], [2], [2]])
-    y = y_tensor.encrypt(pub)
+    y = y_tensor.encrypt(encryption_method="paillier", public_key=pub)
 
     z = (x_tensor.mm(y)).decrypt(pri)
 


Code review comments:
diff_hunk: @@ -992,19 +992,57 @@ def torch_type(self):
         else:
             return self.child.torch_type()
 
-    def encrypt(self, public_key):
-        """This method will encrypt each value in the tensor using Paillier
-        homomorphic encryption.
+    def encrypt(self, encryption_method="mpc", **kwargs):
+        """
+        This method will encrypt each value in the tensor using Multi Party
+        Computation (default) or Paillier Homomorphic Encryption
 
         Args:
-            *public_key a public key created using
-                syft.frameworks.torch.he.paillier.keygen()
+            encryption_method (str): Currently supports 'mpc' for Multi Party
+                Computation and 'paillier' for Paillier Homomorphic Encryption
+            **kwargs:
+                With Respect to MPC accepts:
+                    workers (list): Parties involved in the sharing of the Tensor
+                    crypto_provider (syft.VirtualWorker): Worker responsible for the
+                        generation of the random numbers for encryption
+                    args_fix_prec (dict): To be parsed as a kwarg for the .fix_prec() method
+
+                With Respect to Paillier accepts:
+                    public_key (phe.paillier.PaillierPublicKey): Can be obtained using
+                        ```public_key, private_key = sy.frameworks.torch.he.paillier.keygen()```
+        Returns:
+            An encrypted version of the Tensor following the protocol specified
+
+        Raises:
+            NotImplementedError: If protocols other than the ones mentioned above are queried
+
         """
+        if encryption_method.lower() == "mpc":
+            args_fix_prec = kwargs.pop("args_fix_prec")
+            workers = kwargs.pop("workers")
+            crypto_provider = kwargs.pop("crypto_provider")
+
+            if args_fix_prec:
+                x_shared = self.fix_prec(**args_fix_prec).share(
+                    *workers, crypto_provider=crypto_provider, **kwargs
path: syft/frameworks/torch/tensors/interpreters/native.py
text: @user1:
with my suggestion, then you will had to remove kwargs here I guess

diff_hunk: @@ -205,3 +205,24 @@ def forward(self, x):
 
     ## Forward on the remote model
     pred = model_net(tensor_remote)
+
+
+def test_encrypt():
+    hook = syft.TorchHook(torch)
path: test/torch/tensors/test_native.py
text: @user1:
Check other examples you can improt workers directly:

ex: in test_additive_sharing.py
```
def test_eq(workers):
    me, alice, bob, crypto_provider = (
        workers["me"],
        workers["alice"],
        workers["bob"],
        workers["james"],
    )
```

diff_hunk: @@ -205,3 +205,24 @@ def forward(self, x):
 
     ## Forward on the remote model
     pred = model_net(tensor_remote)
+
+
+def test_encrypt():
+    hook = syft.TorchHook(torch)
+
+    bob = syft.VirtualWorker(hook, id="bob")
+    alice = syft.VirtualWorker(hook, id="alice")
+    james = syft.VirtualWorker(hook, id="james")
+
+    x = torch.randint(10, (1, 5), dtype=torch.float32)
+    x_encrypted = x.encrypt(workers=[bob, alice], crypto_provider=james, args_fix_prec={"base": 10})
+    assert torch.all(torch.eq(x_encrypted.get().float_prec(), x))
path: test/torch/tensors/test_native.py
text: @user1:
Oh we've forgotten to add a generic decrypt method!!
Would you be willing to do also this one? :) 
for MPC .decrypt() would call `.get().float_prec()` 
and for paillier, it would'nt change

@author:
Sure. Do I submit it in another PR?
I've incorporated the suggestions mentioned. 

@user1:
In this PR would be more coherent if you can do it !

diff_hunk: @@ -992,19 +992,57 @@ def torch_type(self):
         else:
             return self.child.torch_type()
 
-    def encrypt(self, public_key):
-        """This method will encrypt each value in the tensor using Paillier
-        homomorphic encryption.
+    def encrypt(self, encryption_method="mpc", **kwargs):
+        """
+        This method will encrypt each value in the tensor using Multi Party
+        Computation (default) or Paillier Homomorphic Encryption
 
         Args:
-            *public_key a public key created using
-                syft.frameworks.torch.he.paillier.keygen()
+            encryption_method (str): Currently supports 'mpc' for Multi Party
+                Computation and 'paillier' for Paillier Homomorphic Encryption
+            **kwargs:
+                With Respect to MPC accepts:
+                    workers (list): Parties involved in the sharing of the Tensor
+                    crypto_provider (syft.VirtualWorker): Worker responsible for the
+                        generation of the random numbers for encryption
+                    args_fix_prec (dict): To be parsed as a kwarg for the .fix_prec() method
+
+                With Respect to Paillier accepts:
+                    public_key (phe.paillier.PaillierPublicKey): Can be obtained using
+                        ```public_key, private_key = sy.frameworks.torch.he.paillier.keygen()```
+        Returns:
+            An encrypted version of the Tensor following the protocol specified
+
+        Raises:
+            NotImplementedError: If protocols other than the ones mentioned above are queried
+
         """
+        if encryption_method.lower() == "mpc":
+            args_fix_prec = kwargs.pop("args_fix_prec")
path: syft/frameworks/torch/tensors/interpreters/native.py
text: @user1:
Use .get() when the parameter is not compulsory (like in fix prec)

diff_hunk: @@ -992,19 +992,57 @@ def torch_type(self):
         else:
             return self.child.torch_type()
 
-    def encrypt(self, public_key):
-        """This method will encrypt each value in the tensor using Paillier
-        homomorphic encryption.
+    def encrypt(self, encryption_method="mpc", **kwargs):
+        """
+        This method will encrypt each value in the tensor using Multi Party
+        Computation (default) or Paillier Homomorphic Encryption
 
         Args:
-            *public_key a public key created using
-                syft.frameworks.torch.he.paillier.keygen()
+            encryption_method (str): Currently supports 'mpc' for Multi Party
+                Computation and 'paillier' for Paillier Homomorphic Encryption
+            **kwargs:
+                With Respect to MPC accepts:
+                    workers (list): Parties involved in the sharing of the Tensor
+                    crypto_provider (syft.VirtualWorker): Worker responsible for the
+                        generation of the random numbers for encryption
+                    args_fix_prec (dict): To be parsed as a kwarg for the .fix_prec() method
path: syft/frameworks/torch/tensors/interpreters/native.py
text: @user1:
Could you also allow fix_precision kwargs to be given alongside the share ones (workers & crypto_provider)? Would be simpler for the end user :) 

@user1:
(I think we would provide kwargs and not args for fix_precision)

diff_hunk: @@ -992,19 +992,57 @@ def torch_type(self):
         else:
             return self.child.torch_type()
 
-    def encrypt(self, public_key):
-        """This method will encrypt each value in the tensor using Paillier
-        homomorphic encryption.
+    def encrypt(self, encryption_method="mpc", **kwargs):
+        """
+        This method will encrypt each value in the tensor using Multi Party
+        Computation (default) or Paillier Homomorphic Encryption
 
         Args:
-            *public_key a public key created using
-                syft.frameworks.torch.he.paillier.keygen()
+            encryption_method (str): Currently supports 'mpc' for Multi Party
+                Computation and 'paillier' for Paillier Homomorphic Encryption
+            **kwargs:
+                With Respect to MPC accepts:
+                    workers (list): Parties involved in the sharing of the Tensor
+                    crypto_provider (syft.VirtualWorker): Worker responsible for the
+                        generation of the random numbers for encryption
+                    args_fix_prec (dict): To be parsed as a kwarg for the .fix_prec() method
+
+                With Respect to Paillier accepts:
+                    public_key (phe.paillier.PaillierPublicKey): Can be obtained using
+                        ```public_key, private_key = sy.frameworks.torch.he.paillier.keygen()```
+        Returns:
+            An encrypted version of the Tensor following the protocol specified
+
+        Raises:
+            NotImplementedError: If protocols other than the ones mentioned above are queried
+
         """
+        if encryption_method.lower() == "mpc":
+            args_fix_prec = kwargs.pop("args_fix_prec")
+            workers = kwargs.pop("workers")
+            crypto_provider = kwargs.pop("crypto_provider")
+
+            if args_fix_prec:
+                x_shared = self.fix_prec(**args_fix_prec).share(
+                    *workers, crypto_provider=crypto_provider, **kwargs
+                )
+            else:  # If no args for .fix_prec()
+                x_shared = self.fix_prec().share(
+                    *workers, crypto_provider=crypto_provider, **kwargs
+                )
+            return x_shared
+
+        elif encryption_method.lower() == "paillier":
+            x = self.copy()
+            x_encrypted = PaillierTensor().on(x)  # Instantiate the class
+            x_encrypted.child.encrypt_(kwargs["public_key"])  # Perform blacHomomorphic Encryption
path: syft/frameworks/torch/tensors/interpreters/native.py
text: @user1:
Extract public_key like you did for mpc (it's nicer :) ) 
and typo: blacHomomorphic

diff_hunk: @@ -992,19 +992,57 @@ def torch_type(self):
         else:
             return self.child.torch_type()
 
-    def encrypt(self, public_key):
-        """This method will encrypt each value in the tensor using Paillier
-        homomorphic encryption.
+    def encrypt(self, encryption_method="mpc", **kwargs):
path: syft/frameworks/torch/tensors/interpreters/native.py
text: @user1:
I have a slight preference for `protocol` instead of `encryption_method` (mostly because its shorter and because we often speak of crypto protocols)