mplang-nightly 0.1.dev265__py3-none-any.whl → 0.1.dev266__py3-none-any.whl

mplang/v2/backends/crypto_impl.py

@@ -614,6 +614,101 @@ def kem_derive_impl(
         return SymmetricKeyValue(suite=suite, key_bytes=secret)
 
 
+@crypto.hkdf_p.def_impl
+def hkdf_impl(
+    interpreter: Interpreter,
+    op: Operation,
+    secret: SymmetricKeyValue | TensorValue,
+) -> SymmetricKeyValue:
+    """HKDF key derivation implementation using SHA-256.
+
+    Implements RFC 5869 HKDF with HMAC-SHA256. This is the NIST SP 800-56C
+    compliant way to derive symmetric keys from ECDH shared secrets.
+
+    Current implementation supports only SHA-256. Future versions will add
+    SHA-512, SHA3-256, and BLAKE2b support.
+
+    Security Notes:
+        - Uses salt=None (defaults to 32-byte all-zero salt per RFC 5869)
+        - ONLY SAFE for high-entropy IKM (e.g., 256-bit ECDH shared secrets)
+        - NOT suitable for: passwords, low-entropy secrets, or repeated key derivations
+        - For session keys with same ECDH pair: use unique 'info' per session
+
+        Per NIST SP 800-56C Rev. 2:
+        "If the IKM is already cryptographically strong (e.g., from ECDH),
+         a salt may not be necessary, but using one does not hurt."
+
+    Args:
+        interpreter: Runtime interpreter context
+        op: Operation node containing attributes (info, hash_algo)
+        secret: Input key material (IKM) as SymmetricKeyValue or TensorValue
+                Must be high-entropy (≥256 bits) for security with salt=None
+
+    Returns:
+        SymmetricKeyValue with suite="hkdf-{hash_algo}" and 32-byte key_bytes
+
+    Raises:
+        TypeError: If secret is not SymmetricKeyValue or TensorValue
+        ValueError: If info parameter is empty (required for domain separation)
+        NotImplementedError: If hash_algo is not "sha256"
+    """
+    from cryptography.hazmat.primitives import hashes
+    from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+
+    # Extract operation attributes
+    info_str = op.attrs.get("info", "")
+    hash_algo = (
+        op.attrs.get("hash_algo", "sha256").lower().replace("-", "").replace("_", "")
+    )
+
+    # Validate info parameter (REQUIRED for domain separation per NIST)
+    if not info_str:
+        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'"
+        )
+
+    info_bytes = info_str.encode("utf-8")
+
+    # Extract input key material (IKM) bytes
+    if isinstance(secret, SymmetricKeyValue):
+        ikm = secret.key_bytes
+    elif isinstance(secret, TensorValue):
+        ikm = secret.unwrap().tobytes()
+    else:
+        raise TypeError(
+            f"hkdf secret must be SymmetricKeyValue or TensorValue, "
+            f"got {type(secret).__name__}"
+        )
+
+    # Validate hash algorithm (currently only SHA-256 implemented)
+    if hash_algo != "sha256":
+        raise NotImplementedError(
+            f"HKDF with hash algorithm '{hash_algo}' is not yet implemented. "
+            f"Currently only 'sha256' is supported. "
+            f"Planned future support: sha512, sha3256, blake2b"
+        )
+
+    # Perform HKDF derivation using cryptography library
+    # Note: salt=None uses 32-byte all-zero salt (not random salt!)
+    # This is secure ONLY because ECDH outputs are already high-entropy (256-bit uniform)
+    # For low-entropy inputs or repeated derivations, a random salt would be required
+    hkdf = HKDF(
+        algorithm=hashes.SHA256(),
+        length=32,  # Output length in bytes (AES-256 key = 32 bytes)
+        salt=None,  # 32-byte zero salt (secure for high-entropy ECDH shared secrets)
+        info=info_bytes,  # Context-specific binding for domain separation
+    )
+
+    derived_key = hkdf.derive(ikm)
+
+    # Return SymmetricKeyValue with composite suite name
+    # Format: "hkdf-{hash_algo}" to indicate derivation method and hash function
+    suite = f"hkdf-{hash_algo}"
+    return SymmetricKeyValue(suite=suite, key_bytes=derived_key)
+
+
 @crypto.random_bytes_p.def_impl
 def random_bytes_impl(interpreter: Interpreter, op: Operation) -> TensorValue:
     """Generate random bytes using os.urandom."""

mplang/v2/dialects/crypto.py

@@ -199,6 +199,9 @@ select_p = el.Primitive[el.Object]("crypto.select")
 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")
 
@@ -336,6 +339,41 @@ def _kem_derive_ae(
     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,))
@@ -472,6 +510,69 @@ def kem_derive(private_key: el.Object, public_key: el.Object) -> el.Object:
     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.
 

mplang/v2/libs/device/api.py

@@ -61,6 +61,9 @@ DEVICE_ATTR_NAME = "__device__"
 # Default KEM suite for TEE session establishment
 _TEE_KEM_SUITE: str = "x25519"
 
