/*
* Copyright (c) 2023, Christian Huitema
*
* Permission is hereby granted, free of charge, to any person obtaining a copy
* of this software and associated documentation files (the "Software"), to
* deal in the Software without restriction, including without limitation the
* rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
* sell copies of the Software, and to permit persons to whom the Software is
* furnished to do so, subject to the following conditions:
*
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
* IN THE SOFTWARE.
*/

#ifdef _WINDOWS
#include "wincompat.h"
#endif

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <picotls.h>
#include <mbedtls/mbedtls_config.h>
#include <mbedtls/build_info.h>
#include <mbedtls/pk.h>
#include <mbedtls/pem.h>
#include <mbedtls/error.h>
#include <psa/crypto.h>
#include <psa/crypto_struct.h>
#include <psa/crypto_values.h>

typedef struct st_ptls_mbedtls_signature_scheme_t {
    uint16_t scheme_id;
    psa_algorithm_t hash_algo;
} ptls_mbedtls_signature_scheme_t;

typedef struct st_ptls_mbedtls_sign_certificate_t {
    ptls_sign_certificate_t super;
    mbedtls_svc_key_id_t key_id;
    psa_key_attributes_t attributes;
    const ptls_mbedtls_signature_scheme_t * schemes;
} ptls_mbedtls_sign_certificate_t;

static const unsigned char ptls_mbedtls_oid_ec_key[] = { 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x02, 0x01 };
static const unsigned char ptls_mbedtls_oid_rsa_key[] = { 0x2a, 0x86, 0x48, 0x86, 0xf7, 0x0d, 0x01, 0x01, 0x01 };
static const unsigned char ptls_mbedtls_oid_ed25519[] = { 0x2b, 0x65, 0x70 };

static const ptls_mbedtls_signature_scheme_t rsa_signature_schemes[] = {
    {PTLS_SIGNATURE_RSA_PSS_RSAE_SHA256, PSA_ALG_SHA_256},
    {PTLS_SIGNATURE_RSA_PSS_RSAE_SHA384, PSA_ALG_SHA_384},
    {PTLS_SIGNATURE_RSA_PSS_RSAE_SHA512, PSA_ALG_SHA_512},
    {UINT16_MAX, PSA_ALG_NONE}};
static const ptls_mbedtls_signature_scheme_t secp256r1_signature_schemes[] = {
    {PTLS_SIGNATURE_ECDSA_SECP256R1_SHA256, PSA_ALG_SHA_256}, {UINT16_MAX, PSA_ALG_NONE}};
static const ptls_mbedtls_signature_scheme_t secp384r1_signature_schemes[] = {
    {PTLS_SIGNATURE_ECDSA_SECP384R1_SHA384, PSA_ALG_SHA_384}, {UINT16_MAX, PSA_ALG_NONE}};
static const ptls_mbedtls_signature_scheme_t secp521r1_signature_schemes[] = {
    {PTLS_SIGNATURE_ECDSA_SECP521R1_SHA512, PSA_ALG_SHA_512}, {UINT16_MAX, PSA_ALG_NONE}};
static const ptls_mbedtls_signature_scheme_t ed25519_signature_schemes[] = {
    {PTLS_SIGNATURE_ED25519, PSA_ALG_NONE}, {UINT16_MAX, PSA_ALG_NONE}};

#if defined(MBEDTLS_PEM_PARSE_C)

static int ptls_mbedtls_parse_der_length(const unsigned char* pem_buf, size_t pem_len, size_t* px, size_t *pl)
{
    int ret = 0;
    size_t x = *px;
    size_t l = pem_buf[x++];

    if (l > 128) {
        size_t ll = l & 0x7F;
        l = 0;
        while (ll > 0 && x + l < pem_len) {
            l *= 256;
            l += pem_buf[x++];
            ll--;
        }
    }

    *pl = l;
    *px = x;

    return ret;
}

