From e88dfad31c430a9d805f30537577d38fa5255976 Mon Sep 17 00:00:00 2001 From: Aniket Dixit Date: Fri, 22 May 2026 18:46:41 +0530 Subject: [PATCH] registry updates --- tee_gateway/controllers/ohttp_controller.py | 2 +- tee_gateway/tee_manager.py | 104 +++++++++++++------- 2 files changed, 70 insertions(+), 36 deletions(-) diff --git a/tee_gateway/controllers/ohttp_controller.py b/tee_gateway/controllers/ohttp_controller.py index 50e5496..087d79c 100644 --- a/tee_gateway/controllers/ohttp_controller.py +++ b/tee_gateway/controllers/ohttp_controller.py @@ -405,7 +405,7 @@ def get_hpke_config(): """ try: tee = get_tee_keys() - return tee.get_hpke_config(), 200 + return tee.get_signed_hpke_config(), 200 except Exception as exc: logger.error("HPKE config error: %s", exc, exc_info=True) return {"error": "Failed to retrieve HPKE config"}, 500 diff --git a/tee_gateway/tee_manager.py b/tee_gateway/tee_manager.py index 98deaa6..b2bf927 100644 --- a/tee_gateway/tee_manager.py +++ b/tee_gateway/tee_manager.py @@ -39,8 +39,7 @@ def __init__(self, register=True): self.tee_id = None self.wallet_address = None # HPKE keypair for OHTTP-style anonymous inference. Generated in the - # same enclave boot so the X25519 public key is covered by the same - # attestation that covers the RSA signing key. + # enclave and bound to the attested signing key by get_signed_hpke_config(). self.hpke_private_key = None self.hpke_public_key_raw: bytes | None = None self._generate_keys() @@ -78,10 +77,8 @@ def _generate_keys(self): self.wallet_address = wallet_account.address # HPKE X25519 keypair — independent random material from the RSA - # signing key. Both public keys are bound to the enclave by the - # nitriding attestation transcript (register_with_nitriding below), - # so verifiers still get a single attested fingerprint covering both, - # without sharing private-key material between them. + # signing key. The HPKE public config is signed by the RSA key for + # registry V2, while Nitro attestation continues to bind the RSA key. self.hpke_private_key, self.hpke_public_key_raw = ohttp.generate_keypair() logger.info("TEE key pair generated successfully") @@ -94,41 +91,18 @@ def _generate_keys(self): def register_with_nitriding(self): """Register public key hash with nitriding. - The hash covers both the RSA signing key (DER-encoded SPKI) and the - raw X25519 HPKE public key. Including both in a single attested digest - means a verifier who validates the attestation document automatically - gets binding for the HPKE config used for anonymous inference — no - separate trust anchor required. + The current on-chain Nitro verifier expects nitriding user_data to + contain SHA256(signing public key DER) as the second hash. OHTTP/HPKE + config is bound separately by an RSA-PSS signature from this same + attested signing key. """ - # Defensive check: the v2 transcript labels both keys, so refusing to - # register is the only safe behavior when one is missing. Falling back - # to b"" would produce a digest that's nominally v2 but only covers - # RSA, and a verifier trusting the label would accept an enclave whose - # HPKE key was never attested. Raise outside the broad try/except - # below so a real misconfiguration isn't masked as a non-TEE - # environment. - if not self.hpke_public_key_raw or len(self.hpke_public_key_raw) != 32: - raise RuntimeError( - "Refusing to register with nitriding: HPKE X25519 public key " - "is missing or wrong length; the v2 attestation transcript " - "requires both RSA and HPKE keys." - ) - try: public_key_der = self.public_key.public_bytes( encoding=serialization.Encoding.DER, format=serialization.PublicFormat.SubjectPublicKeyInfo, ) - # Domain-separated transcript so a future addition of more keys - # can't be confused with the existing layout. - transcript = ( - b"og-tee-keys|v2|rsa-spki=" - + public_key_der - + b"|hpke-x25519=" - + self.hpke_public_key_raw - ) - key_hash = hashlib.sha256(transcript).digest() + key_hash = hashlib.sha256(public_key_der).digest() key_hash_b64 = base64.b64encode(key_hash).decode("utf-8") logger.info(f"Public key DER length: {len(public_key_der)} bytes") @@ -216,13 +190,40 @@ def get_hpke_config(self) -> dict: ).decode("ascii"), } + def get_signed_hpke_config(self) -> dict: + """Return HPKE config plus an RSA-PSS signature for registry V2. + + The signature covers the same ABI-encoded hash computed by + TEERegistryV2.computeOHTTPConfigHash, so the registry can verify that + the OHTTP key came from the enclave signing key already proven by Nitro + attestation. + """ + config = self.get_hpke_config() + public_key = bytes.fromhex(config["public_key"]) + key_config = base64.b64decode(config["key_config"]) + config_hash = compute_ohttp_config_hash( + bytes.fromhex(self.tee_id), + config["key_id"], + config["kem_id"], + config["kdf_id"], + config["aead_id"], + public_key, + key_config, + ) + return { + **config, + "tee_id": f"0x{self.tee_id}", + "signature": self.sign_data(config_hash), + "signature_hash": f"0x{config_hash.hex()}", + } + def get_attestation_document(self) -> dict: """Return TEE attestation document.""" return { "public_key": self.public_key_pem, "tee_id": f"0x{self.tee_id}", "wallet_address": self.wallet_address, - "hpke": self.get_hpke_config() if self.hpke_public_key_raw else None, + "hpke": self.get_signed_hpke_config() if self.hpke_public_key_raw else None, "timestamp": datetime.now(UTC).isoformat(), "enclave_info": { "platform": "aws-nitro", @@ -233,6 +234,39 @@ def get_attestation_document(self) -> dict: } +def compute_ohttp_config_hash( + tee_id: bytes, + key_id: int, + kem_id: int, + kdf_id: int, + aead_id: int, + ohttp_public_key: bytes, + ohttp_key_config: bytes, +) -> bytes: + """Compute TEERegistryV2.computeOHTTPConfigHash off-chain.""" + if len(tee_id) != 32: + raise ValueError("tee_id must be 32 bytes") + + def word(value: int) -> bytes: + return value.to_bytes(32, "big") + + domain = keccak(b"OPENGRADIENT_TEE_OHTTP_CONFIG_V1") + return keccak( + b"".join( + [ + domain, + tee_id, + word(key_id), + word(kem_id), + word(kdf_id), + word(aead_id), + keccak(ohttp_public_key), + keccak(ohttp_key_config), + ] + ) + ) + + def signal_ready(): """Signal to nitriding that enclave is ready to accept traffic.""" try: