/*

 * Copyright 2024-2025 The OpenSSL Project Authors. All Rights Reserved.

 *

 * Licensed under the Apache License 2.0 (the "License").  You may not use

 * this file except in compliance with the License.  You can obtain a copy

 * in the file LICENSE in the source distribution or at

 * https://www.openssl.org/source/license.html

 */

#include <openssl/byteorder.h>

#include <openssl/err.h>

#include <openssl/evp.h>

#include <openssl/proverr.h>

#include "ml_dsa_hash.h"

#include "ml_dsa_key.h"

#include "ml_dsa_sign.h"

#include "internal/packet.h"

#define POLY_COEFF_NUM_BYTES(bits)  ((bits) * (ML_DSA_NUM_POLY_COEFFICIENTS / 8))

/* Cast mod_sub result in support of left-shifts that create 64-bit values. */

#define mod_sub_64(a, b) ((uint64_t) mod_sub(a, b))

typedef int (ENCODE_FN)(const POLY *s, WPACKET *pkt);

typedef int (DECODE_FN)(POLY *s, PACKET *pkt);

static ENCODE_FN poly_encode_signed_2;

static ENCODE_FN poly_encode_signed_4;

static ENCODE_FN poly_encode_signed_two_to_power_17;

static ENCODE_FN poly_encode_signed_two_to_power_19;

static DECODE_FN poly_decode_signed_2;

static DECODE_FN poly_decode_signed_4;

static DECODE_FN poly_decode_signed_two_to_power_17;

static DECODE_FN poly_decode_signed_two_to_power_19;

/* Bit packing Algorithms */

/*

 * Encodes a polynomial into a byte string, assuming that all coefficients are

 * in the range 0..15 (4 bits).

 *

 * See FIPS 204, Algorithm 16, SimpleBitPack(w, b) where b = 4 bits

 *

 * i.e. Use 4 bits from each coefficient and pack them into bytes

 * So every 2 coefficients fit into 1 byte.

 *

 * This is used to encode w1 when signing with ML-DSA-65 and ML-DSA-87

 *

 * @param p A polynomial with coefficients all in the range (0..15)

 * @param pkt A packet object to write 128 bytes to.

 *

 * @returns 1 on success, or 0 on error.

 */

static int poly_encode_4_bits(const POLY *p, WPACKET *pkt)

{

    uint8_t *out;

    const uint32_t *in = p->coeff, *end = in + ML_DSA_NUM_POLY_COEFFICIENTS;

    if (!WPACKET_allocate_bytes(pkt, POLY_COEFF_NUM_BYTES(4), &out))

        return 0;

    do {

        uint32_t z0 = *in++;

        uint32_t z1 = *in++;

        *out++ = z0 | (z1 << 4);

    } while (in < end);

    return 1;

}

/*

 * Encodes a polynomial into a byte string, assuming that all coefficients are

 * in the range 0..43 (6 bits).

 *

 * See FIPS 204, Algorithm 16, SimpleBitPack(w, b) where b = 43

 *

 * i.e. Use 6 bits from each coefficient and pack them into bytes

 * So every 4 coefficients fit into 3 bytes.

 *

 *  |c0||c1||c2||c3|

 *   |  /|  /\  /

 *  |6 2|4 4|2 6|

 *

 * This is used to encode w1 when signing with ML-DSA-44

 *

 * @param p A polynomial with coefficients all in the range (0..43)

 * @param pkt A packet object to write 96 bytes to.

 *

 * @returns 1 on success, or 0 on error.

 */

static int poly_encode_6_bits(const POLY *p, WPACKET *pkt)

{

    uint8_t *out;

    const uint32_t *in = p->coeff, *end = in + ML_DSA_NUM_POLY_COEFFICIENTS;

    if (!WPACKET_allocate_bytes(pkt, POLY_COEFF_NUM_BYTES(6), &out))

        return 0;

    do {

        uint32_t c0 = *in++;

        uint32_t c1 = *in++;

        uint32_t c2 = *in++;

        uint32_t c3 = *in++;

        *out++ = c0 | (c1 << 6);

        *out++ = (c1 >> 2) | (c2 << 4);

        *out++ = (c2 >> 4) | (c3 << 2);

    } while (in < end);

    return 1;

}

/*

 * Encodes a polynomial into a byte string, assuming that all coefficients are

 * unsigned 10 bit values.

 *

 * See FIPS 204, Algorithm 16, SimpleBitPack(w, b) where b = 10 bits

 *

 * i.e. Use 10 bits from each coefficient and pack them into bytes

 * So every 4 coefficients (c0..c3) fit into 5 bytes.

 *  |c0||c1||c2||c3|

 *   |\  |\  |\  |\

 *  |8|2 6|4 4|6 2|8|

 *

 * This is used to save t1 (the high part of public key polynomial t)

 *

 * @param p A polynomial with coefficients all in the range (0..1023)

 * @param pkt A packet object to write 320 bytes to.

 *

 * @returns 1 on success, or 0 on error.

 */