static int ptls_mbedtls_parse_ecdsa_field(const unsigned char* pem_buf, size_t pem_len, size_t* key_index, size_t* key_length)
{
    int ret = 0;
    size_t x = 0;

    // const unsigned char head = { 0x30, l-2, 0x02, 0x01, 0x01, 0x04 }
    if (pem_len < 16 ||
        pem_buf[x++] != 0x30 /* type = sequence */)
    {
        ret = -1;
    }
    else {
        size_t l = 0;
        ret = ptls_mbedtls_parse_der_length(pem_buf, pem_len, &x, &l);

        if (x + l != pem_len) {
            ret = -1;
        }
    }
    if (ret == 0){
        if (pem_buf[x++] != 0x02 /* type = int */ ||
            pem_buf[x++] != 0x01 /* length of int = 1 */ ||
            pem_buf[x++] != 0x01 /* version = 1 */ ||
            pem_buf[x++] != 0x04 /*octet string */ ||
            pem_buf[x] + x >= pem_len) {
            ret = -1;
        }
        else {
            *key_index = x + 1;
            *key_length = pem_buf[x];
            x += 1 + pem_buf[x];

            if (x < pem_len && pem_buf[x] == 0xa0) {
                /* decode the EC parameters, identify the curve */
                x++;
                if (x + pem_buf[x] >= pem_len) {
                    /* EC parameters extend beyond buffer */
                    ret = -1;
                }
                else {
                    x += pem_buf[x] + 1;
                }
            }

            if (ret == 0 && x < pem_len) {
                /* skip the public key parameter */
                if (pem_buf[x++] != 0xa1 ||
                    x >= pem_len) {
                    ret = -1;
                }
                else {
                    size_t l = 0;
                    ret = ptls_mbedtls_parse_der_length(pem_buf, pem_len, &x, &l);
                    x += l;
                }
            }

            if (x != pem_len) {
                ret = -1;
            }
        }
    }
    return ret;
}

/* On input, key_index points at the "key information" in a
* "private key" message. For EDDSA, this contains an
* octet string carrying the key itself. On return, key index
* and key length are updated to point at the key field.
*/
static int ptls_mbedtls_parse_eddsa_key(const unsigned char* pem_buf, size_t pem_len,
    size_t* key_index, size_t* key_length)
{
    int ret = 0;
    size_t x = *key_index;
    size_t l_key = 0;

    if (*key_length < 2 || pem_buf[x++] != 0x04) {
        ret = -1;
    } else {
        ret = ptls_mbedtls_parse_der_length(pem_buf, pem_len, &x, &l_key);
        if (x + l_key != *key_index + *key_length) {
            ret = -1;
        }
        else {
            *key_index = x;
            *key_length = l_key;
        }
    }
    return ret;
}

/* If using PKCS8 encoding, the "private key" field contains the
* same "ecdsa field" found in PEM "EC PRIVATE KEY" files. We
* use the same parser, but we need to reset indices so they
* reflect the unwrapped key.
*/
int ptls_mbedtls_parse_ec_private_key(const unsigned char* pem_buf, size_t pem_len,
    size_t* key_index, size_t* key_length)
{
    size_t x_offset = 0;
    size_t x_len = 0;
    int ret = ptls_mbedtls_parse_ecdsa_field(pem_buf + *key_index, *key_length, &x_offset, &x_len);

    if (ret == 0) {
        *key_index += x_offset;
        *key_length = x_len;
    }
    return ret;
}