+# HKDF info string for TEE session key derivation (domain separation)
+_TEE_HKDF_INFO: str = "mplang/device/tee/v2"
+
 # Global cache for TEE sessions (keyed by (frm_dev_id, to_dev_id))
 # Each entry is (context_id, sess_frm, sess_tee) where context_id ensures
 # sessions are not reused across different trace/interp contexts.
@@ -412,16 +415,24 @@ def _ensure_tee_session(
     """Ensure a TEE session (sess_frm at sender, sess_tee at TEE) exists.
 
     Performs remote attestation and establishes an encrypted channel between
-    a PPU and a TEE device. The session keys are cached to avoid repeated
-    handshakes within the same execution context.
+    a PPU and a TEE device using NIST SP 800-56C compliant key derivation.
+    Session keys are cached within the same execution context to avoid
+    repeated handshakes.
 
-    The protocol:
+    Protocol (ECDH + Remote Attestation + HKDF):
     1. TEE generates keypair (sk, pk) and creates attestation quote binding pk
-    2. Quote is sent to the sender (PPU) for verification
+    2. Quote is sent to sender (PPU) for verification
     3. Sender verifies quote and extracts TEE's attested public key
-    4. Sender generates its own ephemeral keypair and sends pk to TEE
-    5. Both sides derive shared secret using ECDH (X25519)
-    6. The shared secret is used directly as the session key for AES-GCM
+    4. Sender generates ephemeral keypair and sends pk to TEE
+    5. Both sides derive ECDH shared secret using X25519
+    6. Both sides derive session keys from shared secret using HKDF-SHA256
+       with protocol-specific info string for domain separation
+
+    Security properties:
+    - Remote attestation: TEE identity is cryptographically verified
+    - Ephemeral keys: Perfect forward secrecy (keys not reused across sessions)
+    - HKDF derivation: NIST SP 800-56C compliant (shared secret not used directly)
+    - Domain separation: Info parameter binds keys to TEE protocol v2
 
     Args:
         frm_dev_id: Source device ID (PPU)
@@ -462,10 +473,18 @@ def _ensure_tee_session(
     v_sk, v_pk = simp.pcall_static((frm_rank,), crypto.kem_keygen, _TEE_KEM_SUITE)
     v_pk_at_tee = simp.shuffle_static(v_pk, {tee_rank: frm_rank})
 
-    # 5. Both sides derive shared secret (symmetric key) via ECDH
-    # The shared secret from X25519 ECDH is suitable for direct use as AES key
-    sess_frm = simp.pcall_static((frm_rank,), crypto.kem_derive, v_sk, tee_pk_at_sender)
-    sess_tee = simp.pcall_static((tee_rank,), crypto.kem_derive, tee_sk, v_pk_at_tee)
+    # 5. Both sides derive ECDH shared secret using X25519
+    # Note: kem_derive signature is (private_key, public_key) - suite is in key type
+    shared_frm = simp.pcall_static(
+        (frm_rank,), crypto.kem_derive, v_sk, tee_pk_at_sender
+    )
+    shared_tee = simp.pcall_static((tee_rank,), crypto.kem_derive, tee_sk, v_pk_at_tee)
+
+    # 6. Derive session keys using HKDF-SHA256 for domain separation
+    # Per NIST SP 800-56C: "shared secret SHALL NOT be used directly as a key"
+    # HKDF provides: uniform distribution + protocol-specific context binding
+    sess_frm = simp.pcall_static((frm_rank,), crypto.hkdf, shared_frm, _TEE_HKDF_INFO)
+    sess_tee = simp.pcall_static((tee_rank,), crypto.hkdf, shared_tee, _TEE_HKDF_INFO)
 
     # Cache the session
     _tee_session_cache[key] = (current_context_id, sess_frm, sess_tee)

{mplang_nightly-0.1.dev265.dist-info → mplang_nightly-0.1.dev266.dist-info}/METADATA

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

{mplang_nightly-0.1.dev265.dist-info → mplang_nightly-0.1.dev266.dist-info}/RECORD

@@ -82,7 +82,7 @@ mplang/v2/cli.py,sha256=QtiTFG418k26opRy4GhVV8fwFqRS11xTLH3xRCIIm6M,19665
 mplang/v2/cli_guide.md,sha256=kyoCaqkvIJJ1vsvCyBu3qgOuRSb0txu9BDZoy9GU5S0,3617
 mplang/v2/backends/__init__.py,sha256=H-4-jBEPWBZl6XT7AxBShRINnruF_f_2lB4iaiQoXME,1988
 mplang/v2/backends/bfv_impl.py,sha256=cQPinze3c2xN4CmIIoXxZoIEhu9ynoGaXbdF95z_aTE,25709
-mplang/v2/backends/crypto_impl.py,sha256=-4-ry7rXgPjkuM-KTFc4pG3k-O5dSNTEkcMR4W2LtJg,19327
+mplang/v2/backends/crypto_impl.py,sha256=tU0KdI34hnYvYujKUkiA1XYpvy4lo_MKpidt0NSIKlk,23205
 mplang/v2/backends/field_impl.py,sha256=50sKGOlkUiaTj_IAola86uQeoi-fxV0o7G91BdTCWZA,14788
 mplang/v2/backends/func_impl.py,sha256=R0662cC0gSSfkjuLyevJ_g4bJDJirY76LTFYqEimCkE,3585
 mplang/v2/backends/phe_impl.py,sha256=r836e_qBHGrHhfnFail5IaUDzvS7bABjdEQmJmAtBVI,4127
@@ -106,7 +106,7 @@ mplang/v2/backends/simp_worker/ops.py,sha256=DMQCsKeoMtemy5ozsVZt2eoF8NZlhLeHZMD
 mplang/v2/backends/simp_worker/state.py,sha256=eRUI7MP6gU8KPC9-H5fwcoAPKOsfW2ODWvpoKWbecMk,1554
 mplang/v2/dialects/__init__.py,sha256=hvzAvz6_brfFyDGgKknoPdgh5EY033YNYwotuJK_zoA,1493
 mplang/v2/dialects/bfv.py,sha256=XrE3FX9DHWqNzUVzY0tuwPvNVVRZYpD51JZIZF-q-l4,22350
-mplang/v2/dialects/crypto.py,sha256=rgJh_VO971MR-0D5jryz8osItF0Q4aGYhlEYjV_hBEE,17503
+mplang/v2/dialects/crypto.py,sha256=dH_DtoE3pGAKeOLPHxeyGtXC-nGwBsOs62TKikJEaq0,22197
 mplang/v2/dialects/dtypes.py,sha256=bGM3Jna3BnvE4MPOurWrEmQegGPxd26z1HIWox1rj0U,12104
 mplang/v2/dialects/field.py,sha256=6nBJg08k5WHb2o5msr8XAnxMQLpoTej55VQ7iSRnC4o,6380
 mplang/v2/dialects/func.py,sha256=UlaMof4NEG28VOtiRL7zBRYgFbIX74YTqqgvozbils0,4375
@@ -138,7 +138,7 @@ mplang/v2/kernels/okvs_opt.cpp,sha256=d_HhvMdcebYsG2x7kYzjuFgmEsh9WKLH6SHee3375B
 mplang/v2/kernels/py_kernels.py,sha256=FDsD86IHV-UBzxZLolhSOkrp24PuboHXeb1gBHLOfMo,12073
 mplang/v2/libs/collective.py,sha256=pfXq9tmFUNKjeHhWMTjtzOi-m2Fn1lLru1G6txZVyic,10683
 mplang/v2/libs/device/__init__.py,sha256=mXsSvXrWmlHu6Ch87Vcd85m4L_qdDkbSvJyHyuai2fc,1251
-mplang/v2/libs/device/api.py,sha256=8i0KP3q_XbOfPYB9fzs7pA6vVY4Ib6lxdoiwAi_RlmU,27838
+mplang/v2/libs/device/api.py,sha256=d_Wbka8bxxXsRMW6zDhjzL9LPtChSk2-ryfi-c4Mqsk,28830
 mplang/v2/libs/device/cluster.py,sha256=YUqYZ_IBS6rpV5ejUFP3kTxcTQHSyeDeuaJcsiFY_Js,12508
 mplang/v2/libs/ml/__init__.py,sha256=xTxhC_YwHP32muUEFCEwOjc-3Ml-vmO48NNECU90zm4,787
 mplang/v2/libs/ml/sgb.py,sha256=qoJww01EEbWo9nvwB5w-JZqx9BchcmhSrToypxV7fPg,60297
@@ -170,8 +170,8 @@ mplang/v2/runtime/dialect_state.py,sha256=HxO1i4kSOujS2tQzAF9-WmI3nChSaGgupf2_07
 mplang/v2/runtime/interpreter.py,sha256=UzrM5oepka6H0YKRZncNXhsuwKVm4pliG5J92fFRZMI,32300
 mplang/v2/runtime/object_store.py,sha256=yT6jtKG2GUEJVmpq3gnQ8mCMvUFYzgBciC5A-J5KRdk,5998
 mplang/v2/runtime/value.py,sha256=CMOxElJP78v7pjasPhEpbxWbSgB2KsLbpPmzz0mQX0E,4317
-mplang_nightly-0.1.dev265.dist-info/METADATA,sha256=Dr04ggCRHYxf8ck9AcrmmgPfHhBmrvDwUfWHXjvTS-w,16775
-mplang_nightly-0.1.dev265.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-mplang_nightly-0.1.dev265.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
-mplang_nightly-0.1.dev265.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mplang_nightly-0.1.dev265.dist-info/RECORD,,
+mplang_nightly-0.1.dev266.dist-info/METADATA,sha256=y4CgWHDl3wbczPz2lcTxnqXsTGFkPcuk3bGcnIQoAV0,16775
+mplang_nightly-0.1.dev266.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+mplang_nightly-0.1.dev266.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
+mplang_nightly-0.1.dev266.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mplang_nightly-0.1.dev266.dist-info/RECORD,,