mplang-nightly 0.1.dev156__tar.gz → 0.1.dev157__tar.gz

{mplang_nightly-0.1.dev156 → mplang_nightly-0.1.dev157}/PKG-INFO

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

{mplang_nightly-0.1.dev156 → mplang_nightly-0.1.dev157}/mplang/device.py

@@ -207,8 +207,15 @@ def _d2d(to_dev_id: str, obj: MPObject) -> MPObject:
         assert len(frm_dev.members) == 1 and len(to_dev.members) == 1
         frm_rank = frm_dev.members[0].rank
         tee_rank = to_dev.members[0].rank
+        platform = to_dev.config.get("platform")
+        if not platform:
+            raise ValueError(
+                f"TEE device '{to_dev_id}' is missing 'platform' in its config."
+            )
         # Ensure sessions (both directions) exist for this PPU<->TEE pair
-        sess_p, sess_t = _ensure_tee_session(frm_dev_id, to_dev_id, frm_rank, tee_rank)
+        sess_p, sess_t = _ensure_tee_session(
+            frm_dev_id, to_dev_id, frm_rank, tee_rank, platform
+        )
         # Bytes-only path: pack -> enc -> p2p -> dec -> unpack (with static out type)
         obj_ty = TensorType.from_obj(obj)
         b = simp.runAt(frm_rank, builtin.pack)(obj)
@@ -222,8 +229,15 @@ def _d2d(to_dev_id: str, obj: MPObject) -> MPObject:
         assert len(frm_dev.members) == 1 and len(to_dev.members) == 1
         tee_rank = frm_dev.members[0].rank
         ppu_rank = to_dev.members[0].rank
+        platform = to_dev.config.get("platform")
+        if not platform:
+            raise ValueError(
+                f"TEE device '{to_dev_id}' is missing 'platform' in its config."
+            )
         # Ensure bidirectional session established for this pair
-        sess_p, sess_t = _ensure_tee_session(to_dev_id, frm_dev_id, ppu_rank, tee_rank)
+        sess_p, sess_t = _ensure_tee_session(
+            to_dev_id, frm_dev_id, ppu_rank, tee_rank, platform
+        )
         obj_ty = TensorType.from_obj(obj)
         b = simp.runAt(tee_rank, builtin.pack)(obj)
         ct = simp.runAt(tee_rank, crypto.enc)(b, sess_t)