int test_parse_private_key_field(const unsigned char* pem_buf, size_t pem_len,
    size_t* oid_index, size_t *oid_length,
    size_t* key_index, size_t* key_length)
{
    int ret = 0;
    size_t l_oid = 0;
    size_t x_oid = 0;
    size_t l_key = 0;
    size_t x_key = 0;

    size_t x = 0;
    /*  const unsigned char head = {0x30, l - 2, 0x02, 0x01, 0x00} */
    if (pem_len < 16 ||
        pem_buf[x++] != 0x30 /* type = sequence */)
    {
        ret = -1;
    }
    else {
        size_t l = 0;
        ret = ptls_mbedtls_parse_der_length(pem_buf, pem_len, &x, &l);

        if (x + l != pem_len) {
            ret = -1;
        }
    }
    if (ret == 0) {
        if (pem_buf[x++] != 0x02 /* type = int */ ||
            pem_buf[x++] != 0x01 /* length of int = 1 */ ||
            pem_buf[x++] != 0x00 /* version = 0 */ ||
            pem_buf[x++] != 0x30 /* sequence */){
            ret = -1;
        }
        else {
            /* the sequence contains the OID and optional key attributes,
            * which we ignore for now.
            */
            size_t l_seq = 0;
            size_t x_seq;
            ret = ptls_mbedtls_parse_der_length(pem_buf, pem_len, &x, &l_seq);
            x_seq = x;
            if (x + l_seq >= pem_len ||
                pem_buf[x++] != 0x06) {
                ret = -1;
            }
            else {
                l_oid = pem_buf[x++];
                x_oid = x;
                if (x + l_oid > x_seq + l_seq) {
                    ret = -1;
                }
                else {
                    x = x_seq + l_seq;
                }
            }
        }
    }
    if (ret == 0) {
        /* At that point the oid has been identified.
        * The next parameter is an octet string containing the key info.
        */
        if (x + 2 > pem_len ||
            pem_buf[x++]  != 0x04){
            ret = -1;
        }
        else {
            ret = ptls_mbedtls_parse_der_length(pem_buf, pem_len, &x, &l_key);
            x_key = x;
            x += l_key;
            if (x > pem_len) {
                ret = -1;
            }
        }
    }
    *oid_index = x_oid;
    *oid_length = l_oid;
    *key_index = x_key;
    *key_length = l_key;

    return ret;
}

int ptls_mbedtls_get_der_key(mbedtls_pem_context* pem,
    mbedtls_pk_type_t * pk_type,
    const unsigned char* key, size_t keylen,
    const unsigned char* pwd, size_t pwdlen,
    int (*f_rng)(void*, unsigned char*, size_t), void* p_rng)
{
    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
#if defined(MBEDTLS_PEM_PARSE_C)
    size_t len;
#endif

    if (keylen == 0) {
        return MBEDTLS_ERR_PK_KEY_INVALID_FORMAT;
    }

    mbedtls_pem_init(pem);

#if defined(MBEDTLS_RSA_C)
    /* Avoid calling mbedtls_pem_read_buffer() on non-null-terminated string */
    if (key[keylen - 1] != '\0') {
        ret = MBEDTLS_ERR_PEM_NO_HEADER_FOOTER_PRESENT;
    }
    else {
        ret = mbedtls_pem_read_buffer(pem,
            "-----BEGIN RSA PRIVATE KEY-----",
            "-----END RSA PRIVATE KEY-----",
            key, pwd, pwdlen, &len);
    }

    if (ret == 0) {
        * pk_type = MBEDTLS_PK_RSA;
        return ret;
    }
    else if (ret == MBEDTLS_ERR_PEM_PASSWORD_MISMATCH) {
        return MBEDTLS_ERR_PK_PASSWORD_MISMATCH;
    }
    else if (ret == MBEDTLS_ERR_PEM_PASSWORD_REQUIRED) {
        return MBEDTLS_ERR_PK_PASSWORD_REQUIRED;
    }
    else if (ret != MBEDTLS_ERR_PEM_NO_HEADER_FOOTER_PRESENT) {
        return ret;
    }
#endif /* MBEDTLS_RSA_C */

#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
    /* Avoid calling mbedtls_pem_read_buffer() on non-null-terminated string */
    if (key[keylen - 1] != '\0') {
        ret = MBEDTLS_ERR_PEM_NO_HEADER_FOOTER_PRESENT;
    }
    else {
        ret = mbedtls_pem_read_buffer(pem,
            "-----BEGIN EC PRIVATE KEY-----",
            "-----END EC PRIVATE KEY-----",
            key, pwd, pwdlen, &len);
    }
    if (ret == 0) {
        * pk_type = MBEDTLS_PK_ECKEY;
        return ret;
    }
    else if (ret == MBEDTLS_ERR_PEM_PASSWORD_MISMATCH) {
        return MBEDTLS_ERR_PK_PASSWORD_MISMATCH;
    }
    else if (ret == MBEDTLS_ERR_PEM_PASSWORD_REQUIRED) {
        return MBEDTLS_ERR_PK_PASSWORD_REQUIRED;
    }
    else if (ret != MBEDTLS_ERR_PEM_NO_HEADER_FOOTER_PRESENT) {
        return ret;
    }
#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */

    /* Avoid calling mbedtls_pem_read_buffer() on non-null-terminated string */
    if (key[keylen - 1] != '\0') {
        ret = MBEDTLS_ERR_PEM_NO_HEADER_FOOTER_PRESENT;
    }
    else {
        ret = mbedtls_pem_read_buffer(pem,
            "-----BEGIN PRIVATE KEY-----",
            "-----END PRIVATE KEY-----",
            key, NULL, 0, &len);
        if (ret == 0) {
            /* info is unknown */
            return ret;
        }
        else if (ret != MBEDTLS_ERR_PEM_NO_HEADER_FOOTER_PRESENT) {
            return ret;
        }
    }

#if defined(MBEDTLS_PKCS12_C) || defined(MBEDTLS_PKCS5_C)
    /* Avoid calling mbedtls_pem_read_buffer() on non-null-terminated string */
    if (key[keylen - 1] != '\0') {
        ret = MBEDTLS_ERR_PEM_NO_HEADER_FOOTER_PRESENT;
    }
    else {
        ret = mbedtls_pem_read_buffer(pem,
            "-----BEGIN ENCRYPTED PRIVATE KEY-----",
            "-----END ENCRYPTED PRIVATE KEY-----",
            key, NULL, 0, &len);
    }
    if (ret == 0) {
        /* infor is unknown */
        return ret;
    }
    else if (ret != MBEDTLS_ERR_PEM_NO_HEADER_FOOTER_PRESENT) {
        return ret;
    }
#endif /* MBEDTLS_PKCS12_C || MBEDTLS_PKCS5_C */
    return MBEDTLS_ERR_PK_KEY_INVALID_FORMAT;
}
#endif

