ECC key derivation and agreement

Hey,

I am trying to implement ecdh key agreement and derivation, but I am getting PSA_ERROR_INVALID_ARGUMENT (-135) from psa_key_derivation_key_agreement:
Any idea what is causing this, because when I was testing with psa_raw_key_agreement it succeeded.

The test setup was based on crypto ecdh sample and sdk is v3.2.1.

main.c:

/*
 * Copyright (c) 2021 Nordic Semiconductor ASA
 *
 * SPDX-License-Identifier: LicenseRef-Nordic-5-Clause
 */

#include <psa/crypto.h>
#include <psa/crypto_extra.h>
#include <stdio.h>
#include <stdlib.h>
#include <zephyr/kernel.h>
#include <zephyr/logging/log.h>
#include <zephyr/sys/printk.h>

LOG_MODULE_REGISTER(ecdh, LOG_LEVEL_DBG);

#ifdef CONFIG_BUILD_WITH_TFM
#include <tfm_ns_interface.h>
#endif

#include "./server_pub_key.h"

#define APP_SUCCESS (0)
#define APP_ERROR (-1)
#define APP_SUCCESS_MESSAGE "Example finished successfully!"
#define APP_ERROR_MESSAGE "Example exited with error!"

#define PRINT_HEX(p_label, p_text, len)               \
  ({                                                  \
    LOG_INF("---- %s (len: %u): ----", p_label, len); \
    LOG_HEXDUMP_INF(p_text, len, "Content:");         \
    LOG_INF("---- %s end  ----", p_label);            \
  })

#define PRINT_HEXSTRING(p_label, p_data, len)                                \
  ({                                                                         \
    LOG_INF("---- %s (len: %u, str_len: %u) ----", p_label, len, len * 2);   \
    uint8_t string[len * 2 + 1];                                             \
    for (size_t i = 0; i < len; i++) {                                       \
      snprintf(string + 2 * i, (sizeof(string) - 2 * i), "%02x", p_data[i]); \
    }                                                                        \
    LOG_INF("%s", string);                                                   \
    LOG_INF("---- %s end ----", p_label);                                    \
  })

#define KEY_PAIR_TYPE PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1)
#define KEY_SIZE (384)
#define PUB_KEY_TYPE PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(KEY_PAIR_TYPE)
#define PUB_KEY_SIZE PSA_EXPORT_PUBLIC_KEY_OUTPUT_SIZE(PUB_KEY_TYPE, KEY_SIZE)
#define RAW_KEY_AGREEMENT_SIZE PSA_RAW_KEY_AGREEMENT_OUTPUT_SIZE(KEY_PAIR_TYPE, KEY_SIZE)
#define AGREEMENT_ALG PSA_ALG_ECDH
#define DERIVATION_ALG PSA_ALG_HKDF(PSA_ALG_SHA_256)

int crypto_init(void) {
  psa_status_t status;

  /* Initialize PSA Crypto */
  status = psa_crypto_init();
  if (status != PSA_SUCCESS)
    return APP_ERROR;

  return APP_SUCCESS;
}

mbedtls_svc_key_id_t crypto_gen_key() {
  int err;
  mbedtls_svc_key_id_t new_key;
  psa_key_attributes_t attrib = psa_key_attributes_init();
  psa_set_key_algorithm(&attrib, PSA_ALG_KEY_AGREEMENT(AGREEMENT_ALG, DERIVATION_ALG));
  psa_set_key_type(&attrib, KEY_PAIR_TYPE);
  psa_set_key_bits(&attrib, KEY_SIZE);
  psa_set_key_usage_flags(&attrib, PSA_KEY_USAGE_DERIVE);

  err = psa_generate_key(&attrib, &new_key);
  if (err != PSA_SUCCESS) {
    LOG_ERR("Failed to create key pair, %d", err);
    return 0;
  }

  return new_key;
}