@@ -245,7 +259,7 @@ def _d2d(to_dev_id: str, obj: MPObject) -> MPObject:
 
 
 def _ensure_tee_session(
-    frm_dev_id: str, to_dev_id: str, frm_rank: int, tee_rank: int
+    frm_dev_id: str, to_dev_id: str, frm_rank: int, tee_rank: int, platform: str
 ) -> tuple[MPObject, MPObject]:
     """Ensure a TEE session (sess_p at sender, sess_t at TEE) exists.
 
@@ -263,11 +277,11 @@ def _ensure_tee_session(
     # 1) TEE generates (sk, pk) and quote(pk)
     # KEM suite currently constant; future: read from tee device config (e.g. cluster_spec.devices[dev_id].config)
     tee_sk, tee_pk = simp.runAt(tee_rank, crypto.kem_keygen)(_TEE_KEM_SUITE)
-    quote = simp.runAt(tee_rank, tee.quote)(tee_pk)
+    quote = simp.runAt(tee_rank, tee.quote_gen)(tee_pk)
 
     # 2) Send quote to sender and attest to obtain TEE pk
     quote_at_sender = mpi.p2p(tee_rank, frm_rank, quote)
-    tee_pk_at_sender = simp.runAt(frm_rank, tee.attest)(quote_at_sender)
+    tee_pk_at_sender = simp.runAt(frm_rank, tee.attest)(quote_at_sender, platform)
 
     # 3) Sender generates its ephemeral keypair and sends its pk to TEE
     v_sk, v_pk = simp.runAt(frm_rank, crypto.kem_keygen)(_TEE_KEM_SUITE)

{mplang_nightly-0.1.dev156 → mplang_nightly-0.1.dev157}/mplang/kernels/context.py

@@ -91,7 +91,7 @@ _DEFAULT_BINDINGS: dict[str, str] = {
     # generic SQL op; backend-specific kernel id for duckdb
     "sql.run": "duckdb.run_sql",
     # tee
-    # "tee.quote": "mock_tee.quote",
+    # "tee.quote_gen": "mock_tee.quote_gen",
     # "tee.attest": "mock_tee.attest",
 }
 

{mplang_nightly-0.1.dev156 → mplang_nightly-0.1.dev157}/mplang/kernels/mock_tee.py

@@ -45,10 +45,10 @@ def _quote_from_pk(pk: np.ndarray) -> NDArray[np.uint8]:
     return out
 
 
-@kernel_def("mock_tee.quote")
-def _tee_quote(pfunc: PFunction, pk: object) -> NDArray[np.uint8]:
+@kernel_def("mock_tee.quote_gen")
+def _tee_quote_gen(pfunc: PFunction, pk: object) -> NDArray[np.uint8]:
     warnings.warn(
-        "Insecure mock TEE kernel 'mock_tee.quote' in use. NOT secure; for local testing only.",
+        "Insecure mock TEE kernel 'mock_tee.quote_gen' in use. NOT secure; for local testing only.",
         stacklevel=3,
     )
     pk = np.asarray(pk, dtype=np.uint8)
@@ -64,6 +64,10 @@ def _tee_attest(pfunc: PFunction, quote: object) -> NDArray[np.uint8]:
         stacklevel=3,
     )
     quote = np.asarray(quote, dtype=np.uint8)
+    platform = pfunc.attrs.get("platform")
+    if platform is None:
+        raise ValueError("missing required 'platform' attribute in PFunction")
+
     if quote.size != 33:
         raise ValueError("mock quote must be 33 bytes (1 header + 32 pk)")
     return quote[1:33].astype(np.uint8)

{mplang_nightly-0.1.dev156 → mplang_nightly-0.1.dev157}/mplang/ops/tee.py

@@ -14,7 +14,11 @@
 
 from __future__ import annotations
 
+from jax.tree_util import PyTreeDef, tree_flatten
+
 from mplang.core.dtype import UINT8
+from mplang.core.mpobject import MPObject
+from mplang.core.pfunc import PFunction
 from mplang.core.tensor import TensorType
 from mplang.ops.base import stateless_mod
 
@@ -22,21 +26,26 @@ _TEE_MOD = stateless_mod("tee")
 
 
 @_TEE_MOD.simple_op()
-def quote(pk: TensorType) -> TensorType:
-    """TEE quote generation binding the provided ephemeral public key.
-
-    API (mock): quote(pk: u8[32]) -> (quote: u8[33])
-    The mock encodes a 1-byte header + 32-byte pk.
-    """
+def quote_gen(pk: TensorType) -> TensorType:
+    """TEE quote generation binding the provided ephemeral public key."""
     _ = pk  # Mark as used for the decorator
-    return TensorType(UINT8, (33,))
-
-
-@_TEE_MOD.simple_op()
-def attest(quote: TensorType) -> TensorType:
-    """TEE quote verification returning the attested TEE public key.
-
-    API (mock): attest(quote: u8[33]) -> tee_pk: u8[32]
-    """
-    _ = quote  # Mark as used for the decorator
-    return TensorType(UINT8, (32,))
+    return TensorType(UINT8, (-1,))
+
+
+@_TEE_MOD.op_def()
+def attest(
+    quote: MPObject, platform: str
+) -> tuple[PFunction, list[MPObject], PyTreeDef]:
+    """TEE quote verification returning the attested TEE public key."""
+
+    ins_info = [TensorType.from_obj(quote)]
+    outs_info = [TensorType(UINT8, (32,))]  # pk is always 32 bytes for x25519
+    pfunc = PFunction(
+        fn_type="tee.attest",
+        ins_info=ins_info,
+        outs_info=outs_info,
+        platform=platform,
+    )
+    _, treedef = tree_flatten(outs_info[0])
+
+    return pfunc, [quote], treedef

{mplang_nightly-0.1.dev156 → mplang_nightly-0.1.dev157}/tests/ops/test_crypto_tee.py

@@ -25,14 +25,14 @@ def _demo_flow():
     # TEE generates two ephemeral keypairs and quotes binding their pk
     t_sk0, t_pk0 = simp.runAt(P2, crypto.kem_keygen)("x25519")
     t_sk1, t_pk1 = simp.runAt(P2, crypto.kem_keygen)("x25519")
-    q0 = simp.runAt(P2, tee.quote)(t_pk0)
-    q1 = simp.runAt(P2, tee.quote)(t_pk1)
+    q0 = simp.runAt(P2, tee.quote_gen)(t_pk0)
+    q1 = simp.runAt(P2, tee.quote_gen)(t_pk1)
 
     # Send quotes to P0/P1 and attest to obtain TEE public keys
     q0_for_p0 = simp.p2p(P2, P0, q0)
     q1_for_p1 = simp.p2p(P2, P1, q1)
-    t_pk0_for_p0 = simp.runAt(P0, tee.attest)(q0_for_p0)
-    t_pk1_for_p1 = simp.runAt(P1, tee.attest)(q1_for_p1)
+    t_pk0_for_p0 = simp.runAt(P0, tee.attest)(q0_for_p0, "TDX")
+    t_pk1_for_p1 = simp.runAt(P1, tee.attest)(q1_for_p1, "TDX")
 
     # Each party generates its own ephemeral keypair and shares pk with TEE
     v_sk0, v_pk0 = simp.runAt(P0, crypto.kem_keygen)("x25519")
@@ -75,7 +75,10 @@ def _demo_flow():
 
 def test_crypto_enc_dec_and_tee_quote_attest_roundtrip():
     # Create simulator with TEE bindings using the new initial_bindings parameter
-    tee_bindings = {"tee.quote": "mock_tee.quote", "tee.attest": "mock_tee.attest"}
+    tee_bindings = {
+        "tee.quote_gen": "mock_tee.quote_gen",
+        "tee.attest": "mock_tee.attest",
+    }
     sim = mplang.Simulator.simple(3, op_bindings=tee_bindings)
     p0, p1 = mplang.evaluate(sim, _demo_flow)
     a = mplang.fetch(sim, p0)

{mplang_nightly-0.1.dev156 → mplang_nightly-0.1.dev157}/tests/simp/test_sugar.py

@@ -35,7 +35,7 @@ def test_basic_callable_and_namespace():
         a, b = simp.P0(crypto.kem_keygen, "x25519")
         # namespace form (tee side key, then quote)
         t_sk, t_pk = simp.P[2].crypto.kem_keygen("x25519")
-        _ = simp.P[2].tee.quote(t_pk)
+        _ = simp.P[2].tee.quote_gen(t_pk)
         # derive something simple at party 0 to ensure run path works
         _ = simp.P0(lambda x: x + 1, 41)
         return a, b, t_sk, t_pk

{mplang_nightly-0.1.dev156 → mplang_nightly-0.1.dev157}/tutorials/3_device.py

@@ -101,7 +101,10 @@ def run_tee():
     print("-" * 10, "millionaire (TEE)", "-" * 10)
 
     # TEE operations need explicit binding for security
-    tee_bindings = {"tee.quote": "mock_tee.quote", "tee.attest": "mock_tee.attest"}
+    tee_bindings = {
+        "tee.quote_gen": "mock_tee.quote_gen",
+        "tee.attest": "mock_tee.attest",
+    }
     # Apply tee bindings across nodes before constructing simulator
     for n in cluster_spec.nodes.values():
         n.runtime_info.op_bindings.update(tee_bindings)

{mplang_nightly-0.1.dev156 → mplang_nightly-0.1.dev157}/tutorials/9_tee.py

@@ -46,7 +46,11 @@ cluster_spec = ClusterSpec.from_dict({
         },
         "P0": {"kind": "PPU", "members": ["node_0"], "config": {}},
         "P1": {"kind": "PPU", "members": ["node_1"], "config": {}},
-        "TEE0": {"kind": "TEE", "members": ["node_2"], "config": {}},
+        "TEE0": {
+            "kind": "TEE",
+            "members": ["node_2"],
+            "config": {"platform": "TDX"},
+        },
     },
 })
 
@@ -72,8 +76,8 @@ def millionaire_manual():
 
     # P0 <-> TEE handshake and transfer x (using sugar)
     tee_sk0, tee_pk0 = P2.crypto.kem_keygen("x25519")
-    quote0 = P2.tee.quote(tee_pk0)
-    tee_pk0_at_p0 = P0.tee.attest(P2P(P2, P0, quote0))
+    quote0 = P2.tee.quote_gen(tee_pk0)
+    tee_pk0_at_p0 = P0.tee.attest(P2P(P2, P0, quote0), "TDX")
     v_sk0, v_pk0 = P0.crypto.kem_keygen("x25519")
     shared0_p = P0.crypto.kem_derive(v_sk0, tee_pk0_at_p0, "x25519")
     shared0_t = P2.crypto.kem_derive(tee_sk0, P2P(P0, P2, v_pk0), "x25519")
@@ -88,8 +92,8 @@ def millionaire_manual():
 
     # P1 <-> TEE handshake and transfer y (still show original style for contrast)
     tee_sk1, tee_pk1 = P2.crypto.kem_keygen("x25519")
-    quote1 = P2.tee.quote(tee_pk1)
-    tee_pk1_at_p1 = P1.tee.attest(P2P(P2, P1, quote1))
+    quote1 = P2.tee.quote_gen(tee_pk1)
+    tee_pk1_at_p1 = P1.tee.attest(P2P(P2, P1, quote1), "TDX")
     v_sk1, v_pk1 = P1.crypto.kem_keygen("x25519")
     shared1_p = P1.crypto.kem_derive(v_sk1, tee_pk1_at_p1, "x25519")
     shared1_t = P2.crypto.kem_derive(tee_sk1, P2P(P1, P2, v_pk1), "x25519")
@@ -117,7 +121,10 @@ def millionaire_manual():
 def main():
     print("-" * 10, "TEE millionaire: device vs manual (end-to-end IR)", "-" * 10)
     # Create simulator with TEE bindings
-    tee_bindings = {"tee.quote": "mock_tee.quote", "tee.attest": "mock_tee.attest"}
+    tee_bindings = {
+        "tee.quote_gen": "mock_tee.quote_gen",
+        "tee.attest": "mock_tee.attest",
+    }
     # Apply tee_bindings per-node (preferred) then construct Simulator
     for n in cluster_spec.nodes.values():
         n.runtime_info.op_bindings.update(tee_bindings)