PQC: CryptoCb support for KEM algorithm Kyber

Add support for crypto callback and device id for all three Kyber PQC KEM
function calls.

Signed-off-by: Tobias Frauenschläger <t.frauenschlaeger@me.com>

src/internal.c

@@ -27462,7 +27462,8 @@ int CreateDevPrivateKey(void** pkey, byte* data, word32 length, int hsType,
         }
 
         if (label) {
-            ret = wc_dilithium_init_label(dilithiumKey, (char*)data, heap, devId);
+            ret = wc_dilithium_init_label(dilithiumKey, (char*)data,
+                                          heap, devId);
         }
         else if (id) {
             ret = wc_dilithium_init_id(dilithiumKey, data, length, heap, devId);
@@ -27654,7 +27655,8 @@ int DecodePrivateKey(WOLFSSL *ssl, word16* length)
                 }
 
                 /* Return the maximum signature length. */
-                *length = (word16)wc_dilithium_sig_size((dilithium_key*)ssl->hsKey);
+                *length = (word16)wc_dilithium_sig_size(
+                                    (dilithium_key*)ssl->hsKey);
             }
     #else
             ret = NOT_COMPILED_IN;

src/tls.c

@@ -8415,7 +8415,7 @@ static int TLSX_KeyShare_ProcessPqc(WOLFSSL* ssl, KeyShareEntry* keyShareEntry)
         return BAD_FUNC_ARG;
     }
 
