Skip to content

Commit 6432215

Browse files
ElenaTyulenevaipl_ci
authored andcommitted
[SB][Library][Fix] Addressed Crypto Review comments for ML-KEM (#777)
1 parent 683cec4 commit 6432215

File tree

12 files changed

+220
-179
lines changed

12 files changed

+220
-179
lines changed

sources/include/ml_kem_internal/memory_consumption.h

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -44,11 +44,11 @@ __IPPCP_INLINE IppStatus mlkemMemoryConsumption(const IppsMLKEMState* pMLKEMCtx,
4444
sts = ippsHashGetSizeOptimal_rmf(&hashCtxSizeShake, ippsHashMethod_SHAKE256(3 * 256 * 8));
4545
IPP_BADARG_RET((sts != ippStsNoErr), sts);
4646
/* keyGen memory consumption */
47-
locKeygenBytes =
48-
eta1 * 64 + 3 * k * (int)sizeof(Ipp16sPoly) + hashCtxSizeShake + 5 * CP_ML_KEM_ALIGNMENT;
47+
locKeygenBytes = eta1 * 64 + 3 * k * (int)sizeof(Ipp16sPoly) + (int)sizeof(Ipp16sPoly) +
48+
hashCtxSizeShake + 5 * CP_ML_KEM_ALIGNMENT;
4949

5050
/* Encaps memory consumption */
51-
locEncapsBytes = 4 * (k * (int)sizeof(Ipp16sPoly)) + eta1 * 64 + 3 * (int)sizeof(Ipp16sPoly) +
51+
locEncapsBytes = 4 * (k * (int)sizeof(Ipp16sPoly)) + eta1 * 64 + 4 * (int)sizeof(Ipp16sPoly) +
5252
hashCtxSizeShake + 9 * CP_ML_KEM_ALIGNMENT;
5353

5454
/* Decaps memory consumption */
@@ -57,7 +57,7 @@ __IPPCP_INLINE IppStatus mlkemMemoryConsumption(const IppsMLKEMState* pMLKEMCtx,
5757
int hashCtxSizeSHA3_512 = 0;
5858
sts = ippsHashGetSizeOptimal_rmf(&hashCtxSizeSHA3_512, ippsHashMethod_SHA3_512());
5959
IPP_BADARG_RET((sts != ippStsNoErr), sts);
60-
locDecapsBytes = IPP_MAX(hashCtxSizeShake, hashCtxSizeSHA3_512) + 2 * (int)sizeof(Ipp16sPoly) +
60+
locDecapsBytes = IPP_MAX(hashCtxSizeShake, hashCtxSizeSHA3_512) + 3 * (int)sizeof(Ipp16sPoly) +
6161
2 * (k * (int)sizeof(Ipp16sPoly)) + 32 * (d_u * k + d_v) + locEncapsBytes +
6262
3 * CP_ML_KEM_ALIGNMENT;
6363

sources/include/ml_kem_internal/ml_kem.h

Lines changed: 53 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,12 @@ struct _cpMLKEMState {
5454
*/
5555
typedef enum { nttTransform, noNttTransform } nttTransformFlag;
5656

57+
/*
58+
* Stuff enumerator used in cp_matrixAGen() to conditionally generate
59+
* transposed matrix A
60+
*/
61+
typedef enum { matrixAOrigin, matrixATransposed } matrixAGenType;
62+
5763
/* State ID set\check helpers */
5864
#define CP_ML_KEM_SET_ID(pCtx) ((pCtx)->idCtx = (Ipp32u)idCtxMLKEM ^ (Ipp32u)IPP_UINT_PTR(pCtx))
5965
#define CP_ML_KEM_VALID_ID(pCtx) ((((pCtx)->idCtx) ^ (Ipp32u)IPP_UINT_PTR(pCtx)) == idCtxMLKEM)
@@ -85,14 +91,40 @@ typedef struct {
8591
// Stuff functions
8692
//-------------------------------//
8793

88-
#define CP_CHECK_FREE_RET(IN_RET_CONDITION, IN_STATUS, IN_P_STORAGE) \
89-
{ \
90-
if (IN_RET_CONDITION) { \
91-
cp_mlkemStorageReleaseAll(IN_P_STORAGE); \
92-
return IN_STATUS; \
93-
} \
94+
#define CP_CHECK_FREE_RET(IN_RET_CONDITION, IN_STATUS, IN_P_STORAGE) \
95+
{ \
96+
if (IN_RET_CONDITION) { \
97+
IppStatus releaseSts = cp_mlkemStorageReleaseAll(IN_P_STORAGE); \
98+
return (IN_STATUS) | releaseSts; \
99+
} \
94100
}
95101

102+
/* Memory allocation helpers for poly and polyvec */
103+
/* clang-format off */
104+
#define CP_ML_KEM_ALLOCATE_ALIGNED_POLYVEC(NAME, SIZE, STORAGE) \
105+
Ipp16sPoly*(NAME) = (Ipp16sPoly*)cp_mlkemStorageAllocate((STORAGE), \
106+
(SIZE) * sizeof(Ipp16sPoly) + CP_ML_KEM_ALIGNMENT); \
107+
CP_CHECK_FREE_RET((NAME) == NULL, ippStsMemAllocErr, (STORAGE)); \
108+
(NAME) = IPP_ALIGNED_PTR((NAME), CP_ML_KEM_ALIGNMENT);
109+
/* clang-format on */
110+
111+
#define CP_ML_KEM_ALLOCATE_ALIGNED_POLY(NAME, STORAGE) \
112+
CP_ML_KEM_ALLOCATE_ALIGNED_POLYVEC((NAME), 1, (STORAGE))
113+
114+
/* Memory release helpers for poly and polyvec */
115+
#define CP_ML_KEM_RELEASE_ALIGNED_POLYVEC(SIZE, STORAGE, STATUS) \
116+
(STATUS) |= \
117+
cp_mlkemStorageRelease((STORAGE), (SIZE) * sizeof(Ipp16sPoly) + CP_ML_KEM_ALIGNMENT);
118+
119+
#define CP_ML_KEM_RELEASE_ALIGNED_POLY(STORAGE, STATUS) \
120+
CP_ML_KEM_RELEASE_ALIGNED_POLYVEC(1, (STORAGE), (STATUS))
121+
122+
#ifdef __GNUC__
123+
#define ASM_VOLATILE(a) __asm__ __volatile__(a);
124+
#else
125+
#define ASM_VOLATILE(a)
126+
#endif
127+
96128
/*
97129
* Memory allocation primitive working with _cpMLKEMStorage structure.
98130
* Input: storage - pointer to _cpMLKEMStorage structure
@@ -127,7 +159,10 @@ __IPPCP_INLINE IppStatus cp_mlkemStorageRelease(_cpMLKEMStorage* storage, Ipp64s
127159
PurgeBlock(storage->pStorageData + storage->bytesUsed - bytesRelease, (int)bytesRelease);
128160

129161
storage->bytesUsed -= bytesRelease;
130-
return ippStsNoErr;
162+
ASM_VOLATILE("" ::: "memory")
163+
164+
// Check that the memory was released and zeroized as intended
165+
return (*(storage->pStorageData + storage->bytesUsed) == 0) ? ippStsNoErr : ippStsMemAllocErr;
131166
}
132167

133168
/*
@@ -136,18 +171,24 @@ __IPPCP_INLINE IppStatus cp_mlkemStorageRelease(_cpMLKEMStorage* storage, Ipp64s
136171
*
137172
* Note: the operation release the whole buffer, safe to be used for an empty storage.
138173
*/
139-
__IPPCP_INLINE void cp_mlkemStorageReleaseAll(_cpMLKEMStorage* storage)
174+
__IPPCP_INLINE IppStatus cp_mlkemStorageReleaseAll(_cpMLKEMStorage* storage)
140175
{
141176
// Zeroize the buffer (minimum amount between bytesUsed and bytesCapacity)
142177
PurgeBlock(storage->pStorageData,
143178
IPP_MIN((int)storage->bytesUsed, (int)storage->bytesCapacity));
144179
storage->bytesUsed = 0;
180+
ASM_VOLATILE("" ::: "memory")
181+
182+
return (*(storage->pStorageData + storage->bytesUsed) == 0) ? ippStsNoErr : ippStsMemAllocErr;
145183
}
146184

147185
//-------------------------------//
148186
// Kernel functions declaration
149187
//-------------------------------//
150188

189+
#define cp_mlkemBarrettReduce OWNAPI(cp_mlkemBarrettReduce)
190+
IPP_OWN_DECL(Ipp16s, cp_mlkemBarrettReduce, (Ipp32s x))
191+
151192
#define cp_polyAdd OWNAPI(cp_polyAdd)
152193
IPP_OWN_DECL(void, cp_polyAdd, (const Ipp16sPoly* f, const Ipp16sPoly* g, Ipp16sPoly* h))
153194
#define cp_polySub OWNAPI(cp_polySub)
@@ -169,7 +210,9 @@ IPP_OWN_DECL(IppStatus, cp_Decompress, (Ipp16u * out, const Ipp16s in, const Ipp
169210
#define cp_byteEncode OWNAPI(cp_byteEncode)
170211
IPP_OWN_DECL(IppStatus, cp_byteEncode, (Ipp8u * B, const Ipp16u d, const Ipp16sPoly* pPolyF))
171212
#define cp_byteDecode OWNAPI(cp_byteDecode)
172-
IPP_OWN_DECL(IppStatus, cp_byteDecode, (Ipp16sPoly * pPolyF, const Ipp16u d, const Ipp8u* B))
213+
IPP_OWN_DECL(IppStatus,
214+
cp_byteDecode,
215+
(Ipp16sPoly * pPolyF, const Ipp16u d, const Ipp8u* B, const int bByteSize))
173216

174217
#define cp_samplePolyCBD OWNAPI(cp_samplePolyCBD)
175218
IPP_OWN_DECL(IppStatus, cp_samplePolyCBD, (Ipp16sPoly * pPoly, const Ipp8u* pSeed, const Ipp8u eta))
@@ -180,7 +223,7 @@ IPP_OWN_DECL(IppStatus, cp_SampleNTT,
180223

181224
#define cp_matrixAGen OWNAPI(cp_matrixAGen)
182225
IPP_OWN_DECL(IppStatus, cp_matrixAGen,
183-
(Ipp16sPoly * matrixA, Ipp8u rho_j_i[34], IppsMLKEMState* mlkemCtx))
226+
(Ipp16sPoly * matrixA, Ipp8u rho_j_i[34], matrixAGenType matrixType, IppsMLKEMState* mlkemCtx))
184227
/* clang-format on */
185228

186229
#define cp_NTT OWNAPI(cp_NTT)

sources/ippcp/ml_kem/internal/internal_decaps.c

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -69,8 +69,7 @@ IPP_OWN_DEFN(IppStatus, cp_MLKEMdecaps_internal, (Ipp8u K[CP_SHARED_SECRET_BYTES
6969

7070
/* 5: m` <- K-PKE.Decrypt(dkPKE, c) */
7171
Ipp8u message[32];
72-
sts = cp_KPKE_Decrypt(message, pPKE_DecKey, ciphertext, mlkemCtx);
73-
CP_CHECK_FREE_RET(sts != ippStsNoErr, sts, pStorage);
72+
IppStatus decryptSts = cp_KPKE_Decrypt(message, pPKE_DecKey, ciphertext, mlkemCtx);
7473

7574
/* 6: (K`, r`) <- G(m`||h) */
7675

@@ -116,16 +115,12 @@ IPP_OWN_DEFN(IppStatus, cp_MLKEMdecaps_internal, (Ipp8u K[CP_SHARED_SECRET_BYTES
116115

117116
/* 9-10: if c != c` then K` <- K`` */
118117
BNU_CHUNK_T is_equal = cpIsEquBlock_ct(ciphertext, ciphertext1, ciphertext_size);
119-
if (is_equal) {
120-
CopyBlock(K1, K, CP_SHARED_SECRET_BYTES);
121-
} else {
122-
CopyBlock(K2, K, CP_SHARED_SECRET_BYTES);
123-
}
118+
MASKED_COPY_BNU(K, is_equal, K1, K2, CP_SHARED_SECRET_BYTES);
124119

125120
/* Release locally used storage */
126121
sts = cp_mlkemStorageRelease(pStorage, hash_size + CP_ML_KEM_ALIGNMENT); // hash_state
127122
sts |= cp_mlkemStorageRelease(pStorage,
128123
ciphertext_size + CP_ML_KEM_ALIGNMENT); // ciphertext1
129124

130-
return sts;
125+
return sts | decryptSts;
131126
}

sources/ippcp/ml_kem/internal/k_pke/k_pke_decrypt.c

Lines changed: 20 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -44,19 +44,21 @@ IPP_OWN_DEFN(IppStatus, cp_KPKE_Decrypt, (Ipp8u * message,
4444
const Ipp8u d_v = mlkemCtx->params.d_v;
4545
_cpMLKEMStorage* pStorage = &mlkemCtx->storage;
4646

47+
/* Allocate memory for temporary objects */
48+
CP_ML_KEM_ALLOCATE_ALIGNED_POLYVEC(u, k, pStorage)
49+
CP_ML_KEM_ALLOCATE_ALIGNED_POLY(v, pStorage)
50+
CP_ML_KEM_ALLOCATE_ALIGNED_POLYVEC(s, k, pStorage)
51+
CP_ML_KEM_ALLOCATE_ALIGNED_POLY(w, pStorage)
52+
4753
/* 1: c1 <- c[0 : 32*d_{u}*k] */
4854
Ipp8u* c1 = (Ipp8u*)ciphertext;
4955
/* 2: c2 <- c[32*d_{u}*k : 32(d_{u}*k + d_{v})] */
5056
Ipp8u* c2 = c1 + (32 * d_u * k);
5157

5258
/* 3: u` <- Decompress_{d_{u}}(cp_byteDecode_{d_{u}}(c1)) */
53-
Ipp16sPoly* u =
54-
(Ipp16sPoly*)cp_mlkemStorageAllocate(pStorage,
55-
k * sizeof(Ipp16sPoly) + CP_ML_KEM_ALIGNMENT);
56-
CP_CHECK_FREE_RET(u == NULL, ippStsMemAllocErr, pStorage);
57-
u = IPP_ALIGNED_PTR(u, CP_ML_KEM_ALIGNMENT);
5859
for (Ipp8u i = 0; i < k; i++) {
59-
cp_byteDecode(&u[i], d_u, c1 + i * 32 * d_u);
60+
sts = cp_byteDecode(&u[i], d_u, c1 + i * 32 * d_u, 32 * d_u * (k - i));
61+
IPP_BADARG_RET((sts != ippStsNoErr), sts);
6062

6163
for (Ipp32u j = 0; j < 256; j++) {
6264
sts = cp_Decompress((Ipp16u*)&u[i].values[j], u[i].values[j], d_u);
@@ -65,42 +67,29 @@ IPP_OWN_DEFN(IppStatus, cp_KPKE_Decrypt, (Ipp8u * message,
6567
}
6668

6769
/* 4: v` <- Decompress_{d_{v}}(cp_byteDecode_{d_{v}}(c2)) */
68-
Ipp16sPoly* v =
69-
(Ipp16sPoly*)cp_mlkemStorageAllocate(pStorage, sizeof(Ipp16sPoly) + CP_ML_KEM_ALIGNMENT);
70-
CP_CHECK_FREE_RET(v == NULL, ippStsMemAllocErr, pStorage);
71-
v = IPP_ALIGNED_PTR(v, CP_ML_KEM_ALIGNMENT);
72-
cp_byteDecode(v, d_v, c2);
70+
sts = cp_byteDecode(v, d_v, c2, 32 * d_v);
71+
IPP_BADARG_RET((sts != ippStsNoErr), sts);
7372
for (Ipp32u j = 0; j < 256; j++) {
74-
7573
sts = cp_Decompress((Ipp16u*)&v->values[j], v->values[j], d_v);
7674
IPP_BADARG_RET((sts != ippStsNoErr), sts);
7775
}
7876

7977
/* 5: s` <- cp_byteDecode_{12}(dk_{pke}) */
80-
Ipp16sPoly* s =
81-
(Ipp16sPoly*)cp_mlkemStorageAllocate(pStorage,
82-
k * sizeof(Ipp16sPoly) + CP_ML_KEM_ALIGNMENT);
83-
CP_CHECK_FREE_RET(s == NULL, ippStsMemAllocErr, pStorage);
84-
s = IPP_ALIGNED_PTR(s, CP_ML_KEM_ALIGNMENT);
85-
8678
for (Ipp8u i = 0; i < k; i++) {
87-
cp_byteDecode(&s[i], 12, pPKE_DecKey + 384 * i);
79+
sts = cp_byteDecode(&s[i], 12, pPKE_DecKey + 384 * i, 384 * (k - i));
80+
IPP_BADARG_RET((sts != ippStsNoErr), sts);
8881
}
8982

9083
/* 6: w <- v` - cp_NTT^{-1}(s`^{T} * cp_NTT(u`)) */
91-
Ipp16sPoly* w =
92-
(Ipp16sPoly*)cp_mlkemStorageAllocate(pStorage, sizeof(Ipp16sPoly) + CP_ML_KEM_ALIGNMENT);
93-
CP_CHECK_FREE_RET(w == NULL, ippStsMemAllocErr, pStorage);
94-
w = IPP_ALIGNED_PTR(w, CP_ML_KEM_ALIGNMENT);
95-
9684
cp_NTT(&u[0]);
9785
cp_multiplyNTT(&s[0], &u[0], w);
86+
CP_ML_KEM_ALLOCATE_ALIGNED_POLY(tmpPoly, pStorage)
9887
for (Ipp8u i = 1; i < k; i++) {
9988
cp_NTT(&u[i]);
100-
Ipp16sPoly tmpPoly;
101-
cp_multiplyNTT(&s[i], &u[i], &tmpPoly);
102-
cp_polyAdd(&tmpPoly, w, w);
89+
cp_multiplyNTT(&s[i], &u[i], tmpPoly);
90+
cp_polyAdd(tmpPoly, w, w);
10391
}
92+
CP_ML_KEM_RELEASE_ALIGNED_POLY(pStorage, sts) // Ipp16sPoly tmpPoly
10493
cp_inverseNTT(w);
10594
cp_polySub(v, w, w);
10695

@@ -113,16 +102,10 @@ IPP_OWN_DEFN(IppStatus, cp_KPKE_Decrypt, (Ipp8u * message,
113102
IPP_BADARG_RET((sts != ippStsNoErr), sts);
114103

115104
/* Release locally used storage */
116-
/* clang-format off */
117-
sts = cp_mlkemStorageRelease(pStorage, // Ipp16sPoly u[k]
118-
k * sizeof(Ipp16sPoly) + CP_ML_KEM_ALIGNMENT);
119-
sts |= cp_mlkemStorageRelease(pStorage, // Ipp16sPoly v
120-
sizeof(Ipp16sPoly) + CP_ML_KEM_ALIGNMENT);
121-
sts |= cp_mlkemStorageRelease(pStorage, // Ipp16sPoly s[k]
122-
k * sizeof(Ipp16sPoly) + CP_ML_KEM_ALIGNMENT);
123-
sts |= cp_mlkemStorageRelease(pStorage, // Ipp16sPoly w
124-
sizeof(Ipp16sPoly) + CP_ML_KEM_ALIGNMENT);
125-
/* clang-format on */
105+
CP_ML_KEM_RELEASE_ALIGNED_POLYVEC(k, pStorage, sts) // Ipp16sPoly u[k]
106+
CP_ML_KEM_RELEASE_ALIGNED_POLY(pStorage, sts) // Ipp16sPoly v
107+
CP_ML_KEM_RELEASE_ALIGNED_POLYVEC(k, pStorage, sts) // Ipp16sPoly s[k]
108+
CP_ML_KEM_RELEASE_ALIGNED_POLY(pStorage, sts) // Ipp16sPoly w
126109

127110
return sts;
128111
}

0 commit comments

Comments
 (0)