/*
 * 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 2019 Nexenta Systems, Inc.  All rights reserved.
 * Copyright 2024 RackTop Systems, Inc.
 */

/*
 * Support for SMB3 encryption (message privacy)
 */

#include <sys/param.h>
#include <sys/systm.h>
#include <sys/conf.h>
#include <sys/proc.h>
#include <sys/fcntl.h>
#include <sys/socket.h>
#include <sys/kmem.h>
#include <sys/errno.h>
#include <sys/cmn_err.h>
#include <sys/random.h>
#include <sys/stream.h>
#include <sys/strsun.h>
#include <sys/sdt.h>

#include <netsmb/smb_osdep.h>
#include <netsmb/smb2.h>
#include <netsmb/smb_conn.h>
#include <netsmb/smb_subr.h>
#include <netsmb/smb_dev.h>
#include <netsmb/smb_rq.h>

#include <netsmb/nsmb_kcrypt.h>

#define	SMB3_TFORM_HDR_SIZE	52
#define	SMB3_NONCE_OFFS		20
#define	SMB3_SIG_OFFS		4

static const uint8_t SMB3_CRYPT_SIG[4] = { 0xFD, 'S', 'M', 'B' };

/*
 * Initialize crypto mechanisms we'll need.
 * Called after negotiate.
 */
void
nsmb_crypt_init_mech(struct smb_vc *vcp)
{
	smb_crypto_mech_t *mech;
	int rc;

	if (vcp->vc3_crypt_mech != NULL)
		return;

	mech = kmem_zalloc(sizeof (*mech), KM_SLEEP);

	/* Always CCM for now. */
	rc = nsmb_aes_ccm_getmech(mech);
	if (rc != 0) {
		kmem_free(mech, sizeof (*mech));
		cmn_err(CE_NOTE, "SMB3 found no AES mechanism"
		    " (encryption disabled)");
		return;
	}
	vcp->vc3_crypt_mech = mech;
}

void
nsmb_crypt_free_mech(struct smb_vc *vcp)
{
	smb_crypto_mech_t *mech;

	if ((mech = vcp->vc3_crypt_mech) == NULL)
		return;

	kmem_free(mech, sizeof (*mech));
}

/*
 * Initialize keys for encryption
 * Called after session setup.
 */
void
nsmb_crypt_init_keys(struct smb_vc *vcp)
{

	/*
	 * If we don't have a session key, we'll fail later when a
	 * request that requires (en/de)cryption can't be (en/de)crypted.
	 * Also don't bother initializing if we don't have a mechanism.
	 */
	if (vcp->vc3_crypt_mech == NULL ||
	    vcp->vc_ssnkeylen <= 0)
		return;

	/*
	 * For SMB3, the encrypt/decrypt keys are derived from
	 * the session key using KDF in counter mode.
	 */
	if (nsmb_kdf(vcp->vc3_encrypt_key, SMB3_KEYLEN,
	    vcp->vc_ssnkey, vcp->vc_ssnkeylen,
	    (uint8_t *)"SMB2AESCCM", 11,
	    (uint8_t *)"ServerIn ", 10) != 0)
		return;

	if (nsmb_kdf(vcp->vc3_decrypt_key, SMB3_KEYLEN,
	    vcp->vc_ssnkey, vcp->vc_ssnkeylen,
	    (uint8_t *)"SMB2AESCCM", 11,
	    (uint8_t *)"ServerOut", 10) != 0)
		return;

	vcp->vc3_encrypt_key_len = SMB3_KEYLEN;
	vcp->vc3_decrypt_key_len = SMB3_KEYLEN;

	(void) random_get_pseudo_bytes(
	    (uint8_t *)&vcp->vc3_nonce_low,
	    sizeof (vcp->vc3_nonce_low));
	(void) random_get_pseudo_bytes(
	    (uint8_t *)&vcp->vc3_nonce_high,
	    sizeof (vcp->vc3_nonce_high));
}

/*
 * Encrypt the message in *mpp, in place, prepending the
 * SMB3 transform header.
 *
 * Any non-zero return is an error (values not used).
 */