static int poly_encode_10_bits(const POLY *p, WPACKET *pkt)

{

    uint8_t *out;

    const uint32_t *in = p->coeff, *end = in + ML_DSA_NUM_POLY_COEFFICIENTS;

    if (!WPACKET_allocate_bytes(pkt, POLY_COEFF_NUM_BYTES(10), &out))

        return 0;

    do {

        uint32_t c0 = *in++;

        uint32_t c1 = *in++;

        uint32_t c2 = *in++;

        uint32_t c3 = *in++;

        *out++ = (uint8_t)c0;

        *out++ = (uint8_t)((c0 >> 8) | (c1 << 2));

        *out++ = (uint8_t)((c1 >> 6) | (c2 << 4));

        *out++ = (uint8_t)((c2 >> 4) | (c3 << 6));

        *out++ = (uint8_t)(c3 >> 2);

    } while (in < end);

    return 1;

}

/*

 * @brief Reverses the procedure of poly_encode_10_bits().

 * See FIPS 204, Algorithm 18, SimpleBitUnpack(v, b) where b = 10.

 *

 * @param p A polynomial to write coefficients to.

 * @param pkt A packet object to read 320 bytes from.

 *

 * @returns 1 on success, or 0 on error.

 */

static int poly_decode_10_bits(POLY *p, PACKET *pkt)

{

    const uint8_t *in = NULL;

    uint32_t v, w, mask = 0x3ff; /* 10 bits */

    uint32_t *out = p->coeff, *end = out + ML_DSA_NUM_POLY_COEFFICIENTS;

    do {

        if (!PACKET_get_bytes(pkt, &in, 5))

            return 0;

        in = OPENSSL_load_u32_le(&v, in);

        w = *in;

        *out++ = v & mask;

        *out++ = (v >> 10) & mask;

        *out++ = (v >> 20) & mask;

        *out++ = (v >> 30) | (w << 2);

    } while (out < end);

    return 1;

}

/*

 * @brief Encodes a polynomial into a byte string, assuming that all

 * coefficients are in the range -4..4.

 * See FIPS 204, Algorithm 17, BitPack(w, a, b). (a = 4, b = 4)

 *

 * It uses a nibble from each coefficient and packs them into bytes

 * So every 2 coefficients fit into 1 byte.

 *

 * This is used to encode the private key polynomial elements of s1 and s2

 * for ML-DSA-65 (i.e. eta = 4)

 *

 * @param p An array of 256 coefficients all in the range -4..4

 * @param pkt A packet to write 128 bytes of encoded polynomial coefficients to.

 *

 * @returns 1 on success, or 0 on error.

 */

static int poly_encode_signed_4(const POLY *p, WPACKET *pkt)

{

    uint8_t *out;

    const uint32_t *in = p->coeff, *end = in + ML_DSA_NUM_POLY_COEFFICIENTS;

    if (!WPACKET_allocate_bytes(pkt, 32 * 4, &out))

        return 0;

    do {

        uint32_t z = mod_sub(4, *in++);

        *out++ = z | (mod_sub(4, *in++) << 4);

    } while (in < end);

    return 1;

}

/*

 * @brief Reverses the procedure of poly_encode_signed_4().

 * See FIPS 204, Algorithm 19, BitUnpack(v, a, b) where a = b = 4.

 *

 * @param p A polynomial to write coefficients to.

 * @param pkt A packet object to read 128 bytes from.

 *

 * @returns 1 on success, or 0 on error. An error will occur if any of the

 *          coefficients are not in the correct range.

 */

static int poly_decode_signed_4(POLY *p, PACKET *pkt)

{

    int i, ret = 0;

    uint32_t v, *out = p->coeff;

    const uint8_t *in;

    uint32_t msbs, mask;

    for (i = 0; i < (ML_DSA_NUM_POLY_COEFFICIENTS / 8); i++) {

        if (!PACKET_get_bytes(pkt, &in, 4))

            goto err;

        in = OPENSSL_load_u32_le(&v, in);

        /*

         * None of the nibbles may be >= 9. So if the MSB of any nibble is set,

         * none of the other bits may be set. First, select all the MSBs.

         */

        msbs = v & 0x88888888u;

        /* For each nibble where the MSB is set, form a mask of all the other bits. */

        mask = (msbs >> 1) | (msbs >> 2) | (msbs >> 3);

        /*

         * A nibble is only out of range in the case of invalid input, in which case

         * it is okay to leak the value.

         */

        if (value_barrier_32((mask & v) != 0))

            goto err;

        *out++ = mod_sub(4, v & 15);

        *out++ = mod_sub(4, (v >> 4) & 15);

        *out++ = mod_sub(4, (v >> 8) & 15);

        *out++ = mod_sub(4, (v >> 12) & 15);

        *out++ = mod_sub(4, (v >> 16) & 15);

        *out++ = mod_sub(4, (v >> 20) & 15);

        *out++ = mod_sub(4, (v >> 24) & 15);

        *out++ = mod_sub(4, v >> 28);

    }

    ret = 1;

 err:

    return ret;

}

