/*
 * This file and its contents are supplied under the terms of the
 * Common Development and Distribution License ("CDDL"), version 1.0.
 * You may only use this file in accordance with the terms of version
 * 1.0 of the CDDL.
 *
 * A full copy of the text of the CDDL should have accompanied this
 * source.  A copy of the CDDL is also available via the Internet at
 * http://www.illumos.org/license/CDDL.
 */

/*
 * Copyright 2018-2021 Tintri by DDN, Inc. All rights reserved.
 * Copyright 2021 RackTop Systems, Inc.
 */

/*
 * Helper functions for SMB3 encryption using PKCS#11
 *
 * There are two implementations of these functions:
 * This one (for user space) and another for kernel.
 * See: uts/common/fs/smbsrv/smb3_encrypt_kcf.c
 *
 * Contrary to what one might assume from the file name,
 * there should be NO SMB implementation knowledge here
 * beyond a few carefully selected things (smb_kcrypt.h).
 */

#include <security/cryptoki.h>
#include <security/pkcs11.h>
#include <smbsrv/smb_kcrypt.h>

#include <sys/cmn_err.h>
#include <sys/debug.h>
#include <stdlib.h>
#include <strings.h>

/*
 * Common function to see if a mech is available.
 */
static int
find_mech(CK_MECHANISM_TYPE id)
{
        CK_SESSION_HANDLE hdl;
        CK_RV rv;

        rv = SUNW_C_GetMechSession(id, &hdl);
        if (rv != CKR_OK) {
                return (-1);
        }
        (void) C_CloseSession(hdl);

        return (0);
}

/*
 * SMB3 encryption helpers:
 * (getmech, init, update, final)
 */

int
smb3_aes_ccm_getmech(smb_crypto_mech_t *mech)
{

        if (find_mech(CKM_AES_CCM) != 0) {
                cmn_err(CE_NOTE, "PKCS#11: no mech AES_CCM");
                return (-1);
        }

        mech->mechanism = CKM_AES_CCM;
        return (0);
}

int
smb3_aes_gcm_getmech(smb_crypto_mech_t *mech)
{

        if (find_mech(CKM_AES_GCM) != 0) {
                cmn_err(CE_NOTE, "PKCS#11: no mech CKM_AES_GCM");
                return (-1);
        }

        mech->mechanism = CKM_AES_GCM;
        return (0);
}

void
smb3_crypto_init_ccm_param(smb_enc_ctx_t *ctx,
    uint8_t *nonce, size_t noncesize,
    uint8_t *auth, size_t authsize,
    size_t datasize)
{

        ASSERT3U(noncesize, >=, SMB3_AES_CCM_NONCE_SIZE);

        /* CK_CCM_PARAMS */
        ctx->param.ccm.ulDataLen = datasize;
        ctx->param.ccm.pNonce = nonce;
        ctx->param.ccm.ulNonceLen = SMB3_AES_CCM_NONCE_SIZE;
        ctx->param.ccm.pAAD = auth;
        ctx->param.ccm.ulAADLen = authsize;
        ctx->param.ccm.ulMACLen = SMB2_SIG_SIZE;

        ctx->mech.pParameter = (caddr_t)&ctx->param.ccm;
        ctx->mech.ulParameterLen = sizeof (ctx->param.ccm);
}

void
smb3_crypto_init_gcm_param(smb_enc_ctx_t *ctx,
    uint8_t *nonce, size_t noncesize,
    uint8_t *auth, size_t authsize)
{

        ASSERT3U(noncesize, >=, SMB3_AES_GCM_NONCE_SIZE);

        /* CK_GCM_PARAMS */
        ctx->param.gcm.pIv = nonce;
        ctx->param.gcm.ulIvLen = SMB3_AES_GCM_NONCE_SIZE;
        ctx->param.gcm.pAAD = auth;                     /* auth data */
        ctx->param.gcm.ulAADLen = authsize;             /* auth data len */
        ctx->param.gcm.ulTagBits = SMB2_SIG_SIZE << 3;  /* bytes to bits */

        ctx->mech.pParameter = (caddr_t)&ctx->param.gcm;
        ctx->mech.ulParameterLen = sizeof (ctx->param.gcm);
}

/*
 * Start the KCF encrypt session, load the key
 * If this returns zero, the caller should call
 * smb3_enc_ctx_done to cleanup the context,
 * even if there are intervening errors.
 */
int
smb3_encrypt_init(smb_enc_ctx_t *ctxp,
    uint8_t *key, size_t keylen)
{
        CK_OBJECT_HANDLE hkey = 0;
        CK_MECHANISM *mech = &ctxp->mech;
        CK_RV rv;

        rv = SUNW_C_GetMechSession(mech->mechanism, &ctxp->ctx);
        if (rv != CKR_OK)
                return (-1);

        rv = SUNW_C_KeyToObject(ctxp->ctx, mech->mechanism,
            key, keylen, &hkey);
        if (rv != CKR_OK)
                return (-1);

        rv = C_EncryptInit(ctxp->ctx, mech, hkey);
        if (rv != CKR_OK) {
                cmn_err(CE_WARN, "C_EncryptInit failed: 0x%lx", rv);
        }
        (void) C_DestroyObject(ctxp->ctx, hkey);

        return (rv == CKR_OK ? 0 : -1);
}