int
smb3_msg_encrypt(struct smb_vc *vcp, mblk_t **mpp)
{
	smb_enc_ctx_t ctx;
	mblk_t *body, *thdr, *lastm;
	struct mbchain	mbp_store;
	struct mbchain *mbp = &mbp_store;
	uint32_t bodylen;
	uint8_t *authdata;
	size_t authlen;
	int rc;

	ASSERT(RW_WRITE_HELD(&vcp->iod_rqlock));

	if (vcp->vc3_crypt_mech == NULL ||
	    vcp->vc3_encrypt_key_len != SMB3_KEYLEN) {
		return (ENOTSUP);
	}

	bzero(&ctx, sizeof (ctx));
	ctx.mech = *((smb_crypto_mech_t *)vcp->vc3_crypt_mech);

	body = *mpp;
	bodylen = msgdsize(body);

	/*
	 * Get a new "nonce".  Access to these counters is
	 * serialized by iod_rqlock (assert above).
	 */
	vcp->vc3_nonce_low++;
	if (vcp->vc3_nonce_low == 0) {
		vcp->vc3_nonce_low++;
		vcp->vc3_nonce_high++;
	}

	/*
	 * Build the transform header, keeping pointers to the various
	 * parts of it that we'll need to refer to later.
	 */
	(void) mb_init(mbp);
	thdr = mbp->mb_top;
	ASSERT(MBLKTAIL(thdr) >= SMB3_TFORM_HDR_SIZE);
	mb_put_mem(mbp, SMB3_CRYPT_SIG, 4, MB_MSYSTEM);
	mb_put_mem(mbp, NULL, SMB2_SIG_SIZE, MB_MZERO);	// signature (later)
	mb_put_uint64le(mbp, vcp->vc3_nonce_low);
	mb_put_uint64le(mbp, vcp->vc3_nonce_high);
	/* Zero last 5 bytes of nonce per. spec. */
	bzero(thdr->b_wptr - 5, 5);
	mb_put_uint32le(mbp, bodylen);
	mb_put_uint16le(mbp, 0);	// reserved
	mb_put_uint16le(mbp, 1);	// flags
	mb_put_uint64le(mbp, vcp->vc2_session_id);
	mbp->mb_top = NULL; // keeping thdr
	mb_done(mbp);

	/*
	 * Need pointers to the part of the transfor header
	 * after the signature (starting with the nonce).
	 */
	authdata = thdr->b_rptr + SMB3_NONCE_OFFS;
	authlen = SMB3_TFORM_HDR_SIZE - SMB3_NONCE_OFFS;

	nsmb_crypto_init_ccm_param(&ctx,
	    authdata, SMB2_SIG_SIZE,
	    authdata, authlen, bodylen);

	rc = nsmb_encrypt_init(&ctx,
	    vcp->vc3_encrypt_key, vcp->vc3_encrypt_key_len);
	if (rc != 0)
		goto errout;

	/*
	 * Temporarily append the transform header onto the
	 * body mblk chain with its r/w pointers set to cover
	 * just the signature, needed for how encrypt works.
	 * Could just use linkb() but we need to unlink the
	 * block as well so just find the tail ourselves.
	 */
	ASSERT(MBLKL(thdr) == SMB3_TFORM_HDR_SIZE);
	thdr->b_rptr += SMB3_SIG_OFFS;
	thdr->b_wptr = thdr->b_rptr + SMB2_SIG_SIZE;
	lastm = body;
	while (lastm->b_cont != NULL)
		lastm = lastm->b_cont;
	lastm->b_cont = thdr;

	/*
	 * The mblk chain is ready. Encrypt!
	 */
	rc = nsmb_encrypt_mblks(&ctx, body, bodylen);
	/* check rc below */

	/* Unlink thdr and restore r/w pointers. */
	lastm->b_cont = NULL;
	thdr->b_rptr -= SMB3_SIG_OFFS;
	thdr->b_wptr = thdr->b_rptr + SMB3_TFORM_HDR_SIZE;

	/* Now check rc from encrypt */
	if (rc != 0)
		goto errout;

	/*
	 * Lastly, prepend the transform header.
	 */
	thdr->b_cont = body;
	*mpp = thdr;
	nsmb_enc_ctx_done(&ctx);
	return (0);

errout:
	freeb(thdr);
	nsmb_enc_ctx_done(&ctx);
	return (rc);
}

/*
 * Decrypt the message in *mpp, in place, removing the
 * SMB3 transform header.
 *
 * Any non-zero return is an error (values not used).
 */
