mplang-nightly 0.1.dev261__tar.gz → 0.1.dev262__tar.gz

{mplang_nightly-0.1.dev261 → mplang_nightly-0.1.dev262}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev261
+Version: 0.1.dev262
 Summary: Multi-Party Programming Language
 Author-email: SecretFlow Team <secretflow-contact@service.alipay.com>
 License:                                  Apache License

{mplang_nightly-0.1.dev261 → mplang_nightly-0.1.dev262}/mplang/v2/backends/crypto_impl.py

@@ -271,6 +271,27 @@ def bytes_to_point_impl(
 
 # --- Sym / Hash Impl ---
 
+# Supported symmetric encryption algorithms
+_SUPPORTED_ALGOS = {"aes-gcm"}
+
+
+def _validate_algo(algo: str, operation: str) -> None:
+    """Validate that the algorithm is supported.
+
+    Args:
+        algo: Algorithm name to validate
+        operation: Operation name for error message (e.g., "encryption", "decryption")
+
+    Raises:
+        ValueError: If algo is not supported
+    """
+    if algo not in _SUPPORTED_ALGOS:
+        supported = ", ".join(sorted(_SUPPORTED_ALGOS))
+        raise ValueError(
+            f"Unsupported {operation} algorithm: {algo!r}. "
+            f"Supported algorithms: {supported}"
+        )
+
 
 @crypto.hash_p.def_impl
 def hash_impl(interpreter: Interpreter, op: Operation, data: Value) -> Value:
@@ -336,6 +357,10 @@ def sym_encrypt_impl(
     numpy arrays, scalars, etc.). This supports both high-level API usage
     (with TensorValue) and elementwise operations (with raw scalars).
     """
+    # Read and validate algo parameter (must be provided by frontend)
+    algo = op.attrs["algo"]
+    _validate_algo(algo, "encryption")
+
     # Get raw key bytes - strict type checking
     if isinstance(key, SymmetricKeyValue):
         k = key.key_bytes
@@ -376,6 +401,10 @@ def sym_decrypt_impl(
     on what was encrypted - could be a Value subclass (TensorValue, BytesValue),
     a numpy array, or a scalar (int, float, etc.) when used in elementwise ops.
     """
+    # Read and validate algo parameter (must be provided by frontend)
+    algo = op.attrs["algo"]
+    _validate_algo(algo, "decryption")
+
     # Get raw key bytes - strict type checking
     if isinstance(key, SymmetricKeyValue):
         k = key.key_bytes

{mplang_nightly-0.1.dev261 → mplang_nightly-0.1.dev262}/mplang/v2/dialects/crypto.py

@@ -275,15 +275,44 @@ def _hash_batch_ae(data: elt.BaseType) -> elt.TensorType:
 
 
 @sym_encrypt_p.def_abstract_eval
-def _sym_encrypt_ae(key: elt.BaseType, plaintext: elt.BaseType) -> elt.TensorType:
+def _sym_encrypt_ae(
+    key: elt.BaseType, plaintext: elt.BaseType, *, algo: str = "aes-gcm"
+) -> elt.TensorType:
+    """Abstract evaluation for symmetric encryption.
+
+    Args:
+        key: Symmetric encryption key
+        plaintext: Data to encrypt
+        algo: Encryption algorithm (keyword-only, validated at runtime)
+
+    Returns:
+        Ciphertext as dynamic-length uint8 tensor
+    """
     # Dynamic shape for ciphertext
+    # algo validation is done at backend impl, not here
     return elt.TensorType(elt.u8, (-1,))
 
 
 @sym_decrypt_p.def_abstract_eval
 def _sym_decrypt_ae(
-    key: elt.BaseType, ciphertext: elt.BaseType, target_type: elt.BaseType
+    key: elt.BaseType,
+    ciphertext: elt.BaseType,
+    *,
+    target_type: elt.BaseType,
+    algo: str = "aes-gcm",
 ) -> elt.BaseType:
+    """Abstract evaluation for symmetric decryption.
+
+    Args:
+        key: Symmetric decryption key
+        ciphertext: Encrypted data
+        target_type: Expected type of decrypted plaintext (keyword-only)
+        algo: Decryption algorithm (keyword-only, validated at runtime)
+
+    Returns:
+        Decrypted plaintext with type matching target_type
+    """
+    # algo validation is done at backend impl, not here
     return target_type
 
 
@@ -367,16 +396,50 @@ def hash_batch(data: el.Object) -> el.Object:
     return hash_batch_p.bind(data)
 
 
-def sym_encrypt(key: el.Object, plaintext: el.Object) -> el.Object:
-    """Symmetric encrypt (XOR stream or AES-GCM)."""
-    return sym_encrypt_p.bind(key, plaintext)
+def sym_encrypt(
+    key: el.Object, plaintext: el.Object, *, algo: str = "aes-gcm"
+) -> el.Object:
+    """Symmetric encrypt.
+
+    Args:
+        key: Symmetric encryption key (SymmetricKeyType or bytes).
+        plaintext: Data to encrypt (any serializable object).
+        algo: Encryption algorithm. Currently only "aes-gcm" is supported.
+              Validation is performed at backend execution time.
+
+    Returns:
+        Ciphertext as Tensor[u8, (-1,)].
+
+    Raises:
+        ValueError: At runtime if algo is not supported by the backend.
+    """
+    return sym_encrypt_p.bind(key, plaintext, algo=algo)
 
 
 def sym_decrypt(
-    key: el.Object, ciphertext: el.Object, target_type: elt.BaseType
+    key: el.Object,
+    ciphertext: el.Object,
+    target_type: elt.BaseType,
+    *,
+    algo: str = "aes-gcm",
 ) -> el.Object:
-    """Symmetric decrypt."""
-    return sym_decrypt_p.bind(key, ciphertext, target_type=target_type)
+    """Symmetric decrypt.
+
+    Args:
+        key: Symmetric decryption key (SymmetricKeyType or bytes).
+        ciphertext: Encrypted data.
+        target_type: Expected type of the decrypted plaintext (for type inference).
+        algo: Decryption algorithm. Must match the algorithm used for encryption.
+              Currently only "aes-gcm" is supported.
+              Validation is performed at backend execution time.
+
+    Returns:
+        Decrypted plaintext with type matching target_type.
+
+    Raises:
+        ValueError: At runtime if algo is not supported by the backend.
+    """
+    return sym_decrypt_p.bind(key, ciphertext, target_type=target_type, algo=algo)
 
 
 def select(cond: el.Object, true_val: el.Object, false_val: el.Object) -> el.Object:

{mplang_nightly-0.1.dev261 → mplang_nightly-0.1.dev262}/tests/v2/backends/test_crypto_impl.py

@@ -174,6 +174,44 @@ class TestSymmetricEncryption:
             # Ciphertexts should be different due to random nonce
             assert _unwrap(ct1.runtime_obj) != _unwrap(ct2.runtime_obj)
 
+    def test_encrypt_decrypt_keyword_only_params(self):
+        """Test that algo parameter is properly passed as keyword-only."""
+        with Interpreter():
+            sk, pk = crypto.kem_keygen("x25519")
+            key = crypto.kem_derive(sk, pk)
+
+            message = np.array([1, 2, 3, 4, 5], dtype=np.uint8)
+            message_obj = tensor.constant(message)
+
+            # Test keyword-only algo parameter in encrypt
+            ciphertext = crypto.sym_encrypt(key, message_obj, algo="aes-gcm")
+            assert isinstance(ciphertext.runtime_obj, BytesValue)
+
+            # Test keyword-only parameters in decrypt (target_type and algo)
+            plaintext = crypto.sym_decrypt(
+                key,
+                ciphertext,
+                target_type=elt.TensorType(elt.u8, (5,)),
+                algo="aes-gcm",
+            )
+            np.testing.assert_array_equal(_unwrap(plaintext.runtime_obj), message)
+
+    def test_encrypt_decrypt_default_algo(self):
+        """Test that default algo parameter works correctly."""
+        with Interpreter():
+            sk, pk = crypto.kem_keygen("x25519")
+            key = crypto.kem_derive(sk, pk)
+
+            message = np.array([1, 2, 3, 4, 5], dtype=np.uint8)
+            message_obj = tensor.constant(message)
+
+            # Should work without explicit algo (defaults to "aes-gcm")
+            ciphertext = crypto.sym_encrypt(key, message_obj)
+            plaintext = crypto.sym_decrypt(
+                key, ciphertext, target_type=elt.TensorType(elt.u8, (5,))
+            )
+            np.testing.assert_array_equal(_unwrap(plaintext.runtime_obj), message)
+
 
 class TestDigitalEnvelope:
     """Test complete digital envelope workflow."""
@@ -236,3 +274,67 @@ class TestDigitalEnvelope:
             ct_b = crypto.sym_encrypt(bob_key, msg_b_obj)
             pt_b = crypto.sym_decrypt(alice_key, ct_b, elt.TensorType(elt.u8, (3,)))
             np.testing.assert_array_equal(_unwrap(pt_b.runtime_obj), msg_b)
+
+
+class TestAlgoParameter:
+    """Test algo parameter validation in sym_encrypt/sym_decrypt."""
+
+    def setup_method(self):
+        self.interpreter = Interpreter()
+        self.interpreter.__enter__()
+        sk, pk = crypto.kem_keygen("x25519")
+        self.key = crypto.kem_derive(sk, pk)
+
+    def teardown_method(self):
+        self.interpreter.__exit__(None, None, None)
+
+    def test_default_algo_works(self):
+        """Test default algo='aes-gcm' works without explicit parameter."""
+
+        message = np.array([1, 2, 3], dtype=np.uint8)
+        message_obj = tensor.constant(message)
+
+        # Use default algo (should be aes-gcm)
+        ciphertext = crypto.sym_encrypt(self.key, message_obj)
+        plaintext = crypto.sym_decrypt(
+            self.key, ciphertext, elt.TensorType(elt.u8, (3,))
+        )
+        np.testing.assert_array_equal(_unwrap(plaintext.runtime_obj), message)
+
+    def test_explicit_aes_gcm_algo(self):
+        """Test explicitly passing algo='aes-gcm' works."""
+
+        message = np.array([4, 5, 6], dtype=np.uint8)
+        message_obj = tensor.constant(message)
+
+        # Explicitly pass algo='aes-gcm'
+        ciphertext = crypto.sym_encrypt(self.key, message_obj, algo="aes-gcm")
+        plaintext = crypto.sym_decrypt(
+            self.key, ciphertext, elt.TensorType(elt.u8, (3,)), algo="aes-gcm"
+        )
+        np.testing.assert_array_equal(_unwrap(plaintext.runtime_obj), message)
+
+    def test_unsupported_algo_encrypt_fails(self):
+        """Test that unsupported algo raises ValueError during encryption."""
+
+        message = np.array([7, 8, 9], dtype=np.uint8)
+        message_obj = tensor.constant(message)
+
+        # Try unsupported algorithm
+        with pytest.raises(ValueError, match="Unsupported encryption algorithm"):
+            crypto.sym_encrypt(self.key, message_obj, algo="aes-ctr")
+
+    def test_unsupported_algo_decrypt_fails(self):
+        """Test that unsupported algo raises ValueError during decryption."""
+
+        message = np.array([10, 11, 12], dtype=np.uint8)
+        message_obj = tensor.constant(message)
+
+        # Encrypt with default algo
+        ciphertext = crypto.sym_encrypt(self.key, message_obj)
+
+        # Try to decrypt with unsupported algorithm
+        with pytest.raises(ValueError, match="Unsupported decryption algorithm"):
+            crypto.sym_decrypt(
+                self.key, ciphertext, elt.TensorType(elt.u8, (3,)), algo="sm4-gcm"
+            )