/*

 * @brief Encodes a polynomial into a byte string, assuming that all

 * coefficients are in the range -2..2.

 * See FIPS 204, Algorithm 17, BitPack(w, a, b). where a = b = 2.

 *

 * This is used to encode the private key polynomial elements of s1 and s2

 * for ML-DSA-44 and ML-DSA-87 (i.e. eta = 2)

 *

 * @param pkt A packet to write 128 bytes of encoded polynomial coefficients to.

 * @param p An array of 256 coefficients all in the range -2..2

 *

 * Use 3 bits from each coefficient and pack them into bytes

 * So every 8 coefficients fit into 3 bytes.

 *  |c0 c1 c2 c3 c4 c5 c6 c7|

 *   | /  / | |  / / | |  /

 *  |3 3 2| 1 3 3 1| 2 3 3|

 *

 * @param p An array of 256 coefficients all in the range -2..2

 * @param pkt A packet to write 64 bytes of encoded polynomial coefficients to.

 *

 * @returns 1 on success, or 0 on error.

 */

static int poly_encode_signed_2(const POLY *p, WPACKET *pkt)

{

    uint8_t *out;

    const uint32_t *in = p->coeff, *end = in + ML_DSA_NUM_POLY_COEFFICIENTS;

    if (!WPACKET_allocate_bytes(pkt, POLY_COEFF_NUM_BYTES(3), &out))

        return 0;

    do {

        uint32_t z;

        z = mod_sub(2, *in++);

        z |= mod_sub(2, *in++) << 3;

        z |= mod_sub(2, *in++) << 6;

        z |= mod_sub(2, *in++) << 9;

        z |= mod_sub(2, *in++) << 12;

        z |= mod_sub(2, *in++) << 15;

        z |= mod_sub(2, *in++) << 18;

        z |= mod_sub(2, *in++) << 21;

        out = OPENSSL_store_u16_le(out, (uint16_t) z);

        *out++ = (uint8_t) (z >> 16);

    } while (in < end);

    return 1;

}

/*

 * @brief Reverses the procedure of poly_encode_signed_2().

 * See FIPS 204, Algorithm 19, BitUnpack(v, a, b) where a = b = 2.

 *

 * @param p A polynomial to write coefficients to.

 * @param pkt A packet object to read 64 encoded bytes from.

 *

 * @returns 1 on success, or 0 on error. An error will occur if any of the

 *          coefficients are not in the correct range.

 */

static int poly_decode_signed_2(POLY *p, PACKET *pkt)

{

    int i, ret = 0;

    uint32_t u = 0, v = 0, *out = p->coeff;

    uint32_t msbs, mask;

    const uint8_t *in;

    for (i = 0; i < (ML_DSA_NUM_POLY_COEFFICIENTS / 8); i++) {

        if (!PACKET_get_bytes(pkt, &in, 3))

            goto err;

        memcpy(&u, in, 3);

        OPENSSL_load_u32_le(&v, (uint8_t *)&u);

        /*

         * Each octal value (3 bits) must be <= 4, So if the MSB is set then the

         * bottom 2 bits must not be set.

         * First, select all the MSBs (Use octal representation for the mask)

         */

        msbs = v & 044444444;

        /* For each octal value where the MSB is set, form a mask of the 2 other bits. */

        mask = (msbs >> 1) | (msbs >> 2);

        /*

         * A nibble is only out of range in the case of invalid input, in which

         * case it is okay to leak the value.

         */

        if (value_barrier_32((mask & v) != 0))

            goto err;

        *out++ = mod_sub(2, v & 7);

        *out++ = mod_sub(2, (v >> 3) & 7);

        *out++ = mod_sub(2, (v >> 6) & 7);

        *out++ = mod_sub(2, (v >> 9) & 7);

        *out++ = mod_sub(2, (v >> 12) & 7);

        *out++ = mod_sub(2, (v >> 15) & 7);

        *out++ = mod_sub(2, (v >> 18) & 7);

        *out++ = mod_sub(2, (v >> 21) & 7);

    }

    ret = 1;

 err:

    return ret;

}

/*

 * @brief Encodes a polynomial into a byte string, assuming that all

 * coefficients are in the range (-2^12 + 1)..2^12.

 * See FIPS 204, Algorithm 17, BitPack(w, a, b). where a = 2^12 - 1, b = 2^12.

 *

 * This is used to encode the LSB of the public key polynomial elements of t0

 * (which are encoded as part of the encoded private key).

 *

 * Use 13 bits from each coefficient and pack them into bytes

 *

 * The code below packs them into 2 64 bits blocks by doing..

 *  z0 z1 z2 z3  z4  z5 z6  z7 0

 *  |   |  | |   / \  |  |  |  |

 * |13 13 13 13 12 |1 13 13 13 24

 *

 * @param p An array of 256 coefficients all in the range -2^12+1..2^12

 * @param pkt A packet to write 416 (13 * 256 / 8) bytes of encoded polynomial

 *            coefficients to.

 *

 * @returns 1 on success, or 0 on error.

 */

