/* 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 "secoid.h"

#include "secasn1.h"

#include "pkcs11.h"

#include "pk11func.h"

#include "pk11sdr.h"

/*

 * Data structure and template for encoding the result of an SDR operation

 *  This is temporary.  It should include the algorithm ID of the encryption mechanism

 */

struct SDRResult {

    SECItem keyid;

    SECAlgorithmID alg;

    SECItem data;

};

typedef struct SDRResult SDRResult;

SEC_ASN1_MKSUB(SECOID_AlgorithmIDTemplate)

static SEC_ASN1Template template[] = {

    { SEC_ASN1_SEQUENCE, 0, NULL, sizeof(SDRResult) },

    { SEC_ASN1_OCTET_STRING, offsetof(SDRResult, keyid) },

    { SEC_ASN1_INLINE | SEC_ASN1_XTRN, offsetof(SDRResult, alg),

      SEC_ASN1_SUB(SECOID_AlgorithmIDTemplate) },

    { SEC_ASN1_OCTET_STRING, offsetof(SDRResult, data) },

    { 0 }

};

static unsigned char keyID[] = {

    0xF8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,

    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01

};

static SECItem keyIDItem = {

    0,

    keyID,

    sizeof keyID

};

/* local utility function for padding an incoming data block

 * to the mechanism block size.

 */

static SECStatus

padBlock(SECItem *data, int blockSize, SECItem *result)

{

    SECStatus rv = SECSuccess;

    int padLength;

    unsigned int i;

    result->data = 0;

    result->len = 0;

    /* This algorithm always adds to the block (to indicate the number

     * of pad bytes).  So allocate a block large enough.

     */

    padLength = blockSize - (data->len % blockSize);

    result->len = data->len + padLength;

    result->data = (unsigned char *)PORT_Alloc(result->len);

    /* Copy the data */

    PORT_Memcpy(result->data, data->data, data->len);

    /* Add the pad values */

    for (i = data->len; i < result->len; i++)

        result->data[i] = (unsigned char)padLength;

    return rv;

}

static SECStatus

unpadBlock(SECItem *data, int blockSize, SECItem *result)

{

    SECStatus rv = SECSuccess;

    int padLength;

    unsigned int i;

    result->data = 0;

    result->len = 0;

    /* Remove the padding from the end if the input data */

    if (data->len == 0 || data->len % blockSize != 0) {

        rv = SECFailure;

        goto loser;

    }

    padLength = data->data[data->len - 1];

    if (padLength > blockSize) {

        rv = SECFailure;

        goto loser;

    }

    /* verify padding */

    for (i = data->len - padLength; i < data->len; i++) {

        if (data->data[i] != padLength) {

            rv = SECFailure;

            goto loser;

        }

    }

    result->len = data->len - padLength;

    result->data = (unsigned char *)PORT_Alloc(result->len);

    if (!result->data) {

        rv = SECFailure;

        goto loser;

    }

    PORT_Memcpy(result->data, data->data, result->len);

    if (padLength < 2) {

        return SECWouldBlock;

    }

loser:

    return rv;

}

static PRLock *pk11sdrLock = NULL;

void

pk11sdr_Init(void)

{

    pk11sdrLock = PR_NewLock();

}

void

pk11sdr_Shutdown(void)

{

    if (pk11sdrLock) {

        PR_DestroyLock(pk11sdrLock);

        pk11sdrLock = NULL;

    }

}

/*

 * PK11SDR_Encrypt

 *  Deprecated version of PK11SDR_EncryptWithMechanism using DES3_CBC.

 */

SECStatus

PK11SDR_Encrypt(SECItem *keyid, SECItem *data, SECItem *result, void *cx)

{

    return PK11SDR_EncryptWithMechanism(NULL, keyid, CKM_DES3_CBC, data, result, cx);

}

/*

 * PK11SDR_EncryptWithMechanism

 *  Encrypt a block of data using the symmetric key identified and the

 *  encryption mechanism specified (only AES_CBC and DES3_CBC are supported).

 *  The result is an ASN.1 (DER) encoded block of keyid, params and data.

 */

SECStatus

PK11SDR_EncryptWithMechanism(PK11SlotInfo *slot, SECItem *keyid, CK_MECHANISM_TYPE type, SECItem *data, SECItem *result, void *cx)

