allow full offload (with GCM)

Signed-off-by: Sameeh Jubran <sameeh@wolfssl.com>

Author: sameehj

tests/api/test_aes.c

@@ -5284,6 +5284,11 @@ static int test_CryptoCb_Aes_Cb(int devId, wc_CryptoInfo* info, void* ctx)
         /* Store handle in aes->devCtx - this is what wolfSSL will use */
         aes->devCtx = cryptoCbAesMockHandle;
 
+        /* CRITICAL: Declare full AES-GCM offload capability.
+         * This tells wolfSSL to skip GCM table generation and expect
+         * all GCM operations to go through CryptoCB. */
+        info->cipher.aessetkey.capabilities = WC_CRYPTOCB_AES_GCM;
+
         cryptoCbAesSetKeyCalled++;
 
         return 0;
@@ -5486,6 +5491,9 @@ int test_wc_CryptoCb_AesSetKey(void)
     /* Verify callback was invoked */
     ExpectIntEQ(cryptoCbAesSetKeyCalled, 1);
 
+    /* Verify GCM offload flag was set (callback declared capability) */
+    ExpectIntEQ(aes->gcmCryptoCbOffload, 1);
+
     /* Verify handle stored in devCtx */
     ExpectPtrEq(aes->devCtx, cryptoCbAesMockHandle);
 
@@ -5628,6 +5636,11 @@ static int test_CryptoCb_AesGcm_Offload_Cb(int devId, wc_CryptoInfo* info, void*
         /* Store handle in aes->devCtx - this is what wolfSSL will use */
         aes->devCtx = cryptoCbAesGcmMockHandle;
 
+        /* CRITICAL: Declare full AES-GCM offload capability.
+         * This tells wolfSSL to skip GCM table generation and expect
+         * all GCM operations to go through CryptoCB. */
+        info->cipher.aessetkey.capabilities = WC_CRYPTOCB_AES_GCM;
+
         cryptoCbAesGcmSetKeyCalled++;
 
         return 0;
@@ -5863,6 +5876,9 @@ int test_wc_CryptoCb_AesGcm_EncryptDecrypt(void)
     /* Verify SetKey callback was invoked */
     ExpectIntEQ(cryptoCbAesGcmSetKeyCalled, 1);
 
+    /* Verify GCM offload flag was set (callback declared capability) */
+    ExpectIntEQ(aes->gcmCryptoCbOffload, 1);
+
     /* Verify handle stored in devCtx */
     ExpectPtrEq(aes->devCtx, cryptoCbAesGcmMockHandle);
 
@@ -5950,5 +5966,201 @@ int test_wc_CryptoCb_AesGcm_EncryptDecrypt(void)
     return EXPECT_RESULT();
 }
 
