// SPDX-License-Identifier: GPL-2.0

#include <linux/ceph/ceph_debug.h>

#include <linux/err.h>
#include <linux/scatterlist.h>
#include <linux/sched.h>
#include <linux/slab.h>
#include <crypto/aes.h>
#include <crypto/krb5.h>
#include <crypto/skcipher.h>
#include <linux/key-type.h>
#include <linux/sched/mm.h>

#include <keys/ceph-type.h>
#include <keys/user-type.h>
#include <linux/ceph/decode.h>
#include "crypto.h"

static int set_aes_tfm(struct ceph_crypto_key *key)
{
        unsigned int noio_flag;
        int ret;

        noio_flag = memalloc_noio_save();
        key->aes_tfm = crypto_alloc_sync_skcipher("cbc(aes)", 0, 0);
        memalloc_noio_restore(noio_flag);
        if (IS_ERR(key->aes_tfm)) {
                ret = PTR_ERR(key->aes_tfm);
                key->aes_tfm = NULL;
                return ret;
        }

        ret = crypto_sync_skcipher_setkey(key->aes_tfm, key->key, key->len);
        if (ret)
                return ret;

        return 0;
}

static int set_krb5_tfms(struct ceph_crypto_key *key, const u32 *key_usages,
                         int key_usage_cnt)
{
        struct krb5_buffer TK = { .len = key->len, .data = key->key };
        unsigned int noio_flag;
        int ret = 0;
        int i;

        if (WARN_ON_ONCE(key_usage_cnt > ARRAY_SIZE(key->krb5_tfms)))
                return -EINVAL;

        key->krb5_type = crypto_krb5_find_enctype(
                             KRB5_ENCTYPE_AES256_CTS_HMAC_SHA384_192);
        if (!key->krb5_type)
                return -ENOPKG;

        /*
         * Despite crypto_krb5_prepare_encryption() taking a gfp mask,
         * crypto_alloc_aead() inside of it allocates with GFP_KERNEL.
         */
        noio_flag = memalloc_noio_save();
        for (i = 0; i < key_usage_cnt; i++) {
                key->krb5_tfms[i] = crypto_krb5_prepare_encryption(
                                        key->krb5_type, &TK, key_usages[i],
                                        GFP_NOIO);
                if (IS_ERR(key->krb5_tfms[i])) {
                        ret = PTR_ERR(key->krb5_tfms[i]);
                        key->krb5_tfms[i] = NULL;
                        goto out_flag;
                }
        }

out_flag:
        memalloc_noio_restore(noio_flag);
        return ret;
}

int ceph_crypto_key_prepare(struct ceph_crypto_key *key,
                            const u32 *key_usages, int key_usage_cnt)
{
        switch (key->type) {
        case CEPH_CRYPTO_NONE:
                return 0; /* nothing to do */
        case CEPH_CRYPTO_AES:
                return set_aes_tfm(key);
        case CEPH_CRYPTO_AES256KRB5:
                hmac_sha256_preparekey(&key->hmac_key, key->key, key->len);
                return set_krb5_tfms(key, key_usages, key_usage_cnt);
        default:
                return -ENOTSUPP;
        }
}

/*
 * @dst should be zeroed before this function is called.
 */
int ceph_crypto_key_clone(struct ceph_crypto_key *dst,
                          const struct ceph_crypto_key *src)
{
        dst->type = src->type;
        dst->created = src->created;
        dst->len = src->len;

        dst->key = kmemdup(src->key, src->len, GFP_NOIO);
        if (!dst->key)
                return -ENOMEM;

        return 0;
}

/*
 * @key should be zeroed before this function is called.
 */
int ceph_crypto_key_decode(struct ceph_crypto_key *key, void **p, void *end)
{
        ceph_decode_need(p, end, 2*sizeof(u16) + sizeof(key->created), bad);
        key->type = ceph_decode_16(p);
        ceph_decode_copy(p, &key->created, sizeof(key->created));
        key->len = ceph_decode_16(p);
        ceph_decode_need(p, end, key->len, bad);
        if (key->len > CEPH_MAX_KEY_LEN) {
                pr_err("secret too big %d\n", key->len);
                return -EINVAL;
        }

        key->key = kmemdup(*p, key->len, GFP_NOIO);
        if (!key->key)
                return -ENOMEM;

        memzero_explicit(*p, key->len);
        *p += key->len;
        return 0;

bad:
        dout("failed to decode crypto key\n");
        return -EINVAL;
}