const ptls_mbedtls_signature_scheme_t* ptls_mbedtls_select_signature_scheme(
    const ptls_mbedtls_signature_scheme_t *available,
    const uint16_t *algorithms, size_t num_algorithms)
{
    const ptls_mbedtls_signature_scheme_t* scheme;
    /* select the algorithm, driven by server-isde preference of `available` */
    for (scheme = available; scheme->scheme_id != UINT16_MAX; ++scheme) {
        for (size_t i = 0; i != num_algorithms; ++i) {
            if (algorithms[i] == scheme->scheme_id) {
                return scheme;
            }
        }
    }
    return NULL;
}

int ptls_mbedtls_set_available_schemes(
    ptls_mbedtls_sign_certificate_t* signer)
{
    int ret = 0;
    psa_algorithm_t algo = psa_get_key_algorithm(&signer->attributes);
    size_t nb_bits = psa_get_key_bits(&signer->attributes);

    switch (algo) {
    case PSA_ALG_RSA_PKCS1V15_SIGN_RAW:
        signer->schemes = rsa_signature_schemes;
        break;
    case PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_256):
        signer->schemes = secp256r1_signature_schemes;
        break;
    case PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_384):
        signer->schemes = secp384r1_signature_schemes;
        break;
    case PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_512):
        signer->schemes = secp521r1_signature_schemes;
        break;
    case PSA_ALG_ECDSA_BASE:
        switch (nb_bits) {
        case 521:
            signer->schemes = secp521r1_signature_schemes;
            break;
        case 384:
            signer->schemes = secp384r1_signature_schemes;
            break;
        case 256:
            signer->schemes = secp256r1_signature_schemes;
            break;
        default:
            signer->schemes = secp256r1_signature_schemes;
            ret = -1;
            break;
        }
        break;
    case PSA_ALG_ED25519PH:
        signer->schemes = ed25519_signature_schemes;
        break;
    default:
        printf("Unknown algo: %x\n", algo);
        ret = -1;
    }

    return ret;
}

/*
* Sign a certificate
* - step1, selected a signature algorithm compatible with the public key algorithm
*   and with the list specified by the application.
* - step2, compute the hash with the specified algorithm.
* - step3, compute the signature of the hash using psa_sign_hash.
* 
* In the case of RSA, we use the algorithm PSA_ALG_RSA_PKCS1V15_SIGN_RAW, which
* pads the hash according to PKCS1V15 before doing the private key operation.
* The implementation of RSA/PKCS1V15 also includes a verification step to protect
* against key attacks through partial faults.
* 
* MBEDTLS has a "psa_sign_message" that combines step2 and step3. However, it
* requires specifying an algorithm type that exactly specifies the signature
* algorithm, such as "RSA with SHA384". This is not compatible with the
* "RSA sign raw" algorithm. Instead, we decompose the operation in two steps.
* There is no performance penalty doing so, as "psa_sign_message" is only
* a convenience API.
*/