static int poly_encode_signed_two_to_power_12(const POLY *p, WPACKET *pkt)

{

    static const uint32_t range = 1u << 12;

    const uint32_t *in = p->coeff, *end = in + ML_DSA_NUM_POLY_COEFFICIENTS;

    do {

        uint8_t *out;

        uint64_t a1, a2;

        if (!WPACKET_allocate_bytes(pkt, 13, &out))

            return 0;

        a1 = mod_sub_64(range, *in++);

        a1 |= mod_sub_64(range, *in++) << 13;

        a1 |= mod_sub_64(range, *in++) << 26;

        a1 |= mod_sub_64(range, *in++) << 39;

        a1 |= (a2 = mod_sub_64(range, *in++)) << 52;

        a2 = (a2 >> 12) | (mod_sub_64(range, *in++) << 1);

        a2 |= mod_sub_64(range, *in++) << 14;

        a2 |= mod_sub_64(range, *in++) << 27;

        out = OPENSSL_store_u64_le(out, a1);

        out = OPENSSL_store_u32_le(out, (uint32_t) a2);

        *out = (uint8_t) (a2 >> 32);

    } while (in < end);

    return 1;

}

/*

 * @brief Reverses the procedure of poly_encode_signed_two_to_power_12().

 * See FIPS 204, Algorithm 19, BitUnpack(v, a, b) where a = 2^12 - 1, b = 2^12.

 *

 * @param p A polynomial to write coefficients to.

 * @param pkt A packet object to read 416 encoded bytes from.

 *

 * @returns 1 on success, or 0 on error.

 */

static int poly_decode_signed_two_to_power_12(POLY *p, PACKET *pkt)

{

    int i, ret = 0;

    uint32_t *out = p->coeff;

    const uint8_t *in;

    static const uint32_t range = 1u << 12;

    static const uint32_t mask_13_bits = (1u << 13) - 1;

    for (i = 0; i < (ML_DSA_NUM_POLY_COEFFICIENTS / 8); i++) {

        uint64_t a1;

        uint32_t a2, b13;

        if (!PACKET_get_bytes(pkt, &in, 13))

            goto err;

        in = OPENSSL_load_u64_le(&a1, in);

        in = OPENSSL_load_u32_le(&a2, in);

        b13 = (uint32_t) *in;

        *out++ = mod_sub(range, a1 & mask_13_bits);

        *out++ = mod_sub(range, (a1 >> 13) & mask_13_bits);

        *out++ = mod_sub(range, (a1 >> 26) & mask_13_bits);

        *out++ = mod_sub(range, (a1 >> 39) & mask_13_bits);

        *out++ = mod_sub(range, (a1 >> 52) | ((a2 << 12) & mask_13_bits));

        *out++ = mod_sub(range, (a2 >> 1) & mask_13_bits);

        *out++ = mod_sub(range, (a2 >> 14) & mask_13_bits);

        *out++ = mod_sub(range, (a2 >> 27) | (b13 << 5));

    }

    ret = 1;

 err:

    return ret;

}

/*

 * @brief Encodes a polynomial into a byte string, assuming that all

 * coefficients are in the range (-2^19 + 1)..2^19.

 * See FIPS 204, Algorithm 17, BitPack(w, a, b). where a = 2^19 - 1, b = 2^19.

 *

 * This is used to encode signatures for ML-DSA-65 & ML-DSA-87 (gamma1 = 2^19)

 *

 * Use 20 bits from each coefficient and pack them into bytes

 *

 * The code below packs every 4 (20 bit) coefficients into 10 bytes

 *  z0  z1  z2 z3

 *  |   |\  |  | \

 * |20 12|8 20 4|16

 *

 * @param p An array of 256 coefficients all in the range -2^19+1..2^19

 * @param pkt A packet to write 640 (20 * 256 / 8) bytes of encoded polynomial

 *            coefficients to.

 *

 * @returns 1 on success, or 0 on error.

 */

static int poly_encode_signed_two_to_power_19(const POLY *p, WPACKET *pkt)

{

    static const uint32_t range = 1u << 19;

    const uint32_t *in = p->coeff, *end = in + ML_DSA_NUM_POLY_COEFFICIENTS;

    do {

        uint32_t z0, z1, z2;

        uint8_t *out;

        if (!WPACKET_allocate_bytes(pkt, 10, &out))

            return 0;

        z0 = mod_sub(range, *in++);

        z0 |= (z1 = mod_sub(range, *in++)) << 20;

        z1 = (z1 >> 12) | (mod_sub(range, *in++) << 8);

        z1 |= (z2 = mod_sub(range, *in++)) << 28;

        out = OPENSSL_store_u32_le(out, z0);

        out = OPENSSL_store_u32_le(out, z1);

        out = OPENSSL_store_u16_le(out, (uint16_t) (z2 >> 4));

    } while (in < end);

    return 1;

}

