/*************************************************************************
* Copyright (C) 2021 Intel Corporation
*
* Licensed under the Apache License,  Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* 	http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law  or agreed  to  in  writing,  software
* distributed under  the License  is  distributed  on  an  "AS IS"  BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the  specific  language  governing  permissions  and
* limitations under the License.
*************************************************************************/

#include <crypto_mb/status.h>
#include <crypto_mb/sm4.h>

#include <internal/sm4/sm4_mb.h>
#include <internal/common/ifma_defs.h>
#include <internal/rsa/ifma_rsa_arith.h>

/* FK[] constants */
static const int32u SM4_FK[4] = {
    0xA3B1BAC6,0x56AA3350,0x677D9197,0xB27022DC
};

/* CK[] constants */
static const int32u SM4_CK[32] = {
    0x00070E15,0x1C232A31,0x383F464D,0x545B6269,
    0x70777E85,0x8C939AA1,0xA8AFB6BD,0xC4CBD2D9,
    0xE0E7EEF5,0xFC030A11,0x181F262D,0x343B4249,
    0x50575E65,0x6C737A81,0x888F969D,0xA4ABB2B9,
    0xC0C7CED5,0xDCE3EAF1,0xF8FF060D,0x141B2229,
    0x30373E45,0x4C535A61,0x686F767D,0x848B9299,
    0xA0A7AEB5,0xBCC3CAD1,0xD8DFE6ED,0xF4FB0209,
    0x10171E25,0x2C333A41,0x484F565D,0x646B7279
};

#define SM4_ONE_RK(K0, K1, K2, K3, TMP, CK, OUT) {  \
    /* (Ki+1 ^ Ki+2 ^ Ki+3 ^ CKi) */                \
    TMP = _mm512_xor_epi32(_mm512_xor_epi32(_mm512_xor_epi32(K1, K2), K3), _mm512_set1_epi32(CK)); \
    /* T'(Ki+1 ^ Ki+2 ^ Ki+3 ^ CKi) */              \
    TMP = sBox512(TMP);                             \
    TMP = _mm512_xor_epi32(TMP, Lkey512(TMP));      \
    /* Ki+4 = Ki ^ T'(Ki+1 ^ Ki+2 ^ Ki+3 ^ CKi) */  \
    K0 = _mm512_xor_epi32(K0, TMP);                 \
    _mm512_storeu_si512((void*)OUT, K0);                   \
}

#define SM4_FOUR_RK(K0, K1, K2, K3, TMP, CK, OUT) {                         \
    SM4_ONE_RK(K0, K1, K2, K3, TMP, CK[0], OUT);          \
    SM4_ONE_RK(K1, K2, K3, K0, TMP, CK[1], (OUT + 1));      \
    SM4_ONE_RK(K2, K3, K0, K1, TMP, CK[2], (OUT + 2));     \
    SM4_ONE_RK(K3, K0, K1, K2, TMP, CK[3], (OUT + 3));     \
}


void sm4_set_round_keys_mb16(int32u* key_sched[SM4_ROUNDS], const int8u* pa_inp_key[SM4_LINES], __mmask16 mb_mask)
{
    __m512i rki = _mm512_setzero_si512();
    __m512i z0, z1, z2, z3;

    TRANSPOSE_16x4_I32_EPI32(&z0, &z1, &z2, &z3, pa_inp_key, mb_mask);

    /* (K0, K1, K2, K3) = (MK0 ^ FK0, MK1 ^ FK1, MK2 ^ FK2, MK3 ^ FK3) */
    z0 = _mm512_xor_epi32(z0, _mm512_set1_epi32(SM4_FK[0]));
    z1 = _mm512_xor_epi32(z1, _mm512_set1_epi32(SM4_FK[1]));
    z2 = _mm512_xor_epi32(z2, _mm512_set1_epi32(SM4_FK[2]));
    z3 = _mm512_xor_epi32(z3, _mm512_set1_epi32(SM4_FK[3]));

    const int32u* pCK = SM4_CK;
    const __m512i* p_rk = (const __m512i*)key_sched;

    int itr;
    for (itr = 0; itr < SM4_ROUNDS; itr += 4, pCK += 4, p_rk += 4)
        SM4_FOUR_RK(z0, z1, z2, z3, rki, pCK, p_rk);

    /* clear copies of sensitive data and round keys */
    zero_mb8((int64u(*)[8])&z0, 1);
    zero_mb8((int64u(*)[8])&z1, 1);
    zero_mb8((int64u(*)[8])&z2, 1);
    zero_mb8((int64u(*)[8])&z3, 1);
    zero_mb8((int64u(*)[8])&rki, 1);
}

DLL_PUBLIC
mbx_status16 mbx_sm4_set_key_mb16(mbx_sm4_key_schedule* key_sched, const sm4_key* pa_key[SM4_LINES])
{
    int buf_no;
    mbx_status16 status = 0;
    __mmask16 mb_mask = 0xFFFF;

    /* Test input pointers */
    if (NULL == key_sched || NULL == pa_key) {
        status = MBX_SET_STS16_ALL(MBX_STATUS_NULL_PARAM_ERR);
        return status;
    }

    /* Don't process buffers with input pointers equal to zero */
    for (buf_no = 0; buf_no < SM4_LINES; buf_no++) {
        if (pa_key[buf_no] == NULL) {
            status = MBX_SET_STS16(status, buf_no, MBX_STATUS_NULL_PARAM_ERR);
            mb_mask &= ~(0x1 << buf_no);
        }
    }

    if (MBX_IS_ANY_OK_STS16(status))
        sm4_set_round_keys_mb16((int32u**)key_sched, (const int8u**)pa_key, mb_mask);

    return status;
}

DLL_PUBLIC
mbx_status16 mbx_sm4_xts_set_keys_mb16(mbx_sm4_key_schedule* key_sched1,
                                       mbx_sm4_key_schedule* key_sched2,
                                       const sm4_xts_key* pa_key[SM4_LINES])
{
    int buf_no;
    mbx_status16 status = 0;
    __mmask16 mb_mask = 0xFFFF;

    /* Test input pointers */
    if (NULL == key_sched1 || NULL == key_sched2 || NULL == pa_key)
        return MBX_SET_STS16_ALL(MBX_STATUS_NULL_PARAM_ERR);

    /* Don't process buffers with input pointers equal to zero */
    for (buf_no = 0; buf_no < SM4_LINES; buf_no++) {
        if (pa_key[buf_no] == NULL) {
            status = MBX_SET_STS16(status, buf_no, MBX_STATUS_NULL_PARAM_ERR);
            mb_mask &= ~(0x1 << buf_no);
        }
    }

    if (MBX_IS_ANY_OK_STS16(status)) {
        /* Generate round keys for key1 */
        sm4_set_round_keys_mb16((int32u**)key_sched1, (const int8u**)pa_key, mb_mask);

        const sm4_key* pa_key2[SM4_LINES];

        for (int i = 0; i < SM4_LINES; i++)
            pa_key2[i] = (const sm4_key*)&((int8u*)pa_key[i])[16];

        /* Generate round keys for key2 */
        sm4_set_round_keys_mb16((int32u**)key_sched2, (const int8u**)pa_key2, mb_mask);
    }

    return status;
}
