From 6d11362034d74d7ca1209045e1bc245f70bb30c9 Mon Sep 17 00:00:00 2001
From: Enginex0 <enginex0@users.noreply.github.com>
Date: Tue, 19 May 2026 17:45:54 +0100
Subject: [PATCH] feat(intercept): resolve KEY_ID via teeResponses

PR #22's KEY_ID lookup at Keystore2Interceptor.onPreTransact only
scanned generatedKeys, which is populated exclusively by
doSoftwareKeyGen (GENERATE mode and the attest-key override path).
For AUTO packages on TEE-good devices the real TEE handles
generateKey and the response lands in teeResponses via the existing
GENERATE_KEY_TRANSACTION post-hook, so getKeyEntry(KEY_ID) by
TimingSideChannelProbe missed and the call leaked SSE.

Add findTeeResponseByKeyId companion helper that mirrors
findGeneratedKeyByKeyId's shape but scans teeResponses keyed by
response.metadata.key.nspace. Wire it as a fallback after the
existing PR #22 lookup. Behavior unchanged for GENERATE packages.
---
 .../interception/keystore/Keystore2Interceptor.kt        | 9 +++++++++
 .../keystore/shim/KeyMintSecurityLevelInterceptor.kt     | 8 ++++++++
 2 files changed, 17 insertions(+)

diff --git a/app/src/main/java/org/matrix/TEESimulator/interception/keystore/Keystore2Interceptor.kt b/app/src/main/java/org/matrix/TEESimulator/interception/keystore/Keystore2Interceptor.kt
index 82b5648..6ed2570 100644
--- a/app/src/main/java/org/matrix/TEESimulator/interception/keystore/Keystore2Interceptor.kt
+++ b/app/src/main/java/org/matrix/TEESimulator/interception/keystore/Keystore2Interceptor.kt
@@ -237,6 +237,15 @@ object Keystore2Interceptor : AbstractKeystoreInterceptor() {
                         )
                         return InterceptorUtils.createTypedObjectReply(info.response)
                     }
+                    val teeResp = KeyMintSecurityLevelInterceptor.findTeeResponseByKeyId(
+                        callingUid, descriptor.nspace
+                    )
+                    if (teeResp != null) {
+                        SystemLogger.info(
+                            "[TX_ID: $txId] Found TEE response via KEY_ID nspace=${descriptor.nspace}"
+                        )
+                        return InterceptorUtils.createTypedObjectReply(teeResp)
+                    }
                 }
                 return TransactionResult.ContinueAndSkipPost
             }
diff --git a/app/src/main/java/org/matrix/TEESimulator/interception/keystore/shim/KeyMintSecurityLevelInterceptor.kt b/app/src/main/java/org/matrix/TEESimulator/interception/keystore/shim/KeyMintSecurityLevelInterceptor.kt
index f651bd5..a41a8f4 100644
--- a/app/src/main/java/org/matrix/TEESimulator/interception/keystore/shim/KeyMintSecurityLevelInterceptor.kt
+++ b/app/src/main/java/org/matrix/TEESimulator/interception/keystore/shim/KeyMintSecurityLevelInterceptor.kt
@@ -1244,6 +1244,14 @@ class KeyMintSecurityLevelInterceptor(
                 ?.value
         }
 
+        fun findTeeResponseByKeyId(callingUid: Int, nspace: Long?): KeyEntryResponse? {
+            if (nspace == null || nspace == 0L) return null
+            return teeResponses.entries
+                .filter { (keyId, _) -> keyId.uid == callingUid }
+                .find { (_, response) -> response.metadata?.key?.nspace == nspace }
+                ?.value
+        }
+
         fun getPatchedChain(keyId: KeyIdentifier): Array<Certificate>? = patchedChains[keyId]
 
         fun isAttestationKey(keyId: KeyIdentifier): Boolean = attestationKeys.contains(keyId)