{

    SECStatus rv = SECSuccess;

    PK11SymKey *key = 0;

    SECItem *params = 0;

    PK11Context *ctx = 0;

    SDRResult sdrResult;

    SECItem paddedData;

    SECItem *pKeyID;

    PLArenaPool *arena = 0;

    SECOidTag algtag;

    PK11SlotInfo *aSlot = slot;

    /* Initialize */

    paddedData.len = 0;

    paddedData.data = 0;

    arena = PORT_NewArena(SEC_ASN1_DEFAULT_ARENA_SIZE);

    if (!arena) {

        rv = SECFailure;

        goto loser;

    }

    /* 1. Locate the requested keyid, or the default key (which has a keyid)

     * 2. Create an encryption context

     * 3. Encrypt

     * 4. Encode the results (using ASN.1)

     */

    if (!slot) {

        slot = PK11_GetInternalKeySlot();

        if (!slot) {

            rv = SECFailure;

            goto loser;

        }

    }

    /*

     * Login to the internal token before we look for the key, otherwise we

     * won't find it.

     */

    rv = PK11_Authenticate(slot, PR_TRUE, cx);

    if (rv != SECSuccess)

        goto loser;

    /* Find the key to use */

    pKeyID = keyid;

    if (pKeyID->len == 0) {

        int keySize = PK11_GetBestKeyLength(slot, type);

        pKeyID = &keyIDItem;

        /* put in a course lock to prevent a race between not finding the

         * key and creating  one.

         */

        if (pk11sdrLock)

            PR_Lock(pk11sdrLock);

        /* Try to find the key */

        key = PK11_FindFixedKey(slot, type, pKeyID, cx);

        /* If the default key doesn't exist yet, try to create it */

        if (!key)

            key = PK11_TokenKeyGen(slot, type, 0, keySize, pKeyID, PR_TRUE, cx);

        if (pk11sdrLock)

            PR_Unlock(pk11sdrLock);

    } else {

        key = PK11_FindFixedKey(slot, type, pKeyID, cx);

    }

    if (!key) {

        rv = SECFailure;

        goto loser;

    }

    params = PK11_GenerateNewParam(type, key);

    if (!params) {

        rv = SECFailure;

        goto loser;

    }

    ctx = PK11_CreateContextBySymKey(type, CKA_ENCRYPT, key, params);

    if (!ctx) {

        rv = SECFailure;

        goto loser;

    }

    rv = padBlock(data, PK11_GetBlockSize(type, 0), &paddedData);

    if (rv != SECSuccess)

        goto loser;

    sdrResult.data.len = paddedData.len;

    sdrResult.data.data = (unsigned char *)PORT_ArenaAlloc(arena, sdrResult.data.len);

    rv = PK11_CipherOp(ctx, sdrResult.data.data, (int *)&sdrResult.data.len, sdrResult.data.len,

                       paddedData.data, paddedData.len);

    if (rv != SECSuccess)

        goto loser;

    PK11_Finalize(ctx);

    sdrResult.keyid = *pKeyID;

    algtag = SECOID_FindOIDByMechanism(type)->offset;

    rv = PK11_ParamToAlgid(algtag, params, arena, &sdrResult.alg);

    if (rv != SECSuccess)

        goto loser;

    if (!SEC_ASN1EncodeItem(0, result, &sdrResult, template)) {

        rv = SECFailure;

        goto loser;

    }

loser:

    SECITEM_ZfreeItem(&paddedData, PR_FALSE);

    if (arena)

        PORT_FreeArena(arena, PR_TRUE);

    if (ctx)

        PK11_DestroyContext(ctx, PR_TRUE);

    if (params)

        SECITEM_ZfreeItem(params, PR_TRUE);

    if (key)

        PK11_FreeSymKey(key);

    if (slot && !aSlot)

        PK11_FreeSlot(slot);

    return rv;

}

/* decrypt a block */

static SECStatus

pk11Decrypt(PK11SlotInfo *slot, PLArenaPool *arena,

            CK_MECHANISM_TYPE type, PK11SymKey *key,

            SECItem *params, SECItem *in, SECItem *result)

{

    PK11Context *ctx = 0;

    SECItem paddedResult;

    SECStatus rv;

    paddedResult.len = 0;

    paddedResult.data = 0;

    ctx = PK11_CreateContextBySymKey(type, CKA_DECRYPT, key, params);

    if (!ctx) {

        rv = SECFailure;

        goto loser;

    }

    paddedResult.len = in->len;

    paddedResult.data = PORT_ArenaAlloc(arena, paddedResult.len);

    rv = PK11_CipherOp(ctx, paddedResult.data,

                       (int *)&paddedResult.len, paddedResult.len,

                       in->data, in->len);

    if (rv != SECSuccess)

        goto loser;

    PK11_Finalize(ctx);

    /* Remove the padding */

    rv = unpadBlock(&paddedResult, PK11_GetBlockSize(type, 0), result);

    if (rv)

        goto loser;

loser:

    if (ctx)

        PK11_DestroyContext(ctx, PR_TRUE);

    return rv;

}

