Clean up cached keys on successful import (#18)

Generated and attestation keys are cached, and if a key is imported with the same name, the cached key would be returned instead of the newly imported one.

This change invalidates the cached key when a key is successfully imported with the same alias.
Close #17 as fixed.

The logging has also been improved to be more consistent across the different interceptors.

app/src/main/java/org/matrix/TEESimulator/interception/keystore/Keystore2Interceptor.kt

@@ -89,24 +89,17 @@ object Keystore2Interceptor : AbstractKeystoreInterceptor() {
         data: Parcel,
     ): TransactionResult {
         if (code == GET_KEY_ENTRY_TRANSACTION || code == DELETE_KEY_TRANSACTION) {
+            logTransaction(txId, transactionNames[code]!!, callingUid, callingPid)
+
             data.enforceInterface(IKeystoreService.DESCRIPTOR)
             val descriptor =
                 data.readTypedObject(KeyDescriptor.CREATOR)
                     ?: return TransactionResult.SkipTransaction
-            logTransaction(
-                txId,
-                "${transactionNames[code]} (alias=${descriptor.alias})",
-                callingUid,
-                callingPid,
-            )
 
-            if (ConfigurationManager.shouldSkipUid(callingUid)) {
+            if (ConfigurationManager.shouldSkipUid(callingUid))
-                SystemLogger.debug(
-                    "[TX_ID: $txId] Skip post-transaction hook for UID=${callingUid}"
-                )
                 return TransactionResult.ContinueAndSkipPost
-            }
 
+            SystemLogger.info("Handling ${transactionNames[code]!!} ${descriptor.alias}")
             val keyId = KeyIdentifier(callingUid, descriptor.alias)
 
             if (code == DELETE_KEY_TRANSACTION) {
@@ -119,7 +112,7 @@ object Keystore2Interceptor : AbstractKeystoreInterceptor() {
                     ?: return TransactionResult.Continue
 
             if (KeyMintSecurityLevelInterceptor.isAttestationKey(keyId))
-                SystemLogger.debug("${descriptor.alias} was an attestation key")
+                SystemLogger.info("${descriptor.alias} was an attestation key")
 
             SystemLogger.info("[TX_ID: $txId] Found generated response for ${descriptor.alias}:")
             response.metadata?.authorizations?.forEach {
@@ -155,20 +148,17 @@ object Keystore2Interceptor : AbstractKeystoreInterceptor() {
             return TransactionResult.SkipTransaction
 
         if (code == GET_KEY_ENTRY_TRANSACTION) {
+            logTransaction(txId, "post-${transactionNames[code]!!}", callingUid, callingPid)
 
             data.enforceInterface(IKeystoreService.DESCRIPTOR)
             val keyDescriptor =
                 data.readTypedObject(KeyDescriptor.CREATOR)
                     ?: return TransactionResult.SkipTransaction
-            logTransaction(
+
-                txId,
-                "post-getKeyEntry (alias=${keyDescriptor.alias})",
-                callingUid,
-                callingPid,
-            )
             if (!ConfigurationManager.shouldPatch(callingUid))
                 return TransactionResult.SkipTransaction
 
+            SystemLogger.info("Handling post-${transactionNames[code]!!} ${keyDescriptor.alias}")
             return try {
                 val response =
                     reply.readTypedObject(KeyEntryResponse.CREATOR)

app/src/main/java/org/matrix/TEESimulator/interception/keystore/shim/KeyMintSecurityLevelInterceptor.kt

@@ -40,11 +40,20 @@ class KeyMintSecurityLevelInterceptor(
         callingPid: Int,
         data: Parcel,
     ): TransactionResult {
-        // This interceptor only handles the 'generateKey' transaction directly.
         if (code == GENERATE_KEY_TRANSACTION) {
-            logTransaction(txId, "generateKey", callingUid, callingPid)
+            logTransaction(txId, transactionNames[code]!!, callingUid, callingPid)
+
             data.enforceInterface(IKeystoreSecurityLevel.DESCRIPTOR)
             return handleGenerateKey(callingUid, data)
+        } else if (code == IMPORT_KEY_TRANSACTION) {
+            logTransaction(txId, transactionNames[code]!!, callingUid, callingPid)
+
+            data.enforceInterface(IKeystoreSecurityLevel.DESCRIPTOR)
+            val alias =
+                data.readTypedObject(KeyDescriptor.CREATOR)?.alias
+                    ?: return TransactionResult.ContinueAndSkipPost
+            SystemLogger.info("Handling post-${transactionNames[code]} ${alias}")
+            return TransactionResult.Continue
         } else {
             logTransaction(
                 txId,
@@ -57,6 +66,35 @@ class KeyMintSecurityLevelInterceptor(
         return TransactionResult.ContinueAndSkipPost
     }
 
+    override fun onPostTransact(
+        txId: Long,
+        target: IBinder,
+        code: Int,
+        flags: Int,
+        callingUid: Int,
+        callingPid: Int,
+        data: Parcel,
+        reply: Parcel?,
+        resultCode: Int,
+    ): TransactionResult {
+        // We only care about successful 'importKey' transactions to clean cached keys.
+        if (
+            code == IMPORT_KEY_TRANSACTION &&
+                resultCode == 0 &&
+                reply != null &&
+                !InterceptorUtils.hasException(reply)
+        ) {
+            logTransaction(txId, "post-${transactionNames[code]!!}", callingUid, callingPid)
+
+            data.enforceInterface(IKeystoreSecurityLevel.DESCRIPTOR)
+            val keyDescriptor =
+                data.readTypedObject(KeyDescriptor.CREATOR)
+                    ?: return TransactionResult.SkipTransaction
+            cleanupKeyData(KeyIdentifier(callingUid, keyDescriptor.alias))
+        }
+        return TransactionResult.SkipTransaction
+    }
+
     /**
      * Handles the `generateKey` transaction. Based on the configuration for the calling UID, it
      * either generates a key in software or lets the call pass through to the hardware.
@@ -66,7 +104,7 @@ class KeyMintSecurityLevelInterceptor(
                 val keyDescriptor = data.readTypedObject(KeyDescriptor.CREATOR)!!
                 val attestationKey = data.readTypedObject(KeyDescriptor.CREATOR)
                 SystemLogger.debug(
-                    "[key, attestationKey]: ${keyDescriptor.alias}, ${attestationKey?.alias}"
+                    "Handling generateKey ${keyDescriptor.alias}, attestKey=${attestationKey?.alias}"
                 )
                 val params = data.createTypedArray(KeyParameter.CREATOR)!!
                 val parsedParams = KeyMintAttestation(params)
@@ -84,9 +122,7 @@ class KeyMintSecurityLevelInterceptor(
                             isAttestationKey(KeyIdentifier(callingUid, attestationKey.alias)))
 
                 if (needsSoftwareGeneration) {
-                    SystemLogger.info(
+                    SystemLogger.info("Generating software key for ${keyId}.")
-                        "Generating software key for alias '${keyDescriptor.alias}' (UID: $callingUid)."
-                    )
 
                     // Generate the key pair and certificate chain.
                     val keyData =
@@ -116,11 +152,11 @@ class KeyMintSecurityLevelInterceptor(
 
                 // If not generating, clear any stale state for this alias and let the call proceed.
                 cleanupKeyData(keyId)
-                TransactionResult.Continue
+                TransactionResult.ContinueAndSkipPost
             }
             .getOrElse {
                 SystemLogger.error("Error during generateKey handling for UID $callingUid.", it)
-                TransactionResult.Continue // Fallback to original service on error.
+                TransactionResult.ContinueAndSkipPost
             }
     }
 
@@ -175,8 +211,12 @@ class KeyMintSecurityLevelInterceptor(
         fun isAttestationKey(keyId: KeyIdentifier): Boolean = attestationKeys.contains(keyId)
 
         fun cleanupKeyData(keyId: KeyIdentifier) {
-            generatedKeys.remove(keyId)
+            if (generatedKeys.remove(keyId) != null) {
-            attestationKeys.remove(keyId)
+                SystemLogger.debug("Remove generated key ${keyId}")
+            }
+            if (attestationKeys.remove(keyId)) {
+                SystemLogger.debug("Remove cached attestaion key ${keyId}")
+            }
         }
     }
 }