int
smb3_msg_decrypt(struct smb_vc *vcp, mblk_t **mpp)
{
	smb_enc_ctx_t ctx;
	uint8_t th_sig[4];
	mblk_t *body, *thdr, *lastm;
	struct mdchain	mdp_store;
	struct mdchain *mdp = &mdp_store;
	uint64_t th_ssnid;
	uint32_t bodylen, tlen;
	uint16_t th_flags;
	uint8_t *authdata;
	size_t authlen;
	int rc;

	if (vcp->vc3_crypt_mech == NULL ||
	    vcp->vc3_encrypt_key_len != SMB3_KEYLEN) {
		return (ENOTSUP);
	}

	bzero(&ctx, sizeof (ctx));
	ctx.mech = *((smb_crypto_mech_t *)vcp->vc3_crypt_mech);

	/*
	 * Split off the transform header
	 * We need it contiguous.
	 */
	thdr = *mpp;
	body = m_split(thdr, SMB3_TFORM_HDR_SIZE, 1);
	if (body == NULL)
		return (ENOSR);
	thdr = m_pullup(thdr, SMB3_TFORM_HDR_SIZE);
	if (thdr == NULL)
		return (ENOSR);

	/*
	 * Decode the transform header
	 */
	(void) md_initm(mdp, thdr);
	md_get_mem(mdp, th_sig, 4, MB_MSYSTEM);
	md_get_mem(mdp, NULL, SMB2_SIG_SIZE, MB_MZERO); // signature
	md_get_mem(mdp, NULL, SMB2_SIG_SIZE, MB_MZERO); // nonce
	md_get_uint32le(mdp, &bodylen);
	md_get_uint16le(mdp, NULL);	// reserved
	md_get_uint16le(mdp, &th_flags);
	md_get_uint64le(mdp, &th_ssnid);
	mdp->md_top = NULL; // keeping thdr
	md_done(mdp);

	/*
	 * Validate transform header fields
	 */
	if (bcmp(th_sig, SMB3_CRYPT_SIG, 4) != 0) {
		rc = EPROTO;
		goto errout;
	}
	if (th_flags != 1 || th_ssnid != vcp->vc2_session_id) {
		rc = EINVAL;
		goto errout;
	}

	/*
	 * Check actual body length (trim if necessary)
	 */
	tlen = msgdsize(body);
	if (tlen < bodylen) {
		rc = EINVAL;
		goto errout;
	}
	if (tlen > bodylen) {
		/* trim from tail */
		ssize_t adj;

		adj = bodylen - tlen;
		ASSERT(adj < 0);
		(void) adjmsg(body, adj);
	}

	/*
	 * Need pointers to the part of the transfor header
	 * after the signature (starting with the nonce).
	 * tlen is now length of ciphertext
	 */
	authdata = thdr->b_rptr + SMB3_NONCE_OFFS;
	authlen = SMB3_TFORM_HDR_SIZE - SMB3_NONCE_OFFS;
	tlen = bodylen + SMB2_SIG_SIZE;

	nsmb_crypto_init_ccm_param(&ctx,
	    authdata, SMB2_SIG_SIZE,
	    authdata, authlen, tlen);

	rc = nsmb_decrypt_init(&ctx,
	    vcp->vc3_decrypt_key, vcp->vc3_decrypt_key_len);
	if (rc != 0)
		goto errout;

	/*
	 * Temporarily append the transform header onto the
	 * body mblk chain with its r/w pointers set to cover
	 * just the signature, needed for how decrypt works.
	 * Could just use linkb() but we need to unlink the
	 * block as well so just find the tail ourselves.
	 */
	thdr->b_rptr += SMB3_SIG_OFFS;
	thdr->b_wptr = thdr->b_rptr + SMB2_SIG_SIZE;
	lastm = body;
	while (lastm->b_cont != NULL)
		lastm = lastm->b_cont;
	lastm->b_cont = thdr;

	/*
	 * The mblk chain is ready. Decrypt!
	 */
	rc = nsmb_decrypt_mblks(&ctx, body, tlen);
	/* check rc below */

	/* Unlink thdr and restore r/w pointers. */
	lastm->b_cont = NULL;
	thdr->b_rptr -= SMB3_SIG_OFFS;
	thdr->b_wptr = thdr->b_rptr + SMB3_TFORM_HDR_SIZE;

	/* Now check rc from decrypt */
	if (rc != 0)
		goto errout;

	/*
	 * Lastly, discard the transform header
	 * and return the body.
	 */
	freeb(thdr);
	*mpp = body;
	nsmb_enc_ctx_done(&ctx);
	return (0);

errout:
	freeb(thdr);
	nsmb_enc_ctx_done(&ctx);
	return (rc);
}