/* 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 "pk11func.h"
#include "pk11sdr.h"
#include "pkcs11.h"
#include "secasn1.h"
#include "seccomon.h"
#include "secerr.h"
#include "secoid.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;

    if (type != CKM_DES3_CBC && type != CKM_AES_CBC) {
        PORT_SetError(SEC_ERROR_INVALID_ARGS);
        rv = SECFailure;
        goto loser;
    }

    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 */
    }

    algtag = SECOID_GetAlgorithmTag(&sdrResult.alg);
    type = PK11_AlgtagToMechanism(algtag);
    if (type != CKM_DES3_CBC && type != CKM_AES_CBC) {
        PORT_SetError(SEC_ERROR_INVALID_ARGS);
        rv = SECFailure;
        goto loser;
    }

    /* 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;
    }

    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;
}