int ptls_mbedtls_sign_certificate(ptls_sign_certificate_t* _self, ptls_t* tls,
    ptls_async_job_t** async, uint16_t* selected_algorithm,
    ptls_buffer_t* outbuf, ptls_iovec_t input, const uint16_t* algorithms, size_t num_algorithms)
{
    int ret = 0;
    ptls_mbedtls_sign_certificate_t* self = (ptls_mbedtls_sign_certificate_t*)
        (((unsigned char*)_self) - offsetof(struct st_ptls_mbedtls_sign_certificate_t, super));
    /* First, find the set of compatible algorithms */
    const ptls_mbedtls_signature_scheme_t* scheme =
        ptls_mbedtls_select_signature_scheme(self->schemes, algorithms, num_algorithms);

    if (scheme == NULL) {
        ret = PTLS_ERROR_INCOMPATIBLE_KEY;
    }
    else {
        /* First prepare the hash */
        unsigned char hash_buffer[PTLS_MAX_DIGEST_SIZE];
        unsigned char* hash_value = NULL;
        size_t hash_length = 0;

        if (scheme->hash_algo == PSA_ALG_NONE) {
            hash_value = input.base;
            hash_length = input.len;
        }
        else {
            if (psa_hash_compute(scheme->hash_algo, input.base, input.len, hash_buffer, PTLS_MAX_DIGEST_SIZE, &hash_length) != PSA_SUCCESS) {
                ret = PTLS_ERROR_NOT_AVAILABLE;
            }
            else {
                hash_value = hash_buffer;
            }
        }
        if (ret == 0) {
            psa_algorithm_t sign_algo = psa_get_key_algorithm(&self->attributes);
            size_t nb_bits = psa_get_key_bits(&self->attributes);
            size_t nb_bytes = (nb_bits + 7) / 8;
            if (nb_bits == 0) {
                if (sign_algo == PSA_ALG_RSA_PKCS1V15_SIGN_RAW) {
                    /* assume at most 4096 bit key */
                    nb_bytes = 512;
                }
                else {
                    /* Max size assumed, secp521r1 */
                    nb_bytes = 124;
                }
            } else if (sign_algo != PSA_ALG_RSA_PKCS1V15_SIGN_RAW) {
                nb_bytes *= 2;
            }
            if ((ret = ptls_buffer_reserve(outbuf, nb_bytes)) == 0) {
                size_t signature_length = 0;

                if (psa_sign_hash(self->key_id, sign_algo, hash_value, hash_length,
                    outbuf->base + outbuf->off, nb_bytes, &signature_length) != 0) {
                    ret = PTLS_ERROR_INCOMPATIBLE_KEY;
                }
                else {
                    outbuf->off += signature_length;
                }
            }
        }
    }
    return ret;
}