-    ret = wc_KyberKey_Init(type, kem, ssl->heap, INVALID_DEVID);
+    ret = wc_KyberKey_Init(type, kem, ssl->heap, ssl->devId);
     if (ret != 0) {
         wc_ecc_free(&eccpubkey);
         WOLFSSL_MSG("Error creating Kyber KEM");
@@ -8907,7 +8907,7 @@ static int server_generate_pqc_ciphertext(WOLFSSL* ssl,
         return MEMORY_E;
     }
 
-    ret = wc_KyberKey_Init(type, kem, ssl->heap, INVALID_DEVID);
+    ret = wc_KyberKey_Init(type, kem, ssl->heap, ssl->devId);
     if (ret != 0) {
         wc_ecc_free(&eccpubkey);
         WOLFSSL_MSG("Error creating Kyber KEM");

wolfcrypt/src/cryptocb.c

@@ -785,6 +785,129 @@ int wc_CryptoCb_Ed25519Verify(const byte* sig, word32 sigLen,
 }
 #endif /* HAVE_ED25519 */
 
+#if defined(HAVE_PQC) && defined(WOLFSSL_HAVE_KYBER)
+int wc_CryptoCb_PqcKemGetDevId(int type, void* key)
+{
+    int devId = INVALID_DEVID;
+
+    if (key == NULL)
+        return devId;
+
+    /* get devId */
+#if defined(WOLFSSL_HAVE_KYBER)
+    if (type == WC_PQC_KEM_TYPE_KYBER) {
+        devId = ((KyberKey*) key)->devId;
+    }
+#endif
+
+    return devId;
+}
+
+int wc_CryptoCb_MakePqcKemKey(WC_RNG* rng, int type, int keySize, void* key)
+{
+    int ret = CRYPTOCB_UNAVAILABLE;
+    int devId = INVALID_DEVID;
+    CryptoCb* dev;
+
+    if (key == NULL)
+        return ret;
+
+    /* get devId */
+    devId = wc_CryptoCb_PqcKemGetDevId(type, key);
+    if (devId == INVALID_DEVID)
+        return ret;
+
+    /* locate registered callback */
+    dev = wc_CryptoCb_FindDevice(devId, WC_ALGO_TYPE_PK);
+    if (dev && dev->cb) {
+        wc_CryptoInfo cryptoInfo;
+        XMEMSET(&cryptoInfo, 0, sizeof(cryptoInfo));
+        cryptoInfo.algo_type = WC_ALGO_TYPE_PK;
+        cryptoInfo.pk.type = WC_PK_TYPE_PQC_KEM_KEYGEN;
+        cryptoInfo.pk.pqc_kem_kg.rng = rng;
+        cryptoInfo.pk.pqc_kem_kg.size = keySize;
+        cryptoInfo.pk.pqc_kem_kg.key = key;
+        cryptoInfo.pk.pqc_kem_kg.type = type;
+
+        ret = dev->cb(dev->devId, &cryptoInfo, dev->ctx);
+    }
+
+    return wc_CryptoCb_TranslateErrorCode(ret);
+}
+
+int wc_CryptoCb_PqcEncapsulate(byte* ciphertext, word32 ciphertextLen,
+    byte* sharedSecret, word32 sharedSecretLen, WC_RNG* rng, int type,
+    void* key)
+{
+    int ret = CRYPTOCB_UNAVAILABLE;
+    int devId = INVALID_DEVID;
+    CryptoCb* dev;
+
+    if (key == NULL)
+        return ret;
+
+    /* get devId */
+    devId = wc_CryptoCb_PqcKemGetDevId(type, key);
+    if (devId == INVALID_DEVID)
+        return ret;
+
+    /* locate registered callback */
+    dev = wc_CryptoCb_FindDevice(devId, WC_ALGO_TYPE_PK);
+    if (dev && dev->cb) {
+        wc_CryptoInfo cryptoInfo;
+        XMEMSET(&cryptoInfo, 0, sizeof(cryptoInfo));
+        cryptoInfo.algo_type = WC_ALGO_TYPE_PK;
+        cryptoInfo.pk.type = WC_PK_TYPE_PQC_KEM_ENCAPS;
+        cryptoInfo.pk.pqc_encaps.ciphertext = ciphertext;
+        cryptoInfo.pk.pqc_encaps.ciphertextLen = ciphertextLen;
+        cryptoInfo.pk.pqc_encaps.sharedSecret = sharedSecret;
+        cryptoInfo.pk.pqc_encaps.sharedSecretLen = sharedSecretLen;
+        cryptoInfo.pk.pqc_encaps.rng = rng;
+        cryptoInfo.pk.pqc_encaps.key = key;
+        cryptoInfo.pk.pqc_encaps.type = type;
+
+        ret = dev->cb(dev->devId, &cryptoInfo, dev->ctx);
+    }
+
+    return wc_CryptoCb_TranslateErrorCode(ret);
+}
+
+int wc_CryptoCb_PqcDecapsulate(const byte* ciphertext, word32 ciphertextLen,
+    byte* sharedSecret, word32 sharedSecretLen, int type, void* key)
+{
+    int ret = CRYPTOCB_UNAVAILABLE;
+    int devId = INVALID_DEVID;
+    CryptoCb* dev;
+
+    if (key == NULL)
+        return ret;
+
+    /* get devId */
+    devId = wc_CryptoCb_PqcKemGetDevId(type, key);
+    if (devId == INVALID_DEVID)
+        return ret;
+
+    /* locate registered callback */
+    dev = wc_CryptoCb_FindDevice(devId, WC_ALGO_TYPE_PK);
+    if (dev && dev->cb) {
+        wc_CryptoInfo cryptoInfo;
+        XMEMSET(&cryptoInfo, 0, sizeof(cryptoInfo));
+        cryptoInfo.algo_type = WC_ALGO_TYPE_PK;
+        cryptoInfo.pk.type = WC_PK_TYPE_PQC_KEM_DECAPS;
+        cryptoInfo.pk.pqc_decaps.ciphertext = ciphertext;
+        cryptoInfo.pk.pqc_decaps.ciphertextLen = ciphertextLen;
+        cryptoInfo.pk.pqc_decaps.sharedSecret = sharedSecret;
+        cryptoInfo.pk.pqc_decaps.sharedSecretLen = sharedSecretLen;
+        cryptoInfo.pk.pqc_decaps.key = key;
+        cryptoInfo.pk.pqc_decaps.type = type;
+
+        ret = dev->cb(dev->devId, &cryptoInfo, dev->ctx);
+    }
+
+    return wc_CryptoCb_TranslateErrorCode(ret);
+}
+#endif /* HAVE_PQC && WOLFSSL_HAVE_KYBER */
+
 #if defined(HAVE_PQC) && (defined(HAVE_FALCON) || defined(HAVE_DILITHIUM))
 int wc_CryptoCb_PqcSigGetDevId(int type, void* key)
 {

wolfcrypt/src/ext_kyber.c

@@ -101,10 +101,15 @@ int wc_KyberKey_Init(int type, KyberKey* key, void* heap, int devId)
 
         /* Keep type for parameters. */
         key->type = type;
+
+#ifdef WOLF_CRYPTO_CB
+        key->devCtx = NULL;
+        key->devId = devId;
+#endif
     }
 
-    (void)devId;
     (void)heap;
+    (void)devId;
 
     return ret;
 }
@@ -312,13 +317,25 @@ int wc_KyberKey_MakeKey(KyberKey* key, WC_RNG* rng)
     OQS_KEM *kem = NULL;
 #endif
 
-    (void)rng;
-
     /* Validate parameter. */
     if (key == NULL) {
         return BAD_FUNC_ARG;
     }
 
+#ifdef WOLF_CRYPTO_CB
+    #ifndef WOLF_CRYPTO_CB_FIND
+    if (key->devId != INVALID_DEVID)
+    #endif
+    {
+        ret = wc_CryptoCb_MakePqcKemKey(rng, WC_PQC_KEM_TYPE_KYBER,
+                                        key->type, key);
+        if (ret != CRYPTOCB_UNAVAILABLE)
+            return ret;
+        /* fall-through when unavailable */
+        ret = 0;
+    }
+#endif
+
 #ifdef HAVE_LIBOQS
     if (ret == 0) {
         algName = OQS_ID2name(key->type);
@@ -327,12 +344,6 @@ int wc_KyberKey_MakeKey(KyberKey* key, WC_RNG* rng)
         }
     }
 