int ceph_crypto_key_unarmor(struct ceph_crypto_key *key, const char *inkey)
{
        int inlen = strlen(inkey);
        int blen = inlen * 3 / 4;
        void *buf, *p;
        int ret;

        dout("crypto_key_unarmor %s\n", inkey);
        buf = kmalloc(blen, GFP_NOFS);
        if (!buf)
                return -ENOMEM;
        blen = ceph_unarmor(buf, inkey, inkey+inlen);
        if (blen < 0) {
                kfree(buf);
                return blen;
        }

        p = buf;
        ret = ceph_crypto_key_decode(key, &p, p + blen);
        kfree(buf);
        if (ret)
                return ret;
        dout("crypto_key_unarmor key %p type %d len %d\n", key,
             key->type, key->len);
        return 0;
}

void ceph_crypto_key_destroy(struct ceph_crypto_key *key)
{
        int i;

        if (!key)
                return;

        kfree_sensitive(key->key);
        key->key = NULL;

        if (key->type == CEPH_CRYPTO_AES) {
                if (key->aes_tfm) {
                        crypto_free_sync_skcipher(key->aes_tfm);
                        key->aes_tfm = NULL;
                }
        } else if (key->type == CEPH_CRYPTO_AES256KRB5) {
                memzero_explicit(&key->hmac_key, sizeof(key->hmac_key));
                for (i = 0; i < ARRAY_SIZE(key->krb5_tfms); i++) {
                        if (key->krb5_tfms[i]) {
                                crypto_free_aead(key->krb5_tfms[i]);
                                key->krb5_tfms[i] = NULL;
                        }
                }
        }
}

static const u8 *aes_iv = (u8 *)CEPH_AES_IV;

/*
 * Should be used for buffers allocated with kvmalloc().
 * Currently these are encrypt out-buffer (ceph_buffer) and decrypt
 * in-buffer (msg front).
 *
 * Dispose of @sgt with teardown_sgtable().
 *
 * @prealloc_sg is to avoid memory allocation inside sg_alloc_table()
 * in cases where a single sg is sufficient.  No attempt to reduce the
 * number of sgs by squeezing physically contiguous pages together is
 * made though, for simplicity.
 */
static int setup_sgtable(struct sg_table *sgt, struct scatterlist *prealloc_sg,
                         const void *buf, unsigned int buf_len)
{
        struct scatterlist *sg;
        const bool is_vmalloc = is_vmalloc_addr(buf);
        unsigned int off = offset_in_page(buf);
        unsigned int chunk_cnt = 1;
        unsigned int chunk_len = PAGE_ALIGN(off + buf_len);
        int i;
        int ret;

        if (buf_len == 0) {
                memset(sgt, 0, sizeof(*sgt));
                return -EINVAL;
        }

        if (is_vmalloc) {
                chunk_cnt = chunk_len >> PAGE_SHIFT;
                chunk_len = PAGE_SIZE;
        }

        if (chunk_cnt > 1) {
                ret = sg_alloc_table(sgt, chunk_cnt, GFP_NOFS);
                if (ret)
                        return ret;
        } else {
                WARN_ON(chunk_cnt != 1);
                sg_init_table(prealloc_sg, 1);
                sgt->sgl = prealloc_sg;
                sgt->nents = sgt->orig_nents = 1;
        }

        for_each_sg(sgt->sgl, sg, sgt->orig_nents, i) {
                struct page *page;
                unsigned int len = min(chunk_len - off, buf_len);

                if (is_vmalloc)
                        page = vmalloc_to_page(buf);
                else
                        page = virt_to_page(buf);

                sg_set_page(sg, page, len, off);

                off = 0;
                buf += len;
                buf_len -= len;
        }
        WARN_ON(buf_len != 0);

        return 0;
}

static void teardown_sgtable(struct sg_table *sgt)
{
        if (sgt->orig_nents > 1)
                sg_free_table(sgt);
}