/*

 * @brief Reverses the procedure of poly_encode_signed_two_to_power_19().

 * See FIPS 204, Algorithm 19, BitUnpack(v, a, b) where a = 2^19 - 1, b = 2^19.

 *

 * @param p A polynomial to write coefficients to.

 * @param pkt A packet object to read 640 encoded bytes from.

 *

 * @returns 1 on success, or 0 on error.

 */

static int poly_decode_signed_two_to_power_19(POLY *p, PACKET *pkt)

{

    int i, ret = 0;

    uint32_t *out = p->coeff;

    const uint8_t *in;

    static const uint32_t range = 1u << 19;

    static const uint32_t mask_20_bits = (1u << 20) - 1;

    for (i = 0; i < (ML_DSA_NUM_POLY_COEFFICIENTS / 4); i++) {

        uint32_t a1, a2;

        uint16_t a3;

        if (!PACKET_get_bytes(pkt, &in, 10))

            goto err;

        in = OPENSSL_load_u32_le(&a1, in);

        in = OPENSSL_load_u32_le(&a2, in);

        in = OPENSSL_load_u16_le(&a3, in);

        *out++ = mod_sub(range, a1 & mask_20_bits);

        *out++ = mod_sub(range, (a1 >> 20) | ((a2 & 0xFF) << 12));

        *out++ = mod_sub(range, (a2 >> 8) & mask_20_bits);

        *out++ = mod_sub(range, (a2 >> 28) | (a3 << 4));

    }

    ret = 1;

 err:

    return ret;

}

/*

 * @brief Encodes a polynomial into a byte string, assuming that all

 * coefficients are in the range (-2^17 + 1)..2^17.

 * See FIPS 204, Algorithm 17, BitPack(w, a, b). where a = 2^17 - 1, b = 2^17.

 *

 * This is used to encode signatures for ML-DSA-44 (where gamma1 = 2^17)

 *

 * Use 18 bits from each coefficient and pack them into bytes

 *

 * The code below packs every 4 (18 bit) coefficients into 9 bytes

 *  z0  z1  z2 z3

 *  |   |\  |  | \

 * |18 14|4 18 10| 8

 *

 * @param p An array of 256 coefficients all in the range -2^17+1..2^17

 * @param pkt A packet to write 576 (18 * 256 / 8) bytes of encoded polynomial

 *            coefficients to.

 *

 * @returns 1 on success, or 0 on error.

 */

static int poly_encode_signed_two_to_power_17(const POLY *p, WPACKET *pkt)

{

    static const uint32_t range = 1u << 17;

    const uint32_t *in = p->coeff, *end = in + ML_DSA_NUM_POLY_COEFFICIENTS;

    do {

        uint8_t *out;

        uint32_t z0, z1, z2;

        if (!WPACKET_allocate_bytes(pkt, 9, &out))

            return 0;

        z0 = mod_sub(range, *in++);

        z0 |= (z1 = mod_sub(range, *in++)) << 18;

        z1 = (z1 >> 14) | (mod_sub(range, *in++) << 4);

        z1 |= (z2 = mod_sub(range, *in++)) << 22;

        out = OPENSSL_store_u32_le(out, z0);

        out = OPENSSL_store_u32_le(out, z1);

        *out = z2 >> 10;

    } while (in < end);

    return 1;

}

/*

 * @brief Reverses the procedure of poly_encode_signed_two_to_power_17().

 * See FIPS 204, Algorithm 19, BitUnpack(v, a, b) where a = 2^17 - 1, b = 2^17.

 *

 * @param p A polynomial to write coefficients to.

 * @param pkt A packet object to read 576 encoded bytes from.

 *

 * @returns 1 on success, or 0 on error.

 */

static int poly_decode_signed_two_to_power_17(POLY *p, PACKET *pkt)

{

    uint32_t *out = p->coeff;

    const uint32_t *end = out + ML_DSA_NUM_POLY_COEFFICIENTS;

    const uint8_t *in;

    static const uint32_t range = 1u << 17;

    static const uint32_t mask_18_bits = (1u << 18) - 1;

    do {

        uint32_t a1, a2, a3;

        if (!PACKET_get_bytes(pkt, &in, 9))

            return 0;

        in = OPENSSL_load_u32_le(&a1, in);

        in = OPENSSL_load_u32_le(&a2, in);

        a3 = (uint32_t) *in;

        *out++ = mod_sub(range, a1 & mask_18_bits);

        *out++ = mod_sub(range, (a1 >> 18) | ((a2 & 0xF) << 14));

        *out++ = mod_sub(range, (a2 >> 4) & mask_18_bits);

        *out++ = mod_sub(range, (a2 >> 22) | (a3 << 10));

    } while (out < end);

    return 1;

}