+/*
+ * Test: Key import only callback - verifies software fallback works
+ * when callback does NOT declare WC_CRYPTOCB_AES_GCM capability
+ */
+#define TEST_CRYPTOCB_AES_KEYIMPORT_ONLY_DEVID  9
+
+/* Test state tracking for key-import-only test */
+static int cryptoCbKeyImportOnlySetKeyCalled = 0;
+static void* cryptoCbKeyImportOnlyMockHandle = (void*)0xDEADBEEF;
+
+/* Mock SE key storage for key-import-only test */
+static struct {
+    byte key[AES_256_KEY_SIZE];
+    word32 keySz;
+    int valid;
+} mockSeKeyImportOnly;
+
+/* Callback that handles key import but NOT GCM operations */
+static int test_CryptoCb_KeyImportOnly_Cb(int devId, wc_CryptoInfo* info, void* ctx)
+{
+    (void)ctx;
+
+    if (devId != TEST_CRYPTOCB_AES_KEYIMPORT_ONLY_DEVID)
+        return CRYPTOCB_UNAVAILABLE;
+
+    /* AES SetKey operation - simulate SE key import */
+    if (info->algo_type == WC_ALGO_TYPE_CIPHER &&
+        info->cipher.type == WC_CIPHER_AES &&
+        info->cipher.aessetkey.aes != NULL) {
+
+        Aes* aes = info->cipher.aessetkey.aes;
+        const byte* key = info->cipher.aessetkey.key;
+        word32 keySz = info->cipher.aessetkey.keySz;
+
+        /* Validate key */
+        if (key == NULL || keySz == 0 || keySz > AES_256_KEY_SIZE) {
+            return BAD_FUNC_ARG;
+        }
+
+        /* "Import" key to simulated SE storage */
+        XMEMCPY(mockSeKeyImportOnly.key, key, keySz);
+        mockSeKeyImportOnly.keySz = keySz;
+        mockSeKeyImportOnly.valid = 1;
+
+        /* Store handle in aes->devCtx */
+        aes->devCtx = cryptoCbKeyImportOnlyMockHandle;
+
+        /* CRITICAL: DO NOT set WC_CRYPTOCB_AES_GCM capability.
+         * This simulates a device that can import keys but cannot
+         * perform GCM operations. wolfSSL should generate GCM tables
+         * and use software fallback for GCM operations. */
+        /* info->cipher.aessetkey.capabilities stays 0 (default) */
+
+        cryptoCbKeyImportOnlySetKeyCalled++;
+
+        return 0;
+    }
+
+    /* GCM operations not supported - return UNAVAILABLE to trigger software fallback */
+    if (info->algo_type == WC_ALGO_TYPE_CIPHER &&
+        info->cipher.type == WC_CIPHER_AES_GCM) {
+        return CRYPTOCB_UNAVAILABLE;
+    }
+
+    return CRYPTOCB_UNAVAILABLE;
+}
+
+int test_wc_CryptoCb_AesKeyImportOnly(void)
+{
+    EXPECT_DECLS;
+#ifdef WOLFSSL_SMALL_STACK
+    Aes* aes = NULL;
+    byte* key = NULL;
+    byte* iv = NULL;
+    byte* plaintext = NULL;
+    byte* ciphertext = NULL;
+    byte* decrypted = NULL;
+    byte* authTag = NULL;
+    byte* aad = NULL;
+#else
+    Aes aes[1];
+    byte key[AES_128_KEY_SIZE];
+    byte iv[GCM_NONCE_MID_SZ];
+    byte plaintext[32];
+    byte ciphertext[32];
+    byte decrypted[32];
+    byte authTag[AES_BLOCK_SIZE];
+    byte aad[16];
+#endif
+    int ret;
+
+#ifdef WOLFSSL_SMALL_STACK
+    aes = (Aes*)XMALLOC(sizeof(Aes), NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    key = (byte*)XMALLOC(AES_128_KEY_SIZE, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    iv = (byte*)XMALLOC(GCM_NONCE_MID_SZ, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    plaintext = (byte*)XMALLOC(32, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    ciphertext = (byte*)XMALLOC(32, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    decrypted = (byte*)XMALLOC(32, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    authTag = (byte*)XMALLOC(AES_BLOCK_SIZE, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    aad = (byte*)XMALLOC(16, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+
+    if (aes == NULL || key == NULL || iv == NULL || plaintext == NULL ||
+        ciphertext == NULL || decrypted == NULL || authTag == NULL || aad == NULL) {
+        ret = MEMORY_E;
+        goto out;
+    }
+#endif
+
+    /* Initialize test data */
+    XMEMSET(key, 0x01, AES_128_KEY_SIZE);
+    XMEMSET(iv, 0x02, GCM_NONCE_MID_SZ);
+    XMEMSET(plaintext, 0x03, 32);
+    XMEMSET(aad, 0x04, 16);
+    XMEMSET(aes, 0, sizeof(Aes));
+    XMEMSET(&mockSeKeyImportOnly, 0, sizeof(mockSeKeyImportOnly));
+    cryptoCbKeyImportOnlySetKeyCalled = 0;
+
+    /* Register callback for key-import-only device */
+    ret = wc_CryptoCb_RegisterDevice(TEST_CRYPTOCB_AES_KEYIMPORT_ONLY_DEVID,
+                                     test_CryptoCb_KeyImportOnly_Cb, NULL);
+    ExpectIntEQ(ret, 0);
+
+    /* Initialize Aes with device ID */
+    ret = wc_AesInit(aes, NULL, TEST_CRYPTOCB_AES_KEYIMPORT_ONLY_DEVID);
+    ExpectIntEQ(ret, 0);
+    ExpectIntEQ(aes->devId, TEST_CRYPTOCB_AES_KEYIMPORT_ONLY_DEVID);
+
+    /* Set key - should trigger CryptoCB and "import" to mock SE */
+    ret = wc_AesGcmSetKey(aes, key, sizeof(key));
+    ExpectIntEQ(ret, 0);
+
+    /* Verify SetKey callback was invoked (callback is called but then we fall through) */
+    ExpectIntEQ(cryptoCbKeyImportOnlySetKeyCalled, 1);
+
+    /* CRITICAL: Verify GCM offload flag is 0 (capability NOT declared) */
+    ExpectIntEQ(aes->gcmCryptoCbOffload, 0);
+
+    /* Verify keylen is set (software path sets this after falling through) */
+    ExpectIntEQ(aes->keylen, (int)sizeof(key));
+    
+    /* Note: The callback sets devCtx, but since capability is not set, we fall through
+     * to software path. The software path will set up the key in aes->key for GCM
+     * table generation. devCtx may or may not be cleared by software path. */
+
+    /* CRITICAL: Verify GCM tables were generated (H table should be non-zero) */
+    /* If gcmCryptoCbOffload is 0, tables should be generated for software fallback */
+    {
+        int hTableNonZero = 0;
+        int i;
+        for (i = 0; i < WC_AES_BLOCK_SIZE; i++) {
+            if (aes->gcm.H[i] != 0) {
+                hTableNonZero = 1;
+                break;
+            }
+        }
+        ExpectIntEQ(hTableNonZero, 1); /* H table should be generated */
+    }
+
+    /* Encrypt via wolfCrypt API - should fall back to software */
+    /* Callback will return CRYPTOCB_UNAVAILABLE, triggering software path */
+    ret = wc_AesGcmEncrypt(aes, ciphertext, plaintext, 32,
+                           iv, sizeof(iv), authTag, sizeof(authTag),
+                           aad, 16);
+    ExpectIntEQ(ret, 0);
+
+    /* Decrypt via wolfCrypt API - should fall back to software */
+    ret = wc_AesGcmDecrypt(aes, decrypted, ciphertext, 32,
+                           iv, sizeof(iv), authTag, sizeof(authTag),
+                           aad, 16);
+    ExpectIntEQ(ret, 0);
+
+    /* Verify decrypted plaintext matches original */
+    ExpectIntEQ(XMEMCMP(plaintext, decrypted, 32), 0);
+
+    /* Cleanup */
+    wc_AesFree(aes);
+    wc_CryptoCb_UnRegisterDevice(TEST_CRYPTOCB_AES_KEYIMPORT_ONLY_DEVID);
+
+    /* Note: mockSeKeyImportOnly.valid may still be 1 since callback was called,
+     * but the key is not actually used (software path is taken) */
+
+#ifdef WOLFSSL_SMALL_STACK
+out:
+    XFREE(aes, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    XFREE(key, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    XFREE(iv, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    XFREE(aad, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    XFREE(plaintext, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    XFREE(ciphertext, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    XFREE(decrypted, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    XFREE(authTag, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+#endif
+
+    return EXPECT_RESULT();
+}
+
 #endif /* WOLF_CRYPTO_CB && WOLF_CRYPTO_CB_AES_SETKEY && !NO_AES && HAVE_AESGCM */
 

tests/api/test_aes.h

@@ -57,12 +57,14 @@ int test_wc_GmacUpdate(void);
     !defined(NO_AES) && defined(HAVE_AESGCM)
 int test_wc_CryptoCb_AesSetKey(void);
 int test_wc_CryptoCb_AesGcm_EncryptDecrypt(void);
+int test_wc_CryptoCb_AesKeyImportOnly(void);
 #endif
 
 #if defined(WOLF_CRYPTO_CB) && defined(WOLF_CRYPTO_CB_AES_SETKEY) && \
     !defined(NO_AES) && defined(HAVE_AESGCM)
 #define TEST_CRYPTOCB_AES_SETKEY_DECL , TEST_DECL_GROUP("aes", test_wc_CryptoCb_AesSetKey), \
-                                       TEST_DECL_GROUP("aes", test_wc_CryptoCb_AesGcm_EncryptDecrypt)
+                                      TEST_DECL_GROUP("aes", test_wc_CryptoCb_AesGcm_EncryptDecrypt), \
+                                      TEST_DECL_GROUP("aes", test_wc_CryptoCb_AesKeyImportOnly)
 #else
 #define TEST_CRYPTOCB_AES_SETKEY_DECL
 #endif