static int ceph_aes_crypt(const struct ceph_crypto_key *key, bool encrypt,
                          void *buf, int buf_len, int in_len, int *pout_len)
{
        SYNC_SKCIPHER_REQUEST_ON_STACK(req, key->aes_tfm);
        struct sg_table sgt;
        struct scatterlist prealloc_sg;
        char iv[AES_BLOCK_SIZE] __aligned(8);
        int pad_byte = AES_BLOCK_SIZE - (in_len & (AES_BLOCK_SIZE - 1));
        int crypt_len = encrypt ? in_len + pad_byte : in_len;
        int ret;

        WARN_ON(crypt_len > buf_len);
        if (encrypt)
                memset(buf + in_len, pad_byte, pad_byte);
        ret = setup_sgtable(&sgt, &prealloc_sg, buf, crypt_len);
        if (ret)
                return ret;

        memcpy(iv, aes_iv, AES_BLOCK_SIZE);
        skcipher_request_set_sync_tfm(req, key->aes_tfm);
        skcipher_request_set_callback(req, 0, NULL, NULL);
        skcipher_request_set_crypt(req, sgt.sgl, sgt.sgl, crypt_len, iv);

        /*
        print_hex_dump(KERN_ERR, "key: ", DUMP_PREFIX_NONE, 16, 1,
                       key->key, key->len, 1);
        print_hex_dump(KERN_ERR, " in: ", DUMP_PREFIX_NONE, 16, 1,
                       buf, crypt_len, 1);
        */
        if (encrypt)
                ret = crypto_skcipher_encrypt(req);
        else
                ret = crypto_skcipher_decrypt(req);
        skcipher_request_zero(req);
        if (ret) {
                pr_err("%s %scrypt failed: %d\n", __func__,
                       encrypt ? "en" : "de", ret);
                goto out_sgt;
        }
        /*
        print_hex_dump(KERN_ERR, "out: ", DUMP_PREFIX_NONE, 16, 1,
                       buf, crypt_len, 1);
        */

        if (encrypt) {
                *pout_len = crypt_len;
        } else {
                pad_byte = *(char *)(buf + in_len - 1);
                if (pad_byte > 0 && pad_byte <= AES_BLOCK_SIZE &&
                    in_len >= pad_byte) {
                        *pout_len = in_len - pad_byte;
                } else {
                        pr_err("%s got bad padding %d on in_len %d\n",
                               __func__, pad_byte, in_len);
                        ret = -EPERM;
                        goto out_sgt;
                }
        }

out_sgt:
        teardown_sgtable(&sgt);
        return ret;
}

static int ceph_krb5_encrypt(const struct ceph_crypto_key *key, int usage_slot,
                             void *buf, int buf_len, int in_len, int *pout_len)
{
        struct sg_table sgt;
        struct scatterlist prealloc_sg;
        int ret;

        if (WARN_ON_ONCE(usage_slot >= ARRAY_SIZE(key->krb5_tfms)))
                return -EINVAL;

        ret = setup_sgtable(&sgt, &prealloc_sg, buf, buf_len);
        if (ret)
                return ret;

        ret = crypto_krb5_encrypt(key->krb5_type, key->krb5_tfms[usage_slot],
                                  sgt.sgl, sgt.nents, buf_len, AES_BLOCK_SIZE,
                                  in_len, false);
        if (ret < 0) {
                pr_err("%s encrypt failed: %d\n", __func__, ret);
                goto out_sgt;
        }

        *pout_len = ret;
        ret = 0;

out_sgt:
        teardown_sgtable(&sgt);
        return ret;
}

static int ceph_krb5_decrypt(const struct ceph_crypto_key *key, int usage_slot,
                             void *buf, int buf_len, int in_len, int *pout_len)
{
        struct sg_table sgt;
        struct scatterlist prealloc_sg;
        size_t data_off = 0;
        size_t data_len = in_len;
        int ret;

        if (WARN_ON_ONCE(usage_slot >= ARRAY_SIZE(key->krb5_tfms)))
                return -EINVAL;

        ret = setup_sgtable(&sgt, &prealloc_sg, buf, in_len);
        if (ret)
                return ret;

        ret = crypto_krb5_decrypt(key->krb5_type, key->krb5_tfms[usage_slot],
                                  sgt.sgl, sgt.nents, &data_off, &data_len);
        if (ret) {
                pr_err("%s decrypt failed: %d\n", __func__, ret);
                goto out_sgt;
        }

        WARN_ON(data_off != AES_BLOCK_SIZE);
        *pout_len = data_len;

out_sgt:
        teardown_sgtable(&sgt);
        return ret;
}

