/* This Source Code Form is subject to the terms of the Mozilla Public

 * License, v. 2.0. If a copy of the MPL was not distributed with this

 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include "seccomon.h"

#include "secerr.h"

#include "blapi.h"

#include "pkcs11i.h"

#include "softoken.h"

static CK_RV

jpake_mapStatus(SECStatus rv, CK_RV invalidArgsMapping) {

    int err;

    if (rv == SECSuccess)

        return CKR_OK;

    err = PORT_GetError();

    switch (err) {

        /* XXX: SEC_ERROR_INVALID_ARGS might be caused by invalid template

            parameters. */

        case SEC_ERROR_INVALID_ARGS:  return invalidArgsMapping;

        case SEC_ERROR_BAD_SIGNATURE: return CKR_SIGNATURE_INVALID;

        case SEC_ERROR_NO_MEMORY:     return CKR_HOST_MEMORY;

    }

    return CKR_FUNCTION_FAILED;

}

/* If key is not NULL then the gx value will be stored as an attribute with

   the type given by the gxAttrType parameter. */

static CK_RV

jpake_Sign(PLArenaPool * arena, const PQGParams * pqg, HASH_HashType hashType,

           const SECItem * signerID, const SECItem * x,

           CK_NSS_JPAKEPublicValue * out)

{

    SECItem gx, gv, r;

    CK_RV crv;

    PORT_Assert(arena != NULL);

    gx.data = NULL;

    gv.data = NULL;

    r.data = NULL;

    crv = jpake_mapStatus(JPAKE_Sign(arena, pqg, hashType, signerID, x, NULL,

                                     NULL, &gx, &gv, &r),

                          CKR_MECHANISM_PARAM_INVALID);

    if (crv == CKR_OK) {

        if ((out->pGX != NULL && out->ulGXLen >= gx.len) ||

            (out->pGV != NULL && out->ulGVLen >= gv.len) ||

            (out->pR  != NULL && out->ulRLen >= r.len)) {

            PORT_Memcpy(out->pGX, gx.data, gx.len); 

            PORT_Memcpy(out->pGV, gv.data, gv.len); 

            PORT_Memcpy(out->pR, r.data, r.len);

            out->ulGXLen = gx.len;

            out->ulGVLen = gv.len;

            out->ulRLen = r.len;

        } else {

            crv = CKR_MECHANISM_PARAM_INVALID;

        }

    } 

    return crv;

}

static CK_RV

jpake_Verify(PLArenaPool * arena, const PQGParams * pqg,

             HASH_HashType hashType, const SECItem * signerID,

             const CK_BYTE * peerIDData, CK_ULONG peerIDLen,

             const CK_NSS_JPAKEPublicValue * publicValueIn)

{

    SECItem peerID, gx, gv, r;

    peerID.data = (unsigned char *) peerIDData; peerID.len = peerIDLen;

    gx.data = publicValueIn->pGX; gx.len = publicValueIn->ulGXLen;

    gv.data = publicValueIn->pGV; gv.len = publicValueIn->ulGVLen;

    r.data = publicValueIn->pR;   r.len = publicValueIn->ulRLen;

    return jpake_mapStatus(JPAKE_Verify(arena, pqg, hashType, signerID, &peerID,

                                        &gx, &gv, &r),

                           CKR_MECHANISM_PARAM_INVALID);

}

#define NUM_ELEM(x) (sizeof (x) / sizeof (x)[0])

/* If the template has the key type set, ensure that it was set to the correct

 * value. If the template did not have the key type set, set it to the

 * correct value.

 */

static CK_RV