void ptls_mbedtls_dispose_sign_certificate(ptls_sign_certificate_t *_self)
{
    if (_self != NULL) {
        ptls_mbedtls_sign_certificate_t* self = (ptls_mbedtls_sign_certificate_t*)
        (((unsigned char*)_self) - offsetof(struct st_ptls_mbedtls_sign_certificate_t, super));
        /* Destroy the key */
        psa_destroy_key(self->key_id);
        psa_reset_key_attributes(&self->attributes);
        memset(self, 0, sizeof(ptls_mbedtls_sign_certificate_t));
        free(self);
    }
}
/*
* An RSa key is encoded in DER as:
* RSAPrivateKey ::= SEQUENCE {
*   version             INTEGER,  -- must be 0
*   modulus             INTEGER,  -- n
*   publicExponent      INTEGER,  -- e
*   privateExponent     INTEGER,  -- d
*   prime1              INTEGER,  -- p
*   prime2              INTEGER,  -- q
*   exponent1           INTEGER,  -- d mod (p-1)
*   exponent2           INTEGER,  -- d mod (q-1)
*   coefficient         INTEGER,  -- (inverse of q) mod p
* }
* 
* The number of key bits is the size in bits of the integer N.
* We must decode the length in octets of the integer representation,
* then subtract the number of zeros at the beginning of the data.
*/
int ptls_mbedtls_rsa_get_key_bits(const unsigned char* key_value, size_t key_length, size_t * p_nb_bits)
{
    int ret = 0;
    size_t nb_bytes = 0;
    size_t nb_bits = 0;
    size_t x = 0;

    if (key_length > 16 && key_value[x++] == 0x30) {
        /* get the length of the sequence. */
        size_t l = 0;
        ret = ptls_mbedtls_parse_der_length(key_value, key_length, &x, &l);

        if (x + l != key_length) {
            ret = -1;
        }
    }

    if (ret == 0 &&
        key_value[x] == 0x02 &&
        key_value[x + 1] == 0x01 &&
        key_value[x + 2] == 0x00 &&
        key_value[x + 3] == 0x02) {
        x += 4;
        ret = ptls_mbedtls_parse_der_length(key_value, key_length, &x, &nb_bytes);
    }
    else {
        ret = -1;
    }

    if (ret == 0) {
        unsigned char v = key_value[x];
        nb_bits = 8 * nb_bytes;

        if (v == 0) {
            nb_bits -= 8;
        }
        else {
            while ((v & 0x80) == 0) {
                nb_bits--;
                v <<= 1;
            }
        }
    }
    *p_nb_bits = nb_bits;
    return ret;
}

void ptls_mbedtls_set_rsa_key_attributes(ptls_mbedtls_sign_certificate_t* signer,
    const unsigned char * key_value, size_t key_length)
{
    size_t nb_bits = 0;
    psa_set_key_usage_flags(&signer->attributes, PSA_KEY_USAGE_SIGN_HASH);
    psa_set_key_algorithm(&signer->attributes, PSA_ALG_RSA_PKCS1V15_SIGN_RAW);
    psa_set_key_type(&signer->attributes, PSA_KEY_TYPE_RSA_KEY_PAIR);
    if (ptls_mbedtls_rsa_get_key_bits(key_value, key_length, &nb_bits) == 0) {
        psa_set_key_bits(&signer->attributes, nb_bits);
    }
}

int ptls_mbedtls_set_ec_key_attributes(ptls_mbedtls_sign_certificate_t* signer, size_t key_length)
{
    int ret = 0;

    psa_set_key_usage_flags(&signer->attributes, PSA_KEY_USAGE_SIGN_HASH);
    psa_set_key_algorithm(&signer->attributes, PSA_ALG_ECDSA_BASE);
    psa_set_key_type(&signer->attributes, PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1));
    if (key_length == 32) {
        psa_set_key_algorithm(&signer->attributes,
            PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_256));
        psa_set_key_bits(&signer->attributes, 256);
    }
    else if (key_length == 48) {
        psa_set_key_algorithm(&signer->attributes,
            PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_384));
        psa_set_key_bits(&signer->attributes, 384);
    }
    else if (key_length == 66) {
        psa_set_key_algorithm(&signer->attributes,
            PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_SHA_512));
        psa_set_key_bits(&signer->attributes, 521);
    }
    else {
        ret = -1;
    }

    return ret;
}