int ceph_crypt(const struct ceph_crypto_key *key, int usage_slot, bool encrypt,
               void *buf, int buf_len, int in_len, int *pout_len)
{
        switch (key->type) {
        case CEPH_CRYPTO_NONE:
                *pout_len = in_len;
                return 0;
        case CEPH_CRYPTO_AES:
                return ceph_aes_crypt(key, encrypt, buf, buf_len, in_len,
                                      pout_len);
        case CEPH_CRYPTO_AES256KRB5:
                return encrypt ?
                    ceph_krb5_encrypt(key, usage_slot, buf, buf_len, in_len,
                                      pout_len) :
                    ceph_krb5_decrypt(key, usage_slot, buf, buf_len, in_len,
                                      pout_len);
        default:
                return -ENOTSUPP;
        }
}

int ceph_crypt_data_offset(const struct ceph_crypto_key *key)
{
        switch (key->type) {
        case CEPH_CRYPTO_NONE:
        case CEPH_CRYPTO_AES:
                return 0;
        case CEPH_CRYPTO_AES256KRB5:
                /* confounder */
                return AES_BLOCK_SIZE;
        default:
                BUG();
        }
}

int ceph_crypt_buflen(const struct ceph_crypto_key *key, int data_len)
{
        switch (key->type) {
        case CEPH_CRYPTO_NONE:
                return data_len;
        case CEPH_CRYPTO_AES:
                /* PKCS#7 padding at the end */
                return data_len + AES_BLOCK_SIZE -
                       (data_len & (AES_BLOCK_SIZE - 1));
        case CEPH_CRYPTO_AES256KRB5:
                /* confounder at the beginning and 192-bit HMAC at the end */
                return AES_BLOCK_SIZE + data_len + 24;
        default:
                BUG();
        }
}

void ceph_hmac_sha256(const struct ceph_crypto_key *key, const void *buf,
                      int buf_len, u8 hmac[SHA256_DIGEST_SIZE])
{
        switch (key->type) {
        case CEPH_CRYPTO_NONE:
        case CEPH_CRYPTO_AES:
                memset(hmac, 0, SHA256_DIGEST_SIZE);
                return;
        case CEPH_CRYPTO_AES256KRB5:
                hmac_sha256(&key->hmac_key, buf, buf_len, hmac);
                return;
        default:
                BUG();
        }
}

static int ceph_key_preparse(struct key_preparsed_payload *prep)
{
        struct ceph_crypto_key *ckey;
        size_t datalen = prep->datalen;
        int ret;
        void *p;

        ret = -EINVAL;
        if (datalen <= 0 || datalen > 32767 || !prep->data)
                goto err;

        ret = -ENOMEM;
        ckey = kzalloc_obj(*ckey);
        if (!ckey)
                goto err;

        /* TODO ceph_crypto_key_decode should really take const input */
        p = (void *)prep->data;
        ret = ceph_crypto_key_decode(ckey, &p, (char*)prep->data+datalen);
        if (ret < 0)
                goto err_ckey;

        prep->payload.data[0] = ckey;
        prep->quotalen = datalen;
        return 0;

err_ckey:
        kfree(ckey);
err:
        return ret;
}

static void ceph_key_free_preparse(struct key_preparsed_payload *prep)
{
        struct ceph_crypto_key *ckey = prep->payload.data[0];
        ceph_crypto_key_destroy(ckey);
        kfree(ckey);
}

static void ceph_key_destroy(struct key *key)
{
        struct ceph_crypto_key *ckey = key->payload.data[0];

        ceph_crypto_key_destroy(ckey);
        kfree(ckey);
}

struct key_type key_type_ceph = {
        .name           = "ceph",
        .preparse       = ceph_key_preparse,
        .free_preparse  = ceph_key_free_preparse,
        .instantiate    = generic_key_instantiate,
        .destroy        = ceph_key_destroy,
};

int __init ceph_crypto_init(void)
{
        return register_key_type(&key_type_ceph);
}

void ceph_crypto_shutdown(void)
{
        unregister_key_type(&key_type_ceph);
}