jpake_enforceKeyType(SFTKObject * key, CK_KEY_TYPE keyType) {

    CK_RV crv;

    SFTKAttribute * keyTypeAttr = sftk_FindAttribute(key, CKA_KEY_TYPE);

    if (keyTypeAttr != NULL) {

        crv = *(CK_KEY_TYPE *)keyTypeAttr->attrib.pValue == keyType

            ? CKR_OK

            : CKR_TEMPLATE_INCONSISTENT;

        sftk_FreeAttribute(keyTypeAttr);

    } else {

        crv = sftk_forceAttribute(key, CKA_KEY_TYPE, &keyType, sizeof keyType);

    }

    return crv;

}

static CK_RV

jpake_MultipleSecItem2Attribute(SFTKObject * key, const SFTKItemTemplate * attrs,

                                size_t attrsCount)

{

    size_t i;

    for (i = 0; i < attrsCount; ++i) {

        CK_RV crv = sftk_forceAttribute(key, attrs[i].type, attrs[i].item->data,

                                        attrs[i].item->len);

        if (crv != CKR_OK)

            return crv;

    }

    return CKR_OK;

}

CK_RV

jpake_Round1(HASH_HashType hashType, CK_NSS_JPAKERound1Params * params,

             SFTKObject * key)

{

    CK_RV crv;

    PQGParams pqg;

    PLArenaPool * arena;

    SECItem signerID;

    SFTKItemTemplate templateAttrs[] = {

        { CKA_PRIME, &pqg.prime },

        { CKA_SUBPRIME, &pqg.subPrime },

        { CKA_BASE, &pqg.base },

        { CKA_NSS_JPAKE_SIGNERID, &signerID }

    };

    SECItem x2, gx1, gx2;

    const SFTKItemTemplate generatedAttrs[] = {

        { CKA_NSS_JPAKE_X2,  &x2  },

        { CKA_NSS_JPAKE_GX1, &gx1 },

        { CKA_NSS_JPAKE_GX2, &gx2 },

    };

    SECItem x1;

    PORT_Assert(params != NULL);

    PORT_Assert(key != NULL);

    arena = PORT_NewArena(NSS_SOFTOKEN_DEFAULT_CHUNKSIZE);

    if (arena == NULL)

        crv = CKR_HOST_MEMORY;

    crv = sftk_MultipleAttribute2SecItem(arena, key, templateAttrs,

                                         NUM_ELEM(templateAttrs));

    if (crv == CKR_OK && (signerID.data == NULL || signerID.len == 0))

        crv = CKR_TEMPLATE_INCOMPLETE;

    /* generate x1, g^x1 and the proof of knowledge of x1 */

    if (crv == CKR_OK) {

        x1.data = NULL;

        crv = jpake_mapStatus(DSA_NewRandom(arena, &pqg.subPrime, &x1),

                              CKR_TEMPLATE_INCONSISTENT);

    }

    if (crv == CKR_OK)

        crv = jpake_Sign(arena, &pqg, hashType, &signerID, &x1, &params->gx1);

    /* generate x2, g^x2 and the proof of knowledge of x2 */

    if (crv == CKR_OK) {

        x2.data = NULL;

        crv = jpake_mapStatus(DSA_NewRandom(arena, &pqg.subPrime, &x2),

                              CKR_TEMPLATE_INCONSISTENT);

    }

    if (crv == CKR_OK)

        crv = jpake_Sign(arena, &pqg, hashType, &signerID, &x2, &params->gx2);

    /* Save the values needed for round 2 into CKA_VALUE */

    if (crv == CKR_OK) {

        gx1.data = params->gx1.pGX;

        gx1.len = params->gx1.ulGXLen;

        gx2.data = params->gx2.pGX;

        gx2.len = params->gx2.ulGXLen;

        crv = jpake_MultipleSecItem2Attribute(key, generatedAttrs, 

                                              NUM_ELEM(generatedAttrs));

    }

    PORT_FreeArena(arena, PR_TRUE);

    return crv;

}

CK_RV

jpake_Round2(HASH_HashType hashType, CK_NSS_JPAKERound2Params * params,

             SFTKObject * sourceKey, SFTKObject * key)