mbedtls_svc_key_id_t crypto_derive_enc_key(mbedtls_svc_key_id_t private_key) {
  int err;
  psa_key_derivation_operation_t op =
      psa_key_derivation_operation_init();
  psa_key_attributes_t attrib = psa_key_attributes_init();
  mbedtls_svc_key_id_t enc_key = 0;

  err = psa_key_derivation_setup(&op, PSA_ALG_KEY_AGREEMENT(
                                          AGREEMENT_ALG, DERIVATION_ALG));
  if (err != PSA_SUCCESS) {
    LOG_ERR("Failed to setup key derivation, %d", err);
    goto exit;
  }

  err = psa_key_derivation_key_agreement(
      &op, PSA_KEY_DERIVATION_INPUT_SECRET,
      private_key, pub_key, sizeof(pub_key));
  if (err != PSA_SUCCESS) {
    LOG_ERR("Failed to execute key agreement and derivation, %d", err);
    goto exit;
  }

  psa_set_key_algorithm(&attrib, PSA_ALG_CCM);
  psa_set_key_type(&attrib, PSA_KEY_TYPE_AES);
  psa_set_key_bits(&attrib, 256);
  psa_set_key_usage_flags(&attrib,
                          PSA_KEY_USAGE_DECRYPT |
                              PSA_KEY_USAGE_ENCRYPT |
                              PSA_KEY_USAGE_EXPORT);
  err = psa_key_derivation_output_key(&attrib, &op, &enc_key);
  if (err != PSA_SUCCESS) {
    LOG_ERR("Failed to generate AES key, %d", err);
    goto exit;
  }

exit:
  psa_reset_key_attributes(&attrib);
  psa_key_derivation_abort(&op);
  return enc_key;
}

int main(void) {
  int status;

  LOG_INF("Starting ECDH example...");
  /* Init crypto */
  status = crypto_init();
  if (status != APP_SUCCESS) {
    LOG_INF(APP_ERROR_MESSAGE);
    return APP_ERROR;
  }

  __auto_type my_key = crypto_gen_key();
  if (!my_key) return APP_ERROR;
  uint8_t shared_secret[RAW_KEY_AGREEMENT_SIZE];
  size_t out_size;
  status = psa_raw_key_agreement(PSA_ALG_ECDH, my_key, pub_key, sizeof(pub_key),
                                 shared_secret, sizeof(shared_secret), &out_size);
  if (status != PSA_SUCCESS) {
    LOG_ERR("Failed to create key agreement, %d", status);
    // return APP_ERROR;
  } else {
    PRINT_HEXSTRING("secret:", shared_secret, out_size);
    PRINT_HEX("secret:", shared_secret, out_size);
  }

  __auto_type enc_key = crypto_derive_enc_key(my_key);
  if (!enc_key) return APP_ERROR;
  uint8_t enc_key_bytes[PSA_EXPORT_KEY_OUTPUT_SIZE(PSA_KEY_TYPE_AES, 256)];
  size_t enc_export_size;
  status = psa_export_key(enc_key, enc_key_bytes, sizeof(enc_key_bytes), &enc_export_size);
  if (status != PSA_SUCCESS) {
    LOG_ERR("Failed to export AES key");
    return APP_ERROR;
  }
  PRINT_HEXSTRING("enc_key", enc_key_bytes, enc_export_size);
  PRINT_HEX("enc_key", enc_key_bytes, enc_export_size);

  uint8_t own_pub_key[PUB_KEY_SIZE];
  size_t export_size;
  status = psa_export_public_key(my_key, own_pub_key, sizeof(own_pub_key), &export_size);
  if (status != PSA_SUCCESS) {
    LOG_ERR("Failed to export public key, %d", status);
    return APP_ERROR;
  }
  PRINT_HEXSTRING("pubkey:", own_pub_key, export_size);
  PRINT_HEX("pubkey:", own_pub_key, export_size);

  LOG_INF(APP_SUCCESS_MESSAGE);

  return APP_SUCCESS;
}