-    if (ret == 0) {
-        algName = OQS_ID2name(key->type);
-        if (algName == NULL) {
-            ret = BAD_FUNC_ARG;
-        }
-    }
     if (ret == 0) {
         kem = OQS_KEM_new(algName);
         if (kem == NULL) {
@@ -403,6 +414,9 @@ int wc_KyberKey_Encapsulate(KyberKey* key, unsigned char* ct, unsigned char* ss,
     WC_RNG* rng)
 {
     int ret = 0;
+#ifdef WOLF_CRYPTO_CB
+    word32 ctlen = 0;
+#endif
 #ifdef HAVE_LIBOQS
     const char * algName = NULL;
     OQS_KEM *kem = NULL;
@@ -415,6 +429,24 @@ int wc_KyberKey_Encapsulate(KyberKey* key, unsigned char* ct, unsigned char* ss,
         ret = BAD_FUNC_ARG;
     }
 
+#ifdef WOLF_CRYPTO_CB
+    if (ret == 0) {
+        ret = wc_KyberKey_CipherTextSize(key, &ctlen);
+    }
+    if ((ret == 0)
+    #ifndef WOLF_CRYPTO_CB_FIND
+        && (key->devId != INVALID_DEVID)
+    #endif
+    ) {
+        ret = wc_CryptoCb_PqcEncapsulate(ct, ctlen, ss, KYBER_SS_SZ, rng,
+                                         WC_PQC_KEM_TYPE_KYBER, key);
+        if (ret != CRYPTOCB_UNAVAILABLE)
+            return ret;
+        /* fall-through when unavailable */
+        ret = 0;
+    }
+#endif
+
 #ifdef HAVE_LIBOQS
     if (ret == 0) {
         algName = OQS_ID2name(key->type);
@@ -509,6 +541,21 @@ int wc_KyberKey_Decapsulate(KyberKey* key, unsigned char* ss,
         ret = BUFFER_E;
     }
 
+#ifdef WOLF_CRYPTO_CB
+    if ((ret == 0)
+    #ifndef WOLF_CRYPTO_CB_FIND
+        && (key->devId != INVALID_DEVID)
+    #endif
+    ) {
+        ret = wc_CryptoCb_PqcDecapsulate(ct, ctlen, ss, KYBER_SS_SZ,
+                                         WC_PQC_KEM_TYPE_KYBER, key);
+        if (ret != CRYPTOCB_UNAVAILABLE)
+            return ret;
+        /* fall-through when unavailable */
+        ret = 0;
+    }
+#endif
+
 #ifdef HAVE_LIBOQS
     if (ret == 0) {
         algName = OQS_ID2name(key->type);

wolfssl/wolfcrypt/cryptocb.h

@@ -71,6 +71,14 @@
 #if defined(WOLFSSL_SHA512) || defined(WOLFSSL_SHA384)
     #include <wolfssl/wolfcrypt/sha512.h>
 #endif
+#ifdef HAVE_PQC
+    #include <wolfssl/wolfcrypt/kyber.h>
+#ifdef WOLFSSL_WC_KYBER
+    #include <wolfssl/wolfcrypt/wc_kyber.h>
+#elif defined(HAVE_LIBOQS) || defined(HAVE_PQM4)
+    #include <wolfssl/wolfcrypt/ext_kyber.h>
+#endif
+#endif
 #if defined(HAVE_PQC) && defined(HAVE_DILITHIUM)
     #include <wolfssl/wolfcrypt/dilithium.h>
 #endif
@@ -208,6 +216,31 @@ typedef struct wc_CryptoInfo {
                 byte         contextLen;
             } ed25519verify;
         #endif
+        #if defined(HAVE_PQC) && defined(WOLFSSL_HAVE_KYBER)
+            struct {
+                WC_RNG*     rng;
+                int         size;
+                void*       key;
+                int         type; /* enum wc_PqcKemType */
+            } pqc_kem_kg;
+            struct {
+                byte*       ciphertext;
+                word32      ciphertextLen;
+                byte*       sharedSecret;
+                word32      sharedSecretLen;
+                WC_RNG*     rng;
+                void*       key;
+                int         type; /* enum wc_PqcKemType */
+            } pqc_encaps;
+            struct {
+                const byte* ciphertext;
+                word32      ciphertextLen;
+                byte*       sharedSecret;
+                word32      sharedSecretLen;
+                void*       key;
+                int         type; /* enum wc_PqcKemType */
+            } pqc_decaps;
+        #endif
         #if defined(HAVE_PQC) && \
             (defined(HAVE_FALCON) || defined(HAVE_DILITHIUM))
             struct {
@@ -492,6 +525,21 @@ WOLFSSL_LOCAL int wc_CryptoCb_Ed25519Verify(const byte* sig, word32 sigLen,
     const byte* context, byte contextLen);
 #endif /* HAVE_ED25519 */
 
+#if defined(HAVE_PQC) && defined(WOLFSSL_HAVE_KYBER)
+WOLFSSL_LOCAL int wc_CryptoCb_PqcKemGetDevId(int type, void* key);
+
+WOLFSSL_LOCAL int wc_CryptoCb_MakePqcKemKey(WC_RNG* rng, int type,
+    int keySize, void* key);
+
+WOLFSSL_LOCAL int wc_CryptoCb_PqcEncapsulate(byte* ciphertext,
+    word32 ciphertextLen, byte* sharedSecret, word32 sharedSecretLen,
+    WC_RNG* rng, int type, void* key);
+
+WOLFSSL_LOCAL int wc_CryptoCb_PqcDecapsulate(const byte* ciphertext,
+    word32 ciphertextLen, byte* sharedSecret, word32 sharedSecretLen,
+    int type, void* key);
+#endif /* HAVE_PQC && WOLFSSL_HAVE_KYBER */
+
 #if defined(HAVE_PQC) && (defined(HAVE_FALCON) || defined(HAVE_DILITHIUM))
 WOLFSSL_LOCAL int wc_CryptoCb_PqcSigGetDevId(int type, void* key);
 

wolfssl/wolfcrypt/ext_kyber.h

@@ -22,6 +22,10 @@
 #ifndef EXT_KYBER_H
 #define EXT_KYBER_H
 
+#ifdef WOLF_CRYPTO_CB
+    #include <wolfssl/wolfcrypt/cryptocb.h>
+#endif
+
 #ifdef WOLFSSL_HAVE_KYBER
 #include <wolfssl/wolfcrypt/kyber.h>
 
@@ -56,6 +60,12 @@ struct KyberKey {
      * Note we don't save the variant (SHAKE vs AES) as that is decided at
      * configuration time. */
     int type;
+
+#ifdef WOLF_CRYPTO_CB
+    void* devCtx;
+    int   devId;
+#endif
+
     byte priv[EXT_KYBER_MAX_PRIV_SZ];
     byte pub[EXT_KYBER_MAX_PUB_SZ];
 };

wolfssl/wolfcrypt/types.h

@@ -1170,11 +1170,18 @@ typedef struct w64wrapper {
         WC_PK_TYPE_CURVE25519_KEYGEN = 16,
         WC_PK_TYPE_RSA_GET_SIZE = 17,
         #define _WC_PK_TYPE_MAX WC_PK_TYPE_RSA_GET_SIZE
+    #if defined(HAVE_PQC) && defined(WOLFSSL_HAVE_KYBER)
+        WC_PK_TYPE_PQC_KEM_KEYGEN = 18,
+        WC_PK_TYPE_PQC_KEM_ENCAPS = 19,
+        WC_PK_TYPE_PQC_KEM_DECAPS = 20,
+        #undef _WC_PK_TYPE_MAX
+        #define _WC_PK_TYPE_MAX WC_PK_TYPE_PQC_KEM_DECAPS
+    #endif
     #if defined(HAVE_PQC) && (defined(HAVE_DILITHIUM) || defined(HAVE_FALCON))
-        WC_PK_TYPE_PQC_SIG_KEYGEN = 18,
-        WC_PK_TYPE_PQC_SIG_SIGN = 19,
-        WC_PK_TYPE_PQC_SIG_VERIFY = 20,
-        WC_PK_TYPE_PQC_SIG_CHECK_PRIV_KEY = 21,
+        WC_PK_TYPE_PQC_SIG_KEYGEN = 21,
+        WC_PK_TYPE_PQC_SIG_SIGN = 22,
+        WC_PK_TYPE_PQC_SIG_VERIFY = 23,
+        WC_PK_TYPE_PQC_SIG_CHECK_PRIV_KEY = 24,
         #undef _WC_PK_TYPE_MAX
         #define _WC_PK_TYPE_MAX WC_PK_TYPE_PQC_SIG_CHECK_PRIV_KEY
     #endif
@@ -1182,6 +1189,18 @@ typedef struct w64wrapper {
     };
 
     #if defined(HAVE_PQC)
+    /* Post quantum KEM algorithms */
+    enum wc_PqcKemType {
+        WC_PQC_KEM_TYPE_NONE = 0,
+        #define _WC_PQC_KEM_TYPE_MAX WC_PQC_KEM_TYPE_NONE
+    #if defined(WOLFSSL_HAVE_KYBER)
+        WC_PQC_KEM_TYPE_KYBER = 1,
+        #undef _WC_PQC_KEM_TYPE_MAX
+        #define _WC_PQC_KEM_TYPE_MAX WC_PQC_KEM_TYPE_KYBER
+    #endif
+        WC_PQC_KEM_TYPE_MAX = _WC_PQC_KEM_TYPE_MAX
+    };
+
     /* Post quantum signature algorithms */
     enum wc_PqcSignatureType {
         WC_PQC_SIG_TYPE_NONE = 0,