/*

 * @brief Encode the public key as an array of bytes.

 * See FIPS 204, Algorithm 22, pkEncode().

 *

 * @param key A key object containing public key values. The encoded public

 *            key data is stored in this key.

 * @returns 1 if the public key was encoded successfully or 0 otherwise.

 */

int ossl_ml_dsa_pk_encode(ML_DSA_KEY *key)

{

    int ret = 0;

    size_t i, written = 0;

    const POLY *t1 = key->t1.poly;

    size_t t1_len = key->t1.num_poly;

    size_t enc_len = key->params->pk_len;

    uint8_t *enc = OPENSSL_malloc(enc_len);

    WPACKET pkt;

    if (enc == NULL)

        return 0;

    if (!WPACKET_init_static_len(&pkt, enc, enc_len, 0)

            || !WPACKET_memcpy(&pkt, key->rho, sizeof(key->rho)))

        goto err;

    for (i = 0; i < t1_len; i++)

        if (!poly_encode_10_bits(t1 + i, &pkt))

            goto err;

    if (!WPACKET_get_total_written(&pkt, &written)

            || written != enc_len)

        goto err;

    OPENSSL_free(key->pub_encoding);

    key->pub_encoding = enc;

    ret = 1;

err:

    WPACKET_finish(&pkt);

    if (ret == 0)

        OPENSSL_free(enc);

    return ret;

}

/*

 * @brief The reverse of ossl_ml_dsa_pk_encode().

 * See FIPS 204, Algorithm 23, pkDecode().

 *

 * @param in An encoded public key.

 * @param in_len The size of |in|

 * @param key A key object to store the decoded public key into.

 *

 * @returns 1 if the public key was decoded successfully or 0 otherwise.

 */

int ossl_ml_dsa_pk_decode(ML_DSA_KEY *key, const uint8_t *in, size_t in_len)

{

    int ret = 0;

    size_t i;

    PACKET pkt;

    EVP_MD_CTX *ctx;

    if (key->priv_encoding != NULL || key->pub_encoding != NULL)

        return 0; /* Do not allow key mutation */

    if (in_len != key->params->pk_len)

        return 0;

    if (!ossl_ml_dsa_key_pub_alloc(key))

        return 0;

    ctx = EVP_MD_CTX_new();

    if (ctx == NULL)

        goto err;

    if (!PACKET_buf_init(&pkt, in, in_len)

            || !PACKET_copy_bytes(&pkt, key->rho, sizeof(key->rho)))

        goto err;

    for (i = 0; i < key->t1.num_poly; i++)

        if (!poly_decode_10_bits(key->t1.poly + i, &pkt))

            goto err;

    /* cache the hash of the encoded public key */

    if (!shake_xof(ctx, key->shake256_md, in, in_len, key->tr, sizeof(key->tr)))

        goto err;

    key->pub_encoding = OPENSSL_memdup(in, in_len);

    ret = (key->pub_encoding != NULL);

err:

    EVP_MD_CTX_free(ctx);

    return ret;

}

/*

 * @brief Encode the private key as an array of bytes.

 * See FIPS 204, Algorithm 24, skEncode().

 *

 * @param key A key object containing private key values. The encoded private

 *            key data is stored in this key.

 * @returns 1 if the private key was encoded successfully or 0 otherwise.

 */

int ossl_ml_dsa_sk_encode(ML_DSA_KEY *key)

{

    int ret = 0;

    const ML_DSA_PARAMS *params = key->params;

    size_t i, written = 0, k = params->k, l = params->l;

    ENCODE_FN *encode_fn;

    size_t enc_len = params->sk_len;

    const POLY *t0 = key->t0.poly;

    WPACKET pkt;

    uint8_t *enc = OPENSSL_malloc(enc_len);

    if (enc == NULL)

        return 0;

    /* eta is the range of private key coefficients (-eta...eta) */

    if (params->eta == ML_DSA_ETA_4)

        encode_fn = poly_encode_signed_4;

    else

        encode_fn = poly_encode_signed_2;

    if (!WPACKET_init_static_len(&pkt, enc, enc_len, 0)

            || !WPACKET_memcpy(&pkt, key->rho, sizeof(key->rho))

            || !WPACKET_memcpy(&pkt, key->K, sizeof(key->K))

            || !WPACKET_memcpy(&pkt, key->tr, sizeof(key->tr)))

        goto err;

    for (i = 0; i < l; ++i)

        if (!encode_fn(key->s1.poly + i, &pkt))

            goto err;

    for (i = 0; i < k; ++i)

        if (!encode_fn(key->s2.poly + i, &pkt))

            goto err;

    for (i = 0; i < k; ++i)

        if (!poly_encode_signed_two_to_power_12(t0++, &pkt))

            goto err;

    if (!WPACKET_get_total_written(&pkt, &written)

            || written != enc_len)

        goto err;

    OPENSSL_clear_free(key->priv_encoding, enc_len);

    key->priv_encoding = enc;

    ret = 1;

err:

    WPACKET_finish(&pkt);

    if (ret == 0)

        OPENSSL_clear_free(enc, enc_len);

    return ret;

}