{

    CK_RV crv;

    PLArenaPool * arena;

    PQGParams pqg;

    SECItem signerID, x2, gx1, gx2;

    SFTKItemTemplate sourceAttrs[] = { 

        { CKA_PRIME, &pqg.prime },

        { CKA_SUBPRIME, &pqg.subPrime },

        { CKA_BASE, &pqg.base },

        { CKA_NSS_JPAKE_SIGNERID, &signerID },

        { CKA_NSS_JPAKE_X2,  &x2 },

        { CKA_NSS_JPAKE_GX1, &gx1 },

        { CKA_NSS_JPAKE_GX2, &gx2 },

    };

    SECItem x2s, gx3, gx4;

    const SFTKItemTemplate copiedAndGeneratedAttrs[] = {

        { CKA_NSS_JPAKE_SIGNERID, &signerID },

        { CKA_PRIME, &pqg.prime },

        { CKA_SUBPRIME, &pqg.subPrime },

        { CKA_NSS_JPAKE_X2,  &x2  },

        { CKA_NSS_JPAKE_X2S, &x2s },

        { CKA_NSS_JPAKE_GX1, &gx1 },

        { CKA_NSS_JPAKE_GX2, &gx2 },

        { CKA_NSS_JPAKE_GX3, &gx3 },

        { CKA_NSS_JPAKE_GX4, &gx4 }

    };

    SECItem peerID;

    PORT_Assert(params != NULL);

    PORT_Assert(sourceKey != NULL);

    PORT_Assert(key != NULL);

    arena = PORT_NewArena(NSS_SOFTOKEN_DEFAULT_CHUNKSIZE);

    if (arena == NULL)

        crv = CKR_HOST_MEMORY;

    /* TODO: check CKK_NSS_JPAKE_ROUND1 */

    crv = sftk_MultipleAttribute2SecItem(arena, sourceKey, sourceAttrs,

                                         NUM_ELEM(sourceAttrs));

    /* Get the peer's ID out of the template and sanity-check it. */

    if (crv == CKR_OK)

        crv = sftk_Attribute2SecItem(arena, &peerID, key,

                                     CKA_NSS_JPAKE_PEERID);

    if (crv == CKR_OK && (peerID.data == NULL || peerID.len == 0))

        crv = CKR_TEMPLATE_INCOMPLETE;

    if (crv == CKR_OK && SECITEM_CompareItem(&signerID, &peerID) == SECEqual)

        crv = CKR_TEMPLATE_INCONSISTENT;

    /* Verify zero-knowledge proofs for g^x3 and g^x4 */

    if (crv == CKR_OK)

        crv = jpake_Verify(arena, &pqg, hashType, &signerID,

                           peerID.data, peerID.len, &params->gx3);

    if (crv == CKR_OK)

        crv = jpake_Verify(arena, &pqg, hashType, &signerID,

                           peerID.data, peerID.len, &params->gx4);

    /* Calculate the base and x2s for A=base^x2s */

    if (crv == CKR_OK) {

        SECItem s;

        s.data = params->pSharedKey;

        s.len = params->ulSharedKeyLen;

        gx3.data = params->gx3.pGX;

        gx3.len = params->gx3.ulGXLen;

        gx4.data = params->gx4.pGX;

        gx4.len = params->gx4.ulGXLen;

        pqg.base.data = NULL;

        x2s.data = NULL;

        crv = jpake_mapStatus(JPAKE_Round2(arena, &pqg.prime, &pqg.subPrime,

                                           &gx1, &gx3, &gx4, &pqg.base, 

                                           &x2, &s, &x2s),

                              CKR_MECHANISM_PARAM_INVALID);

    }

    /* Generate A=base^x2s and its zero-knowledge proof. */

    if (crv == CKR_OK)

        crv = jpake_Sign(arena, &pqg, hashType, &signerID, &x2s, &params->A);

    /* Copy P and Q from the ROUND1 key to the ROUND2 key and save the values

       needed for the final key material derivation into CKA_VALUE. */

    if (crv == CKR_OK)

        crv = sftk_forceAttribute(key, CKA_PRIME, pqg.prime.data,

                                  pqg.prime.len);

    if (crv == CKR_OK)

        crv = sftk_forceAttribute(key, CKA_SUBPRIME, pqg.subPrime.data,

                                  pqg.subPrime.len);

    if (crv == CKR_OK) {

        crv = jpake_MultipleSecItem2Attribute(key, copiedAndGeneratedAttrs,

                                              NUM_ELEM(copiedAndGeneratedAttrs));

    }

    if (crv == CKR_OK)

        crv = jpake_enforceKeyType(key, CKK_NSS_JPAKE_ROUND2);

    PORT_FreeArena(arena, PR_TRUE);

    return crv;

}