int ptls_mbedtls_load_private_key(ptls_context_t* ctx, char const* pem_fname)
{
    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
    size_t n;
    unsigned char* buf;
    mbedtls_pem_context pem = { 0 };
    mbedtls_pk_type_t pk_type = 0;
    /* mbedtls_svc_key_id_t key_id = 0; */
    size_t key_length = 0;
    size_t key_index = 0;
    ptls_mbedtls_sign_certificate_t* signer = (ptls_mbedtls_sign_certificate_t*)malloc(sizeof(ptls_mbedtls_sign_certificate_t));

    if (signer == NULL) {
        return(PTLS_ERROR_NO_MEMORY);
    }
    memset(signer, 0, sizeof(ptls_mbedtls_sign_certificate_t));
    signer->attributes = psa_key_attributes_init();

    if ((ret = mbedtls_pk_load_file(pem_fname, &buf, &n)) != 0) {
        if (ret == MBEDTLS_ERR_PK_ALLOC_FAILED) {
            return(PTLS_ERROR_NO_MEMORY);
        }
        else {
            return(PTLS_ERROR_NOT_AVAILABLE);
        }
    }
    ret = ptls_mbedtls_get_der_key(&pem, &pk_type, buf, n, NULL, 0, NULL, NULL);

    /* We cannot use the platform API:
    mbedtls_zeroize_and_free(buf, n);
    so we do our own thing.
    */
    memset(buf, 0, n);
    free(buf);

    if (ret == 0) {
        if (pk_type == MBEDTLS_PK_RSA) {
            key_length = pem.private_buflen;
            ptls_mbedtls_set_rsa_key_attributes(signer, pem.private_buf, key_length);
        }
        else if (pk_type == MBEDTLS_PK_ECKEY) {
            ret = ptls_mbedtls_parse_ecdsa_field(pem.private_buf, pem.private_buflen, &key_index, &key_length);
            if (ret == 0) {
                ret = ptls_mbedtls_set_ec_key_attributes(signer, key_length);
            }
        }
        else if (pk_type == MBEDTLS_PK_NONE) {
            /* TODO: not clear whether MBDED TLS supports ED25519 yet. Probably not. */
            /* Should have option to encode RSA or ECDSA using PKCS8 */
            size_t oid_index = 0;
            size_t oid_length = 0;

            psa_set_key_usage_flags(&signer->attributes, PSA_KEY_USAGE_SIGN_HASH);
            ret = test_parse_private_key_field(pem.private_buf, pem.private_buflen, &oid_index, &oid_length, &key_index, &key_length);
            if (ret == 0) {
                /* need to parse the OID in order to set the parameters */

                if (oid_length == sizeof(ptls_mbedtls_oid_ec_key) &&
                    memcmp(pem.private_buf + oid_index, ptls_mbedtls_oid_ec_key, sizeof(ptls_mbedtls_oid_ec_key)) == 0) {
                    ret = ptls_mbedtls_parse_ec_private_key(pem.private_buf, pem.private_buflen, &key_index, &key_length);
                    if (ret == 0) {
                        ret = ptls_mbedtls_set_ec_key_attributes(signer, key_length);
                    }
                }
                else if (oid_length == sizeof(ptls_mbedtls_oid_ed25519) &&
                    memcmp(pem.private_buf + oid_index, ptls_mbedtls_oid_ed25519, sizeof(ptls_mbedtls_oid_ed25519)) == 0) {
                    /* We recognized ED25519 -- PSA_ECC_FAMILY_TWISTED_EDWARDS -- PSA_ALG_ED25519PH */
                    psa_set_key_algorithm(&signer->attributes, PSA_ALG_PURE_EDDSA);
                    psa_set_key_type(&signer->attributes, PSA_ECC_FAMILY_TWISTED_EDWARDS);
                    ret = ptls_mbedtls_parse_eddsa_key(pem.private_buf, pem.private_buflen, &key_index, &key_length);
                    psa_set_key_bits(&signer->attributes, 256);
                }
                else if (oid_length == sizeof(ptls_mbedtls_oid_rsa_key) &&
                    memcmp(pem.private_buf + oid_index, ptls_mbedtls_oid_rsa_key, sizeof(ptls_mbedtls_oid_rsa_key)) == 0) {
                    /* We recognized RSA */
                    key_length = pem.private_buflen;
                    ptls_mbedtls_set_rsa_key_attributes(signer, pem.private_buf, key_length);
                }
                else {
                    ret = PTLS_ERROR_NOT_AVAILABLE;
                }
            }
        }
        else {
            ret = -1;
        }

        if (ret == 0) {
            /* Now that we have the DER or bytes for the key, try import into PSA */
            psa_status_t status = psa_import_key(&signer->attributes, pem.private_buf + key_index, key_length, &signer->key_id);

            if (status != PSA_SUCCESS) {
                ret = -1;
            }
            else {
                ret = ptls_mbedtls_set_available_schemes(signer);
            }
        }
        /* Free the PEM buffer */
        mbedtls_pem_free(&pem);
    }
    if (ret == 0) {
        signer->super.cb = ptls_mbedtls_sign_certificate;
        ctx->sign_certificate = &signer->super;
    } else {
        /* Dispose of what we have allocated. */
        ptls_mbedtls_dispose_sign_certificate(&signer->super);
    }
    return ret;
}