mplang-nightly 0.1.dev192__py3-none-any.whl → 0.1.dev268__py3-none-any.whl

mplang/v2/dialects/crypto.py

@@ -0,0 +1,689 @@
+# Copyright 2025 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Crypto dialect for the EDSL.
+
+Provides cryptographic primitives including ECC, Hashing, and Symmetric Encryption.
+"""
+
+from __future__ import annotations
+
+from typing import Any, ClassVar
+
+import mplang.v2.edsl as el
+import mplang.v2.edsl.typing as elt
+from mplang.v2.edsl import serde
+
+# ==============================================================================
+# --- Type Definitions
+# ==============================================================================
+
+
+@serde.register_class
+class PointType(elt.BaseType):
+    """Type for an ECC Point."""
+
+    def __init__(self, curve: str = "secp256k1"):
+        self.curve = curve
+
+    def __str__(self) -> str:
+        return f"Point[{self.curve}]"
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, PointType):
+            return False
+        return self.curve == other.curve
+
+    def __hash__(self) -> int:
+        return hash(("PointType", self.curve))
+
+    # --- Serde methods ---
+    _serde_kind: ClassVar[str] = "crypto.PointType"
+
+    def to_json(self) -> dict[str, Any]:
+        return {"curve": self.curve}
+
+    @classmethod
+    def from_json(cls, data: dict[str, Any]) -> PointType:
+        return cls(curve=data["curve"])
+
+
+@serde.register_class
+class ScalarType(elt.BaseType):
+    """Type for an ECC Scalar (integer modulo curve order)."""
+
+    def __init__(self, curve: str = "secp256k1"):
+        self.curve = curve
+
+    def __str__(self) -> str:
+        return f"Scalar[{self.curve}]"
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, ScalarType):
+            return False
+        return self.curve == other.curve
+
+    def __hash__(self) -> int:
+        return hash(("ScalarType", self.curve))
+
+    # --- Serde methods ---
+    _serde_kind: ClassVar[str] = "crypto.ScalarType"
+
+    def to_json(self) -> dict[str, Any]:
+        return {"curve": self.curve}
+
+    @classmethod
+    def from_json(cls, data: dict[str, Any]) -> ScalarType:
+        return cls(curve=data["curve"])
+
+
+@serde.register_class
+class PrivateKeyType(elt.BaseType):
+    """Type for a KEM private key."""
+
+    def __init__(self, suite: str = "x25519"):
+        self.suite = suite
+
+    def __str__(self) -> str:
+        return f"PrivateKey[{self.suite}]"
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, PrivateKeyType):
+            return False
+        return self.suite == other.suite
+
+    def __hash__(self) -> int:
+        return hash(("PrivateKeyType", self.suite))
+
+    # --- Serde methods ---
+    _serde_kind: ClassVar[str] = "crypto.PrivateKeyType"
+
+    def to_json(self) -> dict[str, Any]:
+        return {"suite": self.suite}
+
+    @classmethod
+    def from_json(cls, data: dict[str, Any]) -> PrivateKeyType:
+        return cls(suite=data["suite"])
+
+
+@serde.register_class
+class PublicKeyType(elt.BaseType):
+    """Type for a KEM public key."""
+
+    def __init__(self, suite: str = "x25519"):
+        self.suite = suite
+
+    def __str__(self) -> str:
+        return f"PublicKey[{self.suite}]"
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, PublicKeyType):
+            return False
+        return self.suite == other.suite
+
+    def __hash__(self) -> int:
+        return hash(("PublicKeyType", self.suite))
+
+    # --- Serde methods ---
+    _serde_kind: ClassVar[str] = "crypto.PublicKeyType"
+
+    def to_json(self) -> dict[str, Any]:
+        return {"suite": self.suite}
+
+    @classmethod
+    def from_json(cls, data: dict[str, Any]) -> PublicKeyType:
+        return cls(suite=data["suite"])
+
+
+@serde.register_class
+class SymmetricKeyType(elt.BaseType):
+    """Type for a symmetric encryption key (e.g., from KEM derive)."""
+
+    def __init__(self, suite: str = "x25519"):
+        self.suite = suite
+
+    def __str__(self) -> str:
+        return f"SymmetricKey[{self.suite}]"
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, SymmetricKeyType):
+            return False
+        return self.suite == other.suite
+
+    def __hash__(self) -> int:
+        return hash(("SymmetricKeyType", self.suite))
+
+    # --- Serde methods ---
+    _serde_kind: ClassVar[str] = "crypto.SymmetricKeyType"
+
+    def to_json(self) -> dict[str, Any]:
+        return {"suite": self.suite}
+
+    @classmethod
+    def from_json(cls, data: dict[str, Any]) -> SymmetricKeyType:
+        return cls(suite=data["suite"])
+
+
+# ==============================================================================
+# --- Primitives
+# ==============================================================================
+
+# ECC
+generator_p = el.Primitive[el.Object]("crypto.ec_generator")
+mul_p = el.Primitive[el.Object]("crypto.ec_mul")
+add_p = el.Primitive[el.Object]("crypto.ec_add")
+sub_p = el.Primitive[el.Object]("crypto.ec_sub")
+point_to_bytes_p = el.Primitive[el.Object]("crypto.ec_point_to_bytes")
+random_scalar_p = el.Primitive[el.Object]("crypto.ec_random_scalar")
+scalar_from_int_p = el.Primitive[el.Object]("crypto.ec_scalar_from_int")
+
+# Symmetric / Hash
+hash_p = el.Primitive[el.Object]("crypto.hash")
+hash_batch_p = el.Primitive[el.Object]("crypto.hash_batch")
+sym_encrypt_p = el.Primitive[el.Object]("crypto.sym_encrypt")
+sym_decrypt_p = el.Primitive[el.Object]("crypto.sym_decrypt")
+select_p = el.Primitive[el.Object]("crypto.select")
+
+# KEM (Key Encapsulation Mechanism)
+kem_keygen_p = el.Primitive[tuple[el.Object, el.Object]]("crypto.kem_keygen")
+kem_derive_p = el.Primitive[el.Object]("crypto.kem_derive")
+
+# HKDF (Key Derivation Function)
+hkdf_p = el.Primitive[el.Object]("crypto.hkdf")
+
+# Randomness
+random_bytes_p = el.Primitive[el.Object]("crypto.random_bytes")
+
+
+# ==============================================================================
+# --- Abstract Evaluation (Type Inference)
+# ==============================================================================
+
+
+@generator_p.def_abstract_eval
+def _generator_ae(curve: str = "secp256k1") -> PointType:
+    return PointType(curve)
+
+
+@mul_p.def_abstract_eval
+def _mul_ae(point: PointType, scalar: ScalarType) -> PointType:
+    return PointType(point.curve)
+
+
+@add_p.def_abstract_eval
+def _add_ae(p1: PointType, p2: PointType) -> PointType:
+    return PointType(p1.curve)
+
+
+@sub_p.def_abstract_eval
+def _sub_ae(p1: PointType, p2: PointType) -> PointType:
+    return PointType(p1.curve)
+
+
+@point_to_bytes_p.def_abstract_eval
+def _pt_to_bytes_ae(point: elt.BaseType) -> elt.TensorType:
+    if isinstance(point, elt.TensorType):
+        # Vectorized behavior: Tensor[Point, shape] -> Tensor[u8, shape + (65,)]
+        return elt.TensorType(elt.u8, (*point.shape, 65))
+    return elt.TensorType(elt.u8, (65,))
+
+
+@random_scalar_p.def_abstract_eval
+def _random_scalar_ae(curve: str = "secp256k1") -> ScalarType:
+    return ScalarType(curve)
+
+
+@scalar_from_int_p.def_abstract_eval
+def _scalar_from_int_ae(
+    val: elt.TensorType | elt.IntegerType, curve: str = "secp256k1"
+) -> ScalarType:
+    return ScalarType(curve)
+
+
+@hash_p.def_abstract_eval
+def _hash_ae(data: elt.BaseType) -> elt.TensorType:
+    # Strictly single output (blob hash)
+    return elt.TensorType(elt.u8, (32,))
+
+
+@hash_batch_p.def_abstract_eval
+def _hash_batch_ae(data: elt.BaseType) -> elt.TensorType:
+    # Explicit batch hashing: Input (..., D) -> Output (..., 32)
+    # Hashes the last dimension D bytes.
+    if not isinstance(data, elt.TensorType):
+        raise TypeError(f"hash_batch requires TensorType, got {data}")
+
+    # data.shape is tuple[int | None, ...]
+    shape = data.shape
+    if len(shape) < 2:
+        # Fallback/Edge case: (D,) -> (32,)
+        # One could argue this should be an error for *batch* primitive,
+        # but allowing it provides consistency for (N=1, D).
+        return elt.TensorType(elt.u8, (32,))
+
+    # Batch shape is everything except last dim
+    batch_shape = shape[:-1]
+    return elt.TensorType(elt.u8, (*batch_shape, 32))
+
+
+@sym_encrypt_p.def_abstract_eval
+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,
+    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
+
+
+@select_p.def_abstract_eval
+def _select_ae(
+    cond: elt.BaseType, true_val: elt.BaseType, false_val: elt.BaseType
+) -> elt.BaseType:
+    return true_val
+
+
+@kem_keygen_p.def_abstract_eval
+def _kem_keygen_ae(suite: str = "x25519") -> tuple[PrivateKeyType, PublicKeyType]:
+    return (PrivateKeyType(suite), PublicKeyType(suite))
+
+
+@kem_derive_p.def_abstract_eval
+def _kem_derive_ae(
+    private_key: PrivateKeyType, public_key: PublicKeyType
+) -> SymmetricKeyType:
+    suite = getattr(private_key, "suite", "x25519")
+    return SymmetricKeyType(suite)
+
+
+@hkdf_p.def_abstract_eval
+def _hkdf_ae(
+    secret: elt.BaseType, *, info: str, hash_algo: str = "sha256"
+) -> SymmetricKeyType:
+    """Abstract evaluation for HKDF key derivation.
+
+    Args:
+        secret: Input key material (SymmetricKeyType from kem_derive or TensorType[u8])
+        info: Context string for domain separation (required, non-empty, keyword-only)
+        hash_algo: Hash algorithm in lowercase without hyphens (e.g., "sha256", keyword-only)
+
+    Returns:
+        SymmetricKeyType with suite="hkdf-{hash_algo}"
+
+    Raises:
+        TypeError: If info or hash_algo is not a string
+        ValueError: If info is empty (required for domain separation per NIST)
+    """
+    # Validate info and hash_algo at trace time
+    if not isinstance(info, str) or not info:
+        raise ValueError(
+            "HKDF requires non-empty 'info' parameter for domain separation. "
+            "The info string binds the derived key to a specific protocol/context. "
+            "Recommended format: 'namespace/component/purpose/version'"
+        )
+    if not isinstance(hash_algo, str) or not hash_algo:
+        raise TypeError("hash_algo must be a non-empty string")
+
+    # Normalize: lowercase, no hyphens
+    hash_algo_normalized = hash_algo.lower().replace("-", "").replace("_", "")
+
+    # Return SymmetricKeyType with composite suite indicating derivation method
+    return SymmetricKeyType(suite=f"hkdf-{hash_algo_normalized}")
+
+
+@random_bytes_p.def_abstract_eval
+def _random_bytes_ae(length: int) -> elt.TensorType:
+    return elt.TensorType(elt.u8, (length,))
+
+
+# ==============================================================================
+# --- Helper Functions (Ops)
+# ==============================================================================
+
+
+def ec_generator(curve: str = "secp256k1") -> el.Object:
+    """Get the generator point G for the curve."""
+    return generator_p.bind(curve=curve)
+
+
+def ec_mul(point: el.Object, scalar: el.Object) -> el.Object:
+    """Scalar multiplication: point * scalar."""
+    return mul_p.bind(point, scalar)
+
+
+def ec_add(p1: el.Object, p2: el.Object) -> el.Object:
+    """Point addition: p1 + p2."""
+    return add_p.bind(p1, p2)
+
+
+def ec_sub(p1: el.Object, p2: el.Object) -> el.Object:
+    """Point subtraction: p1 - p2."""
+    return sub_p.bind(p1, p2)
+
+
+def ec_point_to_bytes(point: el.Object) -> el.Object:
+    """Serialize point to bytes."""
+    return point_to_bytes_p.bind(point)
+
+
+def ec_random_scalar(curve: str = "secp256k1") -> el.Object:
+    """Generate a random scalar."""
+    return random_scalar_p.bind(curve=curve)
+
+
+def ec_scalar_from_int(val: el.Object, curve: str = "secp256k1") -> el.Object:
+    """Convert an integer tensor to a scalar."""
+    return scalar_from_int_p.bind(val, curve=curve)
+
+
+def hash_bytes(data: el.Object) -> el.Object:
+    """Hash bytes (SHA256). Returns 32-byte tensor."""
+    return hash_p.bind(data)
+
+
+def hash_batch(data: el.Object) -> el.Object:
+    """Hash each row of a tensor independently.
+
+    Treats the last dimension as the data to hash.
+    Input: (N, D) -> Output: (N, 32)
+    Input: (B, N, D) -> Output: (B, N, 32)
+    """
+    return hash_batch_p.bind(data)
+
+
+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,
+    *,
+    algo: str = "aes-gcm",
+) -> el.Object:
+    """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:
+    """Select between two values based on condition."""
+    return select_p.bind(cond, true_val, false_val)
+
+
+def kem_keygen(suite: str = "x25519") -> tuple[el.Object, el.Object]:
+    """Generate a KEM key pair (private_key, public_key).
+
+    Args:
+        suite: The KEM suite to use (e.g., "x25519", "kyber768")
+
+    Returns:
+        A tuple of (private_key, public_key)
+    """
+    return kem_keygen_p.bind(suite=suite)
+
+
+def kem_derive(private_key: el.Object, public_key: el.Object) -> el.Object:
+    """Derive a symmetric key from a private key and a public key (ECDH).
+
+    Args:
+        private_key: The local private key
+        public_key: The remote party's public key
+
+    Returns:
+        A symmetric key suitable for use with sym_encrypt/sym_decrypt
+    """
+    return kem_derive_p.bind(private_key, public_key)
+
+
+def hkdf(secret: el.Object, info: str, *, hash_algo: str = "sha256") -> el.Object:
+    """Derive a cryptographic key from input key material using HKDF.
+
+    HKDF (HMAC-based Key Derivation Function) is specified in RFC 5869 and
+    required by NIST SP 800-56C Rev.2 for deriving symmetric keys from
+    key agreement schemes like ECDH. Per NIST: "The shared secret output
+    from a key-agreement scheme SHALL NOT be used directly as a cryptographic
+    key. A key-derivation function (KDF) SHALL be used."
+
+    Args:
+        secret: Input key material (IKM). Accepts:
+                - SymmetricKeyValue: Typically from crypto.kem_derive (ECDH output)
+                - TensorType[u8, (N,)]: Raw bytes (N-byte secret)
+        info: Application-specific context string for domain separation.
+              REQUIRED and must be non-empty. Different info values produce
+              cryptographically independent keys even from the same secret.
+              Recommended format: "namespace/component/purpose/version"
+              Example: "mplang/device/tee/v2"
+        hash_algo: Hash function to use. Must be lowercase without hyphens.
+                   Currently supported: "sha256" (default)
+                   Future support planned: "sha512", "sha3256", "blake2b"
+                   Default "sha256" provides 128-bit security level.
+
+    Returns:
+        SymmetricKeyValue with:
+        - suite: "hkdf-{hash_algo}" (e.g., "hkdf-sha256")
+        - key_bytes: 32-byte derived key suitable for AES-256-GCM
+
+    Security considerations:
+        - Output length: Fixed at 32 bytes (256 bits) for AES-256 keys
+        - Salt: Uses salt=None (acceptable for ECDH output per NIST guidance)
+        - Info: Provides protocol/context binding (domain separation)
+        - Deterministic: Same (secret, info, hash_algo) always produces same key
+
+    Raises:
+        ValueError:
+            - At abstract evaluation time if hash_algo is unsupported.
+            - At execution time if info is empty.
+        NotImplementedError:
+            - At execution time if hash_algo is not "sha256".
+
+    Examples:
+        >>> # Standard TEE session establishment
+        >>> sk_local, pk_local = crypto.kem_keygen("x25519")
+        >>> sk_remote, pk_remote = crypto.kem_keygen("x25519")
+        >>> # ECDH on both sides
+        >>> shared_local = crypto.kem_derive(sk_local, pk_remote)
+        >>> shared_remote = crypto.kem_derive(sk_remote, pk_local)
+        >>> # HKDF for domain separation
+        >>> sess_local = crypto.hkdf(shared_local, "mplang/device/tee/v2")
+        >>> sess_remote = crypto.hkdf(shared_remote, "mplang/device/tee/v2")
+        >>> # sess_local and sess_remote have identical key_bytes
+        >>> # but suite="hkdf-sha256" (not "x25519")
+        >>>
+        >>> # Derive multiple independent keys from one master secret
+        >>> master_secret = crypto.kem_derive(sk, pk)
+        >>> encryption_key = crypto.hkdf(master_secret, "app/encryption/v1")
+        >>> mac_key = crypto.hkdf(master_secret, "app/mac/v1")
+        >>> # encryption_key ≠ mac_key due to different info strings
+    """
+    return hkdf_p.bind(secret, info=info, hash_algo=hash_algo)
+
+
+def random_bytes(length: int) -> el.Object:
+    """Generate cryptographically secure random bytes at runtime.
+
+    Args:
+        length: Number of bytes to generate.
+
+    Returns:
+        (length,) uint8 Tensor.
+    """
+    return random_bytes_p.bind(length=length)
+
+
+def random_tensor(shape: tuple[int, ...], dtype: elt.ScalarType) -> el.Object:
+    """Generate cryptographically secure random tensor at runtime.
+
+    This is a helper function that composes `random_bytes` with `tensor.run_jax`
+    to produce a tensor of the specified shape and dtype.
+
+    Args:
+        shape: Output tensor shape (e.g., (100,) or (10, 16)).
+        dtype: Element type (e.g., elt.u64, elt.i32, elt.f32).
+
+    Returns:
+        Tensor[dtype, shape] with CSPRNG values.
+
+    Example:
+        >>> # Generate 100 random uint64 values
+        >>> x = crypto.random_tensor((100,), elt.u64)
+        >>> # Generate 10x16 random int32 matrix
+        >>> y = crypto.random_tensor((10, 16), elt.i32)
+    """
+    import math
+    from typing import cast
+
+    from mplang.v2.dialects import dtypes, tensor
+
+    # Get byte size from numpy dtype
+    np_dtype = dtypes.to_numpy(dtype)
+    element_bytes = np_dtype.itemsize
+    total_elements = math.prod(shape)
+    total_bytes = total_elements * element_bytes
+
+    raw = random_bytes(total_bytes)
+
+    jax_dtype = dtypes.to_jax(dtype)
+
+    def _view_reshape(b: Any) -> Any:
+        return b.view(jax_dtype).reshape(shape)
+
+    return cast(el.Object, tensor.run_jax(_view_reshape, raw))
+
+
+def random_bits(n: int) -> el.Object:
+    """Generate n cryptographically secure random bits at runtime.
+
+    Each bit is stored as a uint8 with value 0 or 1 (unpacked representation).
+
+    Args:
+        n: Number of random bits to generate.
+
+    Returns:
+        (n,) uint8 Tensor with values 0 or 1.
+
+    Example:
+        >>> # Generate 1024 random bits for OT selection
+        >>> choice_bits = crypto.random_bits(1024)
+    """
+    from typing import cast
+
+    import jax.numpy as jnp
+
+    from mplang.v2.dialects import tensor
+
+    # Generate enough bytes to cover n bits
+    num_bytes = (n + 7) // 8
+    raw = random_bytes(num_bytes)
+
+    def _unpack_and_slice(b: Any, n: int = n) -> Any:
+        bits = jnp.unpackbits(b, bitorder="little")
+        return bits[:n]
+
+    return cast(el.Object, tensor.run_jax(_unpack_and_slice, raw))
+
+
+# --- Bytes <-> Point Conversions ---
+
+bytes_to_point_p = el.Primitive[el.Object]("crypto.ec_bytes_to_point")
+
+
+@bytes_to_point_p.def_abstract_eval
+def _bytes_to_point_ae(b: elt.TensorType) -> PointType:
+    return PointType("secp256k1")
+
+
+def ec_bytes_to_point(b: el.Object) -> el.Object:
+    """
+    Deserialize bytes to an ECC point.
+
+    Args:
+        b: A (65,) uint8 Tensor representing an uncompressed point in SEC1 format.
+           The first byte must be 0x04, followed by 32 bytes for X and 32 bytes for Y.
+
+    Returns:
+        An ECC point object corresponding to the input bytes.
+
+    Raises:
+        ValueError: If the input is not a valid 65-byte uncompressed point representation.
+
+    Example:
+        >>> # Example: Deserialize a point from bytes
+        >>> point_bytes = jnp.array([0x04] + [0x01] * 32 + [0x02] * 32, dtype=jnp.uint8)
+        >>> point = crypto.ec_bytes_to_point(point_bytes)
+    """
+    return bytes_to_point_p.bind(b)