/*

 * @brief The reverse of ossl_ml_dsa_sk_encode().

 * See FIPS 204, Algorithm 24, skDecode().

 *

 * @param in An encoded private key.

 * @param in_len The size of |in|

 * @param key A key object to store the decoded private key into.

 *

 * @returns 1 if the private key was decoded successfully or 0 otherwise.

 */

int ossl_ml_dsa_sk_decode(ML_DSA_KEY *key, const uint8_t *in, size_t in_len)

{

    DECODE_FN *decode_fn;

    const ML_DSA_PARAMS *params = key->params;

    size_t i, k = params->k, l = params->l;

    uint8_t input_tr[ML_DSA_TR_BYTES];

    PACKET pkt;

    /* When loading from an explicit key, drop the seed. */

    OPENSSL_clear_free(key->seed, ML_DSA_SEED_BYTES);

    key->seed = NULL;

    /* Allow the key encoding to be already set to the provided pointer */

    if ((key->priv_encoding != NULL && key->priv_encoding != in)

        || key->pub_encoding != NULL)

        return 0; /* Do not allow key mutation */

    if (in_len != key->params->sk_len)

        return 0;

    if (!ossl_ml_dsa_key_priv_alloc(key))

        return 0;

    /* eta is the range of private key coefficients (-eta...eta) */

    if (params->eta == ML_DSA_ETA_4)

        decode_fn = poly_decode_signed_4;

    else

        decode_fn = poly_decode_signed_2;

    if (!PACKET_buf_init(&pkt, in, in_len)

            || !PACKET_copy_bytes(&pkt, key->rho, sizeof(key->rho))

            || !PACKET_copy_bytes(&pkt, key->K, sizeof(key->K))

            || !PACKET_copy_bytes(&pkt, input_tr, sizeof(input_tr)))

        return 0;

    for (i = 0; i < l; ++i)

        if (!decode_fn(key->s1.poly + i, &pkt))

            goto err;

    for (i = 0; i < k; ++i)

        if (!decode_fn(key->s2.poly + i, &pkt))

            goto err;

    for (i = 0; i < k; ++i)

        if (!poly_decode_signed_two_to_power_12(key->t0.poly + i, &pkt))

            goto err;

    if (PACKET_remaining(&pkt) != 0)

        goto err;

    if (key->priv_encoding == NULL

        && (key->priv_encoding = OPENSSL_memdup(in, in_len)) == NULL)

        goto err;

    /*

     * Computing the public key also computes its hash, which must be equal to

     * the |tr| value in the private key, else the key was corrupted.

     */

    if (!ossl_ml_dsa_key_public_from_private(key)

            || memcmp(input_tr, key->tr, sizeof(input_tr)) != 0) {

        ERR_raise_data(ERR_LIB_PROV, PROV_R_INVALID_KEY,

                       "%s private key does not match its pubkey part",

                       key->params->alg);

        ossl_ml_dsa_key_reset(key);

        goto err;

    }

    return 1;

 err:

    return 0;

}

/*

 * See FIPS 204, Algorithm 20, HintBitPack().

 * Hint is composed of k polynomials with binary coefficients where only 'omega'

 * of all the coefficients are set to 1.

 * This can be encoded as a byte array of 'omega' polynomial coefficient index

 * positions for the coefficients that are set, followed by

 * k values of the last coefficient index used in each polynomial.

 */

static int hint_bits_encode(const VECTOR *hint, WPACKET *pkt, uint32_t omega)

{

    int i, j, k = hint->num_poly;

    size_t coeff_index = 0;

    POLY *p = hint->poly;

    uint8_t *data;

    if (!WPACKET_allocate_bytes(pkt, omega + k, &data))

        return 0;

    memset(data, 0, omega + k);

    for (i = 0; i < k; i++, p++) {

        for (j = 0; j < ML_DSA_NUM_POLY_COEFFICIENTS; j++)

            if (p->coeff[j] != 0)

                data[coeff_index++] = j;

        data[omega + i] = (uint8_t)coeff_index;

    }

    return 1;

}

/*

 * @brief Reverse the process of hint_bits_encode()

 * See FIPS 204, Algorithm 21, HintBitUnpack()

 *

 * @returns 1 if the hints were successfully unpacked, or 0

 * if 'pkt' is too small or malformed.

 */

static int hint_bits_decode(VECTOR *hint, PACKET *pkt, uint32_t omega)