CK_RV

jpake_Final(HASH_HashType hashType, const CK_NSS_JPAKEFinalParams * param,

            SFTKObject * sourceKey, SFTKObject * key)

{

    PLArenaPool * arena;

    SECItem K;

    PQGParams pqg;

    CK_RV crv;

    SECItem peerID, signerID, x2s, x2, gx1, gx2, gx3, gx4;

    SFTKItemTemplate sourceAttrs[] = {

        { CKA_NSS_JPAKE_PEERID, &peerID },

        { CKA_NSS_JPAKE_SIGNERID, &signerID },

        { CKA_PRIME, &pqg.prime },

        { CKA_SUBPRIME, &pqg.subPrime },

        { CKA_NSS_JPAKE_X2,  &x2  },

        { CKA_NSS_JPAKE_X2S, &x2s },

        { CKA_NSS_JPAKE_GX1, &gx1 },

        { CKA_NSS_JPAKE_GX2, &gx2 },

        { CKA_NSS_JPAKE_GX3, &gx3 },

        { CKA_NSS_JPAKE_GX4, &gx4 }

    };

    PORT_Assert(param != NULL);

    PORT_Assert(sourceKey != NULL);

    PORT_Assert(key != NULL);

    arena = PORT_NewArena(NSS_SOFTOKEN_DEFAULT_CHUNKSIZE);

    if (arena == NULL)

        crv = CKR_HOST_MEMORY;

    /* TODO: verify key type CKK_NSS_JPAKE_ROUND2 */

    crv = sftk_MultipleAttribute2SecItem(arena, sourceKey, sourceAttrs,

                                         NUM_ELEM(sourceAttrs));

    /* Calculate base for B=base^x4s */

    if (crv == CKR_OK) {

        pqg.base.data = NULL;

        crv = jpake_mapStatus(JPAKE_Round2(arena, &pqg.prime, &pqg.subPrime,

                                           &gx1, &gx2, &gx3, &pqg.base,

                                           NULL, NULL, NULL),

                              CKR_MECHANISM_PARAM_INVALID);

    }

    /* Verify zero-knowledge proof for B */

    if (crv == CKR_OK)

        crv = jpake_Verify(arena, &pqg, hashType, &signerID,

                           peerID.data, peerID.len, &param->B);

    if (crv == CKR_OK) {

        SECItem B;

        B.data = param->B.pGX;

        B.len = param->B.ulGXLen;

        K.data = NULL;

        crv = jpake_mapStatus(JPAKE_Final(arena, &pqg.prime, &pqg.subPrime,

                                          &x2, &gx4, &x2s, &B, &K),

                              CKR_MECHANISM_PARAM_INVALID);

    }

    /* Save key material into CKA_VALUE. */

    if (crv == CKR_OK)

        crv = sftk_forceAttribute(key, CKA_VALUE, K.data, K.len);

    if (crv == CKR_OK)

        crv = jpake_enforceKeyType(key, CKK_GENERIC_SECRET);

    PORT_FreeArena(arena, PR_TRUE);

    return crv;

}