feat(wpa_supplicant): Add optimized PSK impl

Signed-off-by: Chien Wong <[email protected]>

components/wpa_supplicant/CMakeLists.txt

@@ -119,6 +119,9 @@ if(CONFIG_ESP_WIFI_MBEDTLS_CRYPTO)
     "esp_supplicant/src/crypto/crypto_mbedtls-bignum.c"
     "esp_supplicant/src/crypto/crypto_mbedtls-rsa.c"
     "esp_supplicant/src/crypto/crypto_mbedtls-ec.c")
+    if(NOT CONFIG_IDF_TARGET_ESP32)
+        list(APPEND crypto_src "esp_supplicant/src/crypto/fastpsk.c")
+    endif()
     # Add internal RC4 as RC4 has been removed from mbedtls
     set(crypto_src ${crypto_src} "src/crypto/rc4.c")
     if(NOT CONFIG_MBEDTLS_DES_C)

components/wpa_supplicant/esp_supplicant/src/crypto/crypto_mbedtls.c

@@ -38,6 +38,7 @@
 
 #ifdef CONFIG_FAST_PBKDF2
 #include "fastpbkdf2.h"
+#include "fastpsk.h"
 #endif
 
 static int digest_vector(mbedtls_md_type_t md_type, size_t num_elem,
@@ -751,9 +752,13 @@ int pbkdf2_sha1(const char *passphrase, const u8 *ssid, size_t ssid_len,
                 int iterations, u8 *buf, size_t buflen)
 {
 #ifdef CONFIG_FAST_PBKDF2
+# if CONFIG_IDF_TARGET_ESP32
     fastpbkdf2_hmac_sha1((const u8 *) passphrase, os_strlen(passphrase),
                          ssid, ssid_len, iterations, buf, buflen);
     return 0;
+# else
+    return esp_fast_psk(passphrase, os_strlen(passphrase), ssid, ssid_len, iterations, buf, buflen);
+# endif
 #else
     int ret = mbedtls_pkcs5_pbkdf2_hmac_ext(MBEDTLS_MD_SHA1, (const u8 *) passphrase,
                                             os_strlen(passphrase), ssid,

components/wpa_supplicant/esp_supplicant/src/crypto/fastpsk.c

@@ -0,0 +1,150 @@
+/*
+ * SPDX-FileCopyrightText: 2015-2024 Espressif Systems (Shanghai) CO LTD
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// Specialized and optimized PBKDF2-SHA1 implementation for Wi-Fi PSK
+
+#include "fastpsk.h"
+
+#include <string.h>
+
+#include <sha/sha_dma.h>
+#include <hal/sha_hal.h>
+
+#ifndef PUT_UINT32_BE
+#define PUT_UINT32_BE(n,b,i)                            \
+{                                                       \
+    (b)[(i)    ] = (unsigned char) ( (n) >> 24 );       \
+    (b)[(i) + 1] = (unsigned char) ( (n) >> 16 );       \
+    (b)[(i) + 2] = (unsigned char) ( (n) >>  8 );       \
+    (b)[(i) + 3] = (unsigned char) ( (n)       );       \
+}
+#endif
+
+#define FAST_PSK_SHA1_BLOCKS           2
+#define SHA1_BLOCK_SZ                  64
+#define SHA1_BLOCK_SZ_WORDS            16
+#define SHA1_OUTPUT_SZ                 20
+#define SHA1_OUTPUT_SZ_WORDS           5
+#define FAST_PSK_SHA1_BLOCKS_BUF_BYTES (FAST_PSK_SHA1_BLOCKS * SHA1_BLOCK_SZ)
+#define FAST_PSK_SHA1_BLOCKS_BUF_WORDS (FAST_PSK_SHA1_BLOCKS * SHA1_BLOCK_SZ / 4)
+
+union hmac_block {
+    union {
+        uint32_t words[SHA1_BLOCK_SZ / 4];
+        uint8_t bytes[SHA1_BLOCK_SZ];
+    } block[FAST_PSK_SHA1_BLOCKS];
+    uint8_t whole_bytes[FAST_PSK_SHA1_BLOCKS_BUF_BYTES];
+    uint32_t whole_words[FAST_PSK_SHA1_BLOCKS_BUF_WORDS];
+};
+_Static_assert(sizeof(union hmac_block) == 128, "Incorrect layout of hmac_block");
+
+struct fast_psk_context {
+    union hmac_block inner, outer;
+    uint32_t sum[SHA1_OUTPUT_SZ_WORDS];
+};
+
+static inline void sha1_setup(void)
+{
+    esp_sha_acquire_hardware();
+}
+
+static inline void sha1_teardown(void)
+{
+    esp_sha_release_hardware();
+}
+
+static void pad_blocks(union hmac_block *ctx, size_t len)
+{
+    size_t bits = len << 3;
+    uint8_t *bytes = ctx->whole_bytes;
+    bytes[len] = 0x80;
+    // Set all remaining bytes to 0
+    memset(&bytes[len + 1], 0, FAST_PSK_SHA1_BLOCKS_BUF_BYTES - (len + 1));
+
+    /*
+     * Simplified PUT_UINT64_BE(bits, bytes, FAST_PSK_SHA1_BLOCKS_BUF_BYTES - 8).
+     * Since len < 128 => bits < 1024, we only need to update the two least significant
+     * bytes, actually.
+     */
+    PUT_UINT32_BE(bits, bytes, FAST_PSK_SHA1_BLOCKS_BUF_BYTES - 4);
+}
+
+void sha1_op(uint32_t blocks[FAST_PSK_SHA1_BLOCKS_BUF_WORDS], uint32_t output[SHA1_OUTPUT_SZ_WORDS])
+{
+    sha_hal_hash_block(SHA1, blocks, SHA1_BLOCK_SZ_WORDS, true);
+    sha_hal_hash_block(SHA1, &blocks[SHA1_BLOCK_SZ_WORDS], SHA1_BLOCK_SZ_WORDS, false);
+    sha_hal_read_digest(SHA1, output);
+}
+
+void fast_psk_f(const char *password, size_t password_len, const uint8_t *ssid, size_t ssid_len, uint32_t count, uint8_t digest[SHA1_OUTPUT_SZ])
+{
+    struct fast_psk_context ctx_, *ctx = &ctx_;
+    size_t i;
+
+    memset(ctx, 0, sizeof(*ctx));
+
+    memset(ctx->outer.block[0].bytes, 0x5c, SHA1_BLOCK_SZ);
+    memset(ctx->inner.block[0].bytes, 0x36, SHA1_BLOCK_SZ);
+
+    for (i = 0; i < password_len; ++i) {
+        ctx->outer.block[0].bytes[i] ^= password[i];
+        ctx->inner.block[0].bytes[i] ^= password[i];
+    }
+
+    // U1 = PRF(P, S || i)
+    memcpy(ctx->inner.block[1].bytes, ssid, ssid_len);
+    PUT_UINT32_BE(count, ctx->inner.block[1].bytes, ssid_len);
+    pad_blocks(&ctx->inner, SHA1_BLOCK_SZ + ssid_len + 4);
+
+    sha1_setup();
+
+    uint32_t *pi, *po;
+    pi = ctx->inner.whole_words;
+    po = ctx->outer.whole_words;
+    // T1 = SHA1(K ^ ipad, S || i)
+    sha1_op(pi, ctx->outer.block[1].words);
+
+    // U1 = SHA1(K ^ opad, T1)
+    pad_blocks(&ctx->outer, SHA1_BLOCK_SZ + SHA1_OUTPUT_SZ);
+    uint32_t *inner_blk1 = ctx->inner.block[1].words;
+    uint32_t *outer_blk1 = ctx->outer.block[1].words;
+    uint32_t *sum = ctx->sum;
+
+    sha1_op(po, inner_blk1);
+    memcpy(sum, inner_blk1, SHA1_OUTPUT_SZ);
+    pad_blocks(&ctx->inner, SHA1_BLOCK_SZ + SHA1_OUTPUT_SZ);
+
+    for (i = 1; i < 4096; ++i) {
+        // Tn = SHA1(K ^ ipad, Un-1)
+        sha1_op(pi, outer_blk1);
+
+        // Un = SHA1(K ^ opad, Tn)
+        sha1_op(po, inner_blk1);
+
+        // F = U1 ^ U2 ^ ... Un
+        for (size_t j = 0; j < SHA1_OUTPUT_SZ_WORDS; ++j) {
+            sum[j] ^= inner_blk1[j];
+        }
+    }
+
+    sha1_teardown();
+
+    memcpy(digest, sum, SHA1_OUTPUT_SZ);
+
+    memset(ctx, 0, sizeof(*ctx));
+}
+
+int esp_fast_psk(const char *password, size_t password_len, const uint8_t *ssid, size_t ssid_len, size_t iterations, uint8_t *output, size_t output_len)
+{
+    if (!(ssid_len <= 32 && password_len <= 63 && iterations == 4096 && output_len == 32)) {
+        return -1;
+    }
+    fast_psk_f(password, password_len, ssid, ssid_len, 2, output);
+    memcpy(output + SHA1_OUTPUT_SZ, output, 32 - SHA1_OUTPUT_SZ);
+    fast_psk_f(password, password_len, ssid, ssid_len, 1, output);
+
+    return 0;
+}

components/wpa_supplicant/esp_supplicant/src/crypto/fastpsk.h

@@ -0,0 +1,32 @@
+/*
+ * SPDX-FileCopyrightText: 2015-2024 Espressif Systems (Shanghai) CO LTD
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <stddef.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * @brief Calculate PSK
+ *
+ * @param password       Password
+ * @param password_len   Length of password, it must be <= 63
+ * @param ssid           SSID
+ * @param ssid_len       Length of SSID, it must be <= 32
+ * @param iterations     Iterations of the PBKDF2-SHA1, this is a dummy param and it must be 4096
+ * @param output         Buffer for calculated PSK, it must be at least 32 bytes
+ * @param output_len     Length of output to return, this is a dummy param and it must be 32
+ * @return 0 on success, non-zero on failure
+ */
+int esp_fast_psk(const char *password, size_t password_len, const uint8_t *ssid, size_t ssid_len, size_t iterations, uint8_t *output, size_t output_len);
+
+#ifdef __cplusplus
+}
+#endif