{

    size_t coeff_index = 0, k = hint->num_poly;

    const uint8_t *in, *limits;

    POLY *p = hint->poly, *end = p + k;

    if (!PACKET_get_bytes(pkt, &in, omega)

            || !PACKET_get_bytes(pkt, &limits, k))

        return 0;

    vector_zero(hint); /* Set all coefficients to zero */

    do {

        const uint32_t limit = *limits++;

        int last = -1;

        if (limit < coeff_index || limit > omega)

            return 0;

        while (coeff_index < limit) {

            int byte = in[coeff_index++];

            if (last >= 0 && byte <= last)

                return 0;

            last = byte;

            p->coeff[byte] = 1;

        }

    } while (++p < end);

    for (; coeff_index < omega; coeff_index++)

        if (in[coeff_index] != 0)

            return 0;

    return 1;

}

/*

 * @brief Encode a ML_DSA signature as an array of bytes.

 * See FIPS 204, Algorithm 26, sigEncode().

 *

 * @param

 * @param

 * @returns 1 if the signature was encoded successfully or 0 otherwise.

 */

int ossl_ml_dsa_sig_encode(const ML_DSA_SIG *sig, const ML_DSA_PARAMS *params,

                           uint8_t *out)

{

    int ret = 0;

    size_t i;

    ENCODE_FN *encode_fn;

    WPACKET pkt;

    if (out == NULL)

        return 0;

    if (params->gamma1 == ML_DSA_GAMMA1_TWO_POWER_19)

        encode_fn = poly_encode_signed_two_to_power_19;

    else

        encode_fn = poly_encode_signed_two_to_power_17;

    if (!WPACKET_init_static_len(&pkt, out, params->sig_len, 0)

            || !WPACKET_memcpy(&pkt, sig->c_tilde, sig->c_tilde_len))

        goto err;

    for (i = 0; i < sig->z.num_poly; ++i)

        if (!encode_fn(sig->z.poly + i, &pkt))

            goto err;

    if (!hint_bits_encode(&sig->hint, &pkt, params->omega))

        goto err;

    ret = 1;

err:

    WPACKET_finish(&pkt);

    return ret;

}

/*

 * @param sig is a initialized signature object to decode into.

 * @param in An encoded signature

 * @param in_len The size of |in|

 * @param params contains constants for an ML-DSA algorithm (such as gamma1)

 * @returns 1 if the signature was successfully decoded or 0 otherwise.

 */

int ossl_ml_dsa_sig_decode(ML_DSA_SIG *sig, const uint8_t *in, size_t in_len,

                           const ML_DSA_PARAMS *params)

{

    int ret = 0;

    size_t i;

    DECODE_FN *decode_fn;

    PACKET pkt;

    if (params->gamma1 == ML_DSA_GAMMA1_TWO_POWER_19)

        decode_fn = poly_decode_signed_two_to_power_19;

    else

        decode_fn = poly_decode_signed_two_to_power_17;

    if (!PACKET_buf_init(&pkt, in, in_len)

            || !PACKET_copy_bytes(&pkt, sig->c_tilde, sig->c_tilde_len))

        goto err;

    for (i = 0; i < sig->z.num_poly; ++i)

        if (!decode_fn(sig->z.poly + i, &pkt))

            goto err;

    if (!hint_bits_decode(&sig->hint, &pkt, params->omega)

            || PACKET_remaining(&pkt) != 0)

        goto err;

    ret = 1;

err:

    return ret;

}

int ossl_ml_dsa_poly_decode_expand_mask(POLY *out,

                                        const uint8_t *in, size_t in_len,

                                        uint32_t gamma1)

{

    PACKET pkt;

    if (!PACKET_buf_init(&pkt, in, in_len))

        return 0;

    if (gamma1 == ML_DSA_GAMMA1_TWO_POWER_19)

        return poly_decode_signed_two_to_power_19(out, &pkt);

    else

        return poly_decode_signed_two_to_power_17(out, &pkt);

}

/*

 * @brief Encode a polynomial vector as an array of bytes.

 * Where the polynomial coefficients have a range of [0..15] or [0..43]

 * depending on the value of gamma2.

 *

 * See FIPS 204, Algorithm 28, w1Encode().

 *

 * @param w1 The vector to convert to bytes

 * @param gamma2 either ML_DSA_GAMMA2_Q_MINUS1_DIV32 or ML_DSA_GAMMA2_Q_MINUS1_DIV88

 * @returns 1 if the signature was encoded successfully or 0 otherwise.

 */

int ossl_ml_dsa_w1_encode(const VECTOR *w1, uint32_t gamma2,

                          uint8_t *out, size_t out_len)

{

    WPACKET pkt;

    ENCODE_FN *encode_fn;

    int ret = 0;

    size_t i;

    if (!WPACKET_init_static_len(&pkt, out, out_len, 0))

        return 0;

    if (gamma2 == ML_DSA_GAMMA2_Q_MINUS1_DIV32)

        encode_fn = poly_encode_4_bits;

    else

        encode_fn = poly_encode_6_bits;

    for (i = 0; i < w1->num_poly; ++i)

        if (!encode_fn(w1->poly + i, &pkt))

            goto err;

    ret = 1;

err:

    WPACKET_finish(&pkt);

    return ret;

}