/*

 * PK11SDR_Decrypt

 *  Decrypt a block of data produced by PK11SDR_EncryptWithMechanism. The key

 *  used is identified by the keyid field within the input.

 */

SECStatus

PK11SDR_Decrypt(SECItem *data, SECItem *result, void *cx)

{

    SECStatus rv = SECSuccess;

    PK11SlotInfo *slot = 0;

    PK11SymKey *key = 0;

    CK_MECHANISM_TYPE type;

    SDRResult sdrResult;

    SECItem *params = 0;

    SECItem possibleResult = { 0, NULL, 0 };

    PLArenaPool *arena = 0;

    SECOidTag algtag;

    arena = PORT_NewArena(SEC_ASN1_DEFAULT_ARENA_SIZE);

    if (!arena) {

        rv = SECFailure;

        goto loser;

    }

    /* Decode the incoming data */

    memset(&sdrResult, 0, sizeof sdrResult);

    rv = SEC_QuickDERDecodeItem(arena, &sdrResult, template, data);

    if (rv != SECSuccess)

        goto loser; /* Invalid format */

    /* Find the slot and key for the given keyid */

    slot = PK11_GetInternalKeySlot();

    if (!slot) {

        rv = SECFailure;

        goto loser;

    }

    rv = PK11_Authenticate(slot, PR_TRUE, cx);

    if (rv != SECSuccess)

        goto loser;

    /* Get the parameter values from the data */

    params = PK11_ParamFromAlgid(&sdrResult.alg);

    if (!params) {

        rv = SECFailure;

        goto loser;

    }

    algtag = SECOID_GetAlgorithmTag(&sdrResult.alg);

    type = PK11_AlgtagToMechanism(algtag);

    key = PK11_FindFixedKey(slot, type, &sdrResult.keyid, cx);

    if (!key) {

        rv = SECFailure;

    } else {

        rv = pk11Decrypt(slot, arena, type, key, params,

                         &sdrResult.data, result);

    }

    /*

     * if the pad value was too small (1 or 2), then it's statistically

     * 'likely' that (1 in 256) that we may not have the correct key.

     * Check the other keys for a better match. If we find none, use

     * this result.

     */

    if (rv == SECWouldBlock) {

        possibleResult = *result;

    }

    /*

     * handle the case where your key indicies may have been broken

     */

    if (rv != SECSuccess) {

        PK11SymKey *keyList = PK11_ListFixedKeysInSlot(slot, NULL, cx);

        PK11SymKey *testKey = NULL;

        PK11SymKey *nextKey = NULL;

        for (testKey = keyList; testKey;

             testKey = PK11_GetNextSymKey(testKey)) {

            if (PK11_GetSymKeyType(testKey) != PK11_GetKeyType(type, 0)) {

                continue;

            }

            rv = pk11Decrypt(slot, arena, type, testKey, params,

                             &sdrResult.data, result);

            if (rv == SECSuccess) {

                break;

            }

            /* found a close match. If it's our first remember it */

            if (rv == SECWouldBlock) {

                if (possibleResult.data) {

                    /* this is unlikely but possible. If we hit this condition,

                     * we have no way of knowing which possibility to prefer.

                     * in this case we just match the key the application

                     * thought was the right one */

                    SECITEM_ZfreeItem(result, PR_FALSE);

                } else {

                    possibleResult = *result;

                }

            }

        }

        /* free the list */

        for (testKey = keyList; testKey; testKey = nextKey) {

            nextKey = PK11_GetNextSymKey(testKey);

            PK11_FreeSymKey(testKey);

        }

    }

    /* we didn't find a better key, use the one with a small pad value */

    if ((rv != SECSuccess) && (possibleResult.data)) {

        *result = possibleResult;

        possibleResult.data = NULL;

        rv = SECSuccess;

    }

loser:

    if (arena)

        PORT_FreeArena(arena, PR_TRUE);

    if (key)

        PK11_FreeSymKey(key);

    if (params)

        SECITEM_ZfreeItem(params, PR_TRUE);

    if (slot)

        PK11_FreeSlot(slot);

    if (possibleResult.data)

        SECITEM_ZfreeItem(&possibleResult, PR_FALSE);

    return rv;

}