server_pub_key.h:
#ifndef DC8A1BF7_016B_49E2_AAD2_025785B64CA1
#define DC8A1BF7_016B_49E2_AAD2_025785B64CA1

unsigned char pub_key[] = {
    0x04, 0xf1, 0x30, 0x2c, 0x19, 0x98, 0xc7, 0x8d, 0x44, 0x29, 0x95, 0x1e,
    0xc7, 0xde, 0x83, 0x25, 0xef, 0x26, 0x83, 0x93, 0x76, 0x8c, 0xdd, 0x68,
    0x6c, 0x4a, 0x05, 0x3a, 0x0d, 0x1f, 0x4f, 0xf5, 0x6f, 0x56, 0xea, 0x9f,
    0x4f, 0x5b, 0xf0, 0x36, 0x78, 0xb7, 0xa2, 0xac, 0x70, 0x86, 0x41, 0xb5,
    0xcd, 0x14, 0x1f, 0xd3, 0xda, 0x90, 0x9d, 0x6e, 0x37, 0x7a, 0xfc, 0xee,
    0x78, 0x89, 0x5a, 0x71, 0x90, 0x89, 0x08, 0x08, 0x45, 0xae, 0xca, 0xf2,
    0x8d, 0x81, 0xba, 0x4a, 0x2c, 0x32, 0x88, 0x24, 0xfb, 0xd2, 0xc5, 0xf0,
    0xf2, 0x05, 0x05, 0xd1, 0x2a, 0x06, 0x58, 0x62, 0x3f, 0x15, 0xdb, 0xe9,
    0x0d};

#endif /* DC8A1BF7_016B_49E2_AAD2_025785B64CA1 */

And conf files:
prj.conf:
#
# Copyright (c) 2024 Nordic Semiconductor ASA
#
# SPDX-License-Identifier: LicenseRef-Nordic-5-Clause
#
# The Zephyr CMSIS emulation assumes that ticks are ms, currently
CONFIG_SYS_CLOCK_TICKS_PER_SEC=1000

CONFIG_MAIN_STACK_SIZE=4096
CONFIG_HEAP_MEM_POOL_SIZE=4096

# Enable logging
CONFIG_CONSOLE=y
CONFIG_LOG=y

# Enable nRF Security backend for PSA Crypto API
CONFIG_NRF_SECURITY=y
CONFIG_MBEDTLS_PSA_CRYPTO_C=y

# Enable cryptographic features
CONFIG_PSA_WANT_GENERATE_RANDOM=y
CONFIG_PSA_WANT_ECC_SECP_R1_256=y
CONFIG_PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_GENERATE=y
CONFIG_PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_IMPORT=y
CONFIG_PSA_WANT_ALG_ECDH=y
CONFIG_PSA_WANT_ALG_HKDF=y
CONFIG_PSA_WANT_KEY_TYPE_AES=y
CONFIG_PSA_WANT_ALG_CCM=y

nrf9151dk_nrf9151_ns.conf:
#
# Copyright (c) 2024 Nordic Semiconductor ASA
#
# SPDX-License-Identifier: LicenseRef-Nordic-5-Clause
#
CONFIG_TFM_PROFILE_TYPE_NOT_SET=y
# CONFIG_TFM_PROFILE_TYPE_MINIMAL=y

# Using hardware crypto accelerator
CONFIG_PSA_CRYPTO_DRIVER_CC3XX=y

nrf9151dk_nrf9151.conf:
#
# Copyright (c) 2024 Nordic Semiconductor ASA
#
# SPDX-License-Identifier: LicenseRef-Nordic-5-Clause
#
# Enable both oberon driver and hardware crypto accelerator
# Key derivation is only supported in software driver but is using chained
# driver to accelerate sub-operations.
CONFIG_PSA_CRYPTO_DRIVER_OBERON=y
CONFIG_PSA_CRYPTO_DRIVER_CC3XX=y

# Mbedtls configuration
CONFIG_MBEDTLS_ENABLE_HEAP=y
CONFIG_MBEDTLS_HEAP_SIZE=8192

Related