/*
 * Start the KCF decrypt session, load the key
 * If this returns zero, the caller should call
 * smb3_enc_ctx_done to cleanup the context,
 * even if there are intervening errors.
 */
int
smb3_decrypt_init(smb_enc_ctx_t *ctxp,
    uint8_t *key, size_t keylen)
{
        CK_OBJECT_HANDLE hkey = 0;
        CK_MECHANISM *mech = &ctxp->mech;
        CK_RV rv;

        rv = SUNW_C_GetMechSession(mech->mechanism, &ctxp->ctx);
        if (rv != CKR_OK)
                return (-1);

        rv = SUNW_C_KeyToObject(ctxp->ctx, mech->mechanism,
            key, keylen, &hkey);
        if (rv != CKR_OK)
                return (-1);

        rv = C_DecryptInit(ctxp->ctx, mech, hkey);
        if (rv != CKR_OK) {
                cmn_err(CE_WARN, "C_DecryptInit failed: 0x%lx", rv);
        }
        (void) C_DestroyObject(ctxp->ctx, hkey);

        return (rv == CKR_OK ? 0 : -1);
}

/*
 * Encrypt a whole message with scatter/gather (UIO)
 *
 * While the PKCS#11 implementation internally has the ability to
 * handle scatter/gather, it currently presents no interface for it.
 * As this library is used primarily for debugging, performance in
 * here is not a big concern, so we'll get around the limitation of
 * libpkcs11 by copying to/from a contiguous working buffer.
 */
int
smb3_encrypt_uio(smb_enc_ctx_t *ctxp, uio_t *in, uio_t *out)
{
        uint8_t *buf = NULL;
        size_t inlen, outlen;
        ulong_t tlen;
        int err, rc = -1;
        CK_RV rv;

        if (in->uio_resid <= 0)
                return (-1);
        inlen = in->uio_resid;
        outlen = inlen + 16;
        buf = malloc(outlen);
        if (buf == NULL)
                return (-1);

        /* Copy from uio segs to buf */
        err = uiomove(buf, inlen, UIO_WRITE, in);
        if (err != 0)
                goto out;

        /* Encrypt in-place in our work buffer. */
        tlen = outlen;
        rv = C_Encrypt(ctxp->ctx, buf, inlen, buf, &tlen);
        if (rv != CKR_OK) {
                cmn_err(CE_WARN, "C_Encrypt failed: 0x%lx", rv);
                goto out;
        }
        if (tlen != outlen) {
                cmn_err(CE_WARN, "smb3_encrypt_uio outlen %d vs %d",
                    (int)tlen, (int)outlen);
                goto out;
        }

        /* Copy from buf to uio segs */
        err = uiomove(buf, outlen, UIO_READ, out);
        if (err != 0)
                goto out;

        rc = 0;
out:
        free(buf);

        return (rc);
}

int
smb3_decrypt_uio(smb_enc_ctx_t *ctxp, uio_t *in, uio_t *out)
{
        uint8_t *buf = NULL;
        size_t inlen, outlen;
        ulong_t tlen;
        int err, rc = -1;
        CK_RV rv;

        if (in->uio_resid <= 16)
                return (-1);
        inlen = in->uio_resid;
        outlen = inlen - 16;
        buf = malloc(inlen);
        if (buf == NULL)
                return (-1);

        /* Copy from uio segs to buf */
        err = uiomove(buf, inlen, UIO_WRITE, in);
        if (err != 0)
                goto out;

        /* Decrypt in-place in our work buffer. */
        tlen = outlen;
        rv = C_Decrypt(ctxp->ctx, buf, inlen, buf, &tlen);
        if (rv != CKR_OK) {
                cmn_err(CE_WARN, "C_Decrypt failed: 0x%lx", rv);
                goto out;
        }
        if (tlen != outlen) {
                cmn_err(CE_WARN, "smb3_decrypt_uio outlen %d vs %d",
                    (int)tlen, (int)outlen);
                goto out;
        }

        /* Copy from buf to uio segs */
        err = uiomove(buf, outlen, UIO_READ, out);
        if (err != 0)
                goto out;

        rc = 0;
out:
        free(buf);

        return (rc);
}

void
smb3_enc_ctx_done(smb_enc_ctx_t *ctxp)
{
        if (ctxp->ctx != 0) {
                (void) C_CloseSession(ctxp->ctx);
                ctxp->ctx = 0;
        }
}