CVE-2026-5598

@@ -3,7 +3,9 @@
 import java.security.SecureRandom;
 
 import org.bouncycastle.crypto.Xof;
+import org.bouncycastle.math.raw.Nat;
 import org.bouncycastle.util.Arrays;
+import org.bouncycastle.util.Bytes;
 import org.bouncycastle.util.Pack;
 
 class FrodoEngine
@@ -110,14 +112,17 @@ private short[] sample_matrix(short[] r, int offset, int n1, int n2)
     private short[] matrix_transpose(short[] X, int n1, int n2)
     {
         short[] res = new short[n1 * n2];
-
         for (int i = 0; i < n2; i++)
+        {
             for (int j = 0; j < n1; j++)
-                res[i*n1 +j] = X[j*n2+ i];
+            {
+                res[i * n1 + j] = X[j * n2 + i];
+            }
+        }
         return res;
     }
 
-    private short[] matrix_mul(short[] X, int Xrow, int Xcol, short[] Y, int Yrow, int Ycol)
+    private short[] matrix_mul(short[] X, int Xrow, int Xcol, short[] Y, int Ycol)
     {
         int qMask = q - 1;
         short[] res = new short[Xrow * Ycol];
@@ -139,11 +144,14 @@ private short[] matrix_mul(short[] X, int Xrow, int Xcol, short[] Y, int Yrow, i
     private short[] matrix_add(short[] X, short[] Y, int n1, int m1)
     {
         int qMask = q - 1;
-        short[] res = new short[n1*m1];
+        short[] res = new short[n1 * m1];
         for (int i = 0; i < n1; i++)
+        {
             for (int j = 0; j < m1; j++)
-                res[i*m1+j] = (short)((X[i*m1+j] + Y[i*m1+j]) & qMask);
-
+            {
+                res[i * m1 + j] = (short)((X[i * m1 + j] + Y[i * m1 + j]) & qMask);
+            }
+        }
         return res;
     }
 
@@ -159,33 +167,31 @@ private byte[] pack(short[] C)
 
         while (i < out.length && (j < n || ((j == n) && (bits > 0))))
         {
-
             byte b = 0;  // bits in out[i] already filled in
             while (b < 8)
             {
                 int nbits = Math.min(8 - b, bits);
-                short mask = (short) ((1 << nbits) - 1);
-                byte t = (byte) ((w >> (bits - nbits)) & mask);  // the bits to copy from w to out
-                out[i] = (byte) (out[i] + (t << (8 - b - nbits)));
+                short mask = (short)((1 << nbits) - 1);
+                byte t = (byte)((w >> (bits - nbits)) & mask); // the bits to copy from w to out
+                out[i] = (byte)(out[i] + (t << (8 - b - nbits)));
                 b += nbits;
                 bits -= nbits;
 
                 if (bits == 0)
                 {
-                    if (j < n)
-                    {
-                        w = C[j];
-                        bits = (byte) D;
-                        j++;
-                    }
-                    else
+                    if (j >= n)
                     {
-                        break;  // the input vector is exhausted
+                        break; // the input vector is exhausted
                     }
+
+                    w = C[j];
+                    bits = (byte)D;
+                    j++;
                 }
             }
             if (b == 8)
-            {  // out[i] is filled in
+            {
+                // out[i] is filled in
                 i++;
             }
         }
@@ -202,24 +208,21 @@ public void kem_keypair(byte[] pk, byte[] sk, SecureRandom random)
         byte[] seedSE = Arrays.copyOfRange(s_seedSE_z, len_s_bytes, len_s_bytes + len_seedSE_bytes);
         byte[] z = Arrays.copyOfRange(s_seedSE_z, len_s_bytes + len_seedSE_bytes, len_s_bytes + len_seedSE_bytes + len_z_bytes);
 
-        // 2. Generate pseudorandom seed seedA = SHAKE(z, len_seedA) (length in bits)
+        // 2. Generate pseudo-random seed seedA = SHAKE(z, len_seedA) (length in bits)
         byte[] seedA = new byte[len_seedA_bytes];
         digest.update(z, 0, z.length);
         digest.doFinal(seedA, 0, seedA.length);
 
         // 3. A = Frodo.Gen(seedA)
-        short[] A = gen.genMatrix(seedA);
+        short[] A = gen.genMatrix(seedA, 0, seedA.length);
 
         // 4. r = SHAKE(0x5F || seedSE, 2*n*nbar*len_chi) (length in bits), parsed as 2*n*nbar len_chi-bit integers in little-endian byte order
         byte[] rbytes = new byte[2 * n * nbar * len_chi_bytes];
-
         digest.update((byte)0x5f);
         digest.update(seedSE, 0, seedSE.length);
         digest.doFinal(rbytes, 0, rbytes.length);
 
-        short[] r = new short[2 * n * nbar];
-        for (int i = 0; i < r.length; i++)
-            r[i] = Pack.littleEndianToShort(rbytes, i * 2);
+        short[] r = Pack.littleEndianToShort(rbytes, 0, rbytes.length / 2);
 
         // 5. S^T = Frodo.SampleMatrix(r[0 .. n*nbar-1], nbar, n)
         short[] S_T = sample_matrix(r, 0, nbar, n);
@@ -229,28 +232,24 @@ public void kem_keypair(byte[] pk, byte[] sk, SecureRandom random)
         short[] E = sample_matrix(r, n * nbar, n, nbar);
 
         // 7. B = A * S + E
-        short[] B = matrix_add(matrix_mul(A, n, n, S, n, nbar), E, n, nbar);
+        short[] B = matrix_add(matrix_mul(A, n, n, S, nbar), E, n, nbar);
 
         // 8. b = Pack(B)
         byte[] b = pack(B);
 
         // 9. pkh = SHAKE(seedA || b, len_pkh) (length in bits)
         // 10. pk = seedA || b
-        System.arraycopy(Arrays.concatenate(seedA, b), 0, pk, 0, len_pk_bytes);
+        System.arraycopy(seedA, 0, pk, 0, len_seedA_bytes);
+        System.arraycopy(b, 0, pk, len_seedA_bytes, len_pk_bytes - len_seedA_bytes);
 
         byte[] pkh = new byte[len_pkh_bytes];
         digest.update(pk, 0, pk.length);
         digest.doFinal(pkh, 0, pkh.length);
 
         //10. sk = (s || seedA || b, S^T, pkh)
-        System.arraycopy(Arrays.concatenate(s, pk), 0,
-                sk, 0, len_s_bytes + len_pk_bytes);
-
-        for (int i = 0; i < nbar; i++)
-            for (int j = 0; j < n; j++)
-                System.arraycopy(Pack.shortToLittleEndian(S_T[i*n+j]), 0,
-                        sk, len_s_bytes + len_pk_bytes + i * n * 2 + j * 2, 2);
-
+        System.arraycopy(s, 0, sk, 0, len_s_bytes);
+        System.arraycopy(pk, 0, sk, len_s_bytes, len_pk_bytes);
+        Pack.shortToLittleEndian(S_T, sk, len_s_bytes + len_pk_bytes);
         System.arraycopy(pkh, 0, sk, len_sk_bytes - len_pkh_bytes, len_pkh_bytes);
     }
 
@@ -269,29 +268,28 @@ private short[] unpack(byte[] in, int n1, int n2)
             while (b < D)
             {
                 int nbits = Math.min(D - b, bits);
-                short mask = (short) (((1 << nbits) - 1) & 0xffff);
-                byte t = (byte) ((((w & 0xff) >>> ((bits & 0xff) - nbits)) & (mask & 0xffff)) & 0xff);  // the bits to copy from w to out
-                out[i] = (short) ((out[i] & 0xffff) + (((t & 0xff) << (D - (b & 0xff) - nbits))) & 0xffff);
+                short mask = (short)(((1 << nbits) - 1) & 0xffff);
+                byte t = (byte)((((w & 0xff) >>> ((bits & 0xff) - nbits)) & (mask & 0xffff)) & 0xff); // the bits to copy from w to out
+                out[i] = (short)((out[i] & 0xffff) + (((t & 0xff) << (D - (b & 0xff) - nbits))) & 0xffff);
                 b += nbits;
                 bits -= nbits;
                 w &= ~(mask << bits);
 
                 if (bits == 0)
                 {
-                    if (j < in.length)
-                    {
-                        w = in[j];
-                        bits = 8;
-                        j++;
-                    }
-                    else
+                    if (j >= in.length)
                     {
-                        break;  // the input vector is exhausted
+                        break; // the input vector is exhausted
                     }
+
+                    w = in[j];
+                    bits = 8;
+                    j++;
                 }
             }
             if (b == D)
-            {  // out[i] is filled in
+            {
+                // out[i] is filled in
                 i++;
             }
         }
@@ -300,19 +298,19 @@ private short[] unpack(byte[] in, int n1, int n2)
 
     private short[] encode(byte[] k)
     {
-        int l, byte_index = 0;
+        int byte_index = 0;
         int bit = 0;
         short[] K = new short[mbar*nbar];
-        int temp;
+
         // 1. for i = 0; i < mbar; i += 1
         for (int i = 0; i < mbar; i++)
         {
             // 2. for j = 0; j < nbar; j += 1
             for (int j = 0; j < nbar; j++)
             {
                 // 3. tmp = sum_{l=0}^{B-1} k_{(i*nbar+j)*B+l} 2^l
-                temp = 0;
-                for (l = 0; l < B; l++)
+                int temp = 0;
+                for (int l = 0; l < B; l++)
                 {
                     temp += ((k[byte_index] >>> bit) & 1) << l;
 
@@ -322,7 +320,7 @@ private short[] encode(byte[] k)
                 }
 
                 // 4. K[i][j] = ec(tmp) = tmp * q/2^B
-                K[i*nbar+j] = (short) (temp * (q / (1 << B)));
+                K[i * nbar + j] = (short)(temp * (q / (1 << B)));
             }
         }
         return K;
@@ -331,7 +329,7 @@ private short[] encode(byte[] k)
     public void kem_enc(byte[] ct, byte[] ss, byte[] pk, SecureRandom random)
     {
         // Parse pk = seedA || b
-        byte[] seedA = Arrays.copyOfRange(pk, 0, len_seedA_bytes);
+//        byte[] seedA = Arrays.copyOfRange(pk, 0, len_seedA_bytes);
         byte[] b = Arrays.copyOfRange(pk, len_seedA_bytes, len_pk_bytes);
 
         // 1. Choose a uniformly random key mu in {0,1}^len_mu (length in bits)
@@ -358,9 +356,7 @@ public void kem_enc(byte[] ct, byte[] ss, byte[] pk, SecureRandom random)
         digest.update(seedSE, 0, seedSE.length);
         digest.doFinal(rbytes, 0, rbytes.length);
 
-        short[] r = new short[rbytes.length / 2];
-        for (int i = 0; i < r.length; i++)
-            r[i] = Pack.littleEndianToShort(rbytes, i * 2);
+        short[] r = Pack.littleEndianToShort(rbytes, 0, rbytes.length / 2);
 
         // 5. S' = Frodo.SampleMatrix(r[0 .. mbar*n-1], mbar, n)
         short[] Sprime = sample_matrix(r, 0, mbar, n);
@@ -369,10 +365,10 @@ public void kem_enc(byte[] ct, byte[] ss, byte[] pk, SecureRandom random)
         short[] Eprime = sample_matrix(r, mbar * n, mbar, n);
 
         // 7. A = Frodo.Gen(seedA)
-        short[] A = gen.genMatrix(seedA);
+        short[] A = gen.genMatrix(pk, 0, len_seedA_bytes);
 
         // 8. B' = S' A + E'
-        short[] Bprime = matrix_add(matrix_mul(Sprime, mbar, n, A, n, n), Eprime, mbar, n);
+        short[] Bprime = matrix_add(matrix_mul(Sprime, mbar, n, A, n), Eprime, mbar, n);
 
         // 9. c1 = Frodo.Pack(B')
         byte[] c1 = pack(Bprime);
@@ -383,9 +379,8 @@ public void kem_enc(byte[] ct, byte[] ss, byte[] pk, SecureRandom random)
         // 11. B = Frodo.Unpack(b, n, nbar)
         short[] B = unpack(b, n, nbar);
 
-
         // 12. V = S' B + E''
-        short[] V = matrix_add(matrix_mul(Sprime, mbar, n, B, n, nbar), Eprimeprime, mbar, nbar);
+        short[] V = matrix_add(matrix_mul(Sprime, mbar, n, B, nbar), Eprimeprime, mbar, nbar);
 
         // 13. C = V + Frodo.Encode(mu)
         short[] EncodedMU = encode(mu);
@@ -395,80 +390,55 @@ public void kem_enc(byte[] ct, byte[] ss, byte[] pk, SecureRandom random)
         byte[] c2 = pack(C);
 
         // 15. ss = SHAKE(c1 || c2 || k, len_ss)
-        // ct = c1 + c2
-        System.arraycopy(Arrays.concatenate(c1, c2), 0, ct, 0, len_ct_bytes);
-        digest.update(c1, 0, c1.length);
-        digest.update(c2, 0, c2.length);
+        // ct = c1 || c2
+        System.arraycopy(c1, 0, ct, 0, c1.length);
+        System.arraycopy(c2, 0, ct, c1.length, len_ct_bytes - c1.length);
+
+        digest.update(ct, 0, len_ct_bytes);
         digest.update(k, 0, len_k_bytes);
         digest.doFinal(ss, 0, len_s_bytes);
     }
 
     private short[] matrix_sub(short[] X, short[] Y, int n1, int n2)
     {
         int qMask = q - 1;
-        short[] res = new short[n1*n2];
+        short[] res = new short[n1 * n2];
         for (int i = 0; i < n1; i++)
+        {
             for (int j = 0; j < n2; j++)
-                res[i*n2+j] = (short)((X[i*n2+j] - Y[i*n2+j]) & qMask);
-
+            {
+                res[i * n2 + j] = (short)((X[i * n2 + j] - Y[i * n2 + j]) & qMask);
+            }
+        }
         return res;
     }
 
     private byte[] decode(short[] in)
     {
-        int i, j, index = 0, npieces_word = 8;
+        int index = 0, npieces_word = 8;
         int nwords = (nbar * nbar) / 8;
-        short temp;
-        short maskex = (short) ((1 << B) - 1);
-        short maskq = (short) ((1 << D) - 1);
+        short maskex = (short)((1 << B) - 1);
+        short maskq = (short)((1 << D) - 1);
         byte[] out = new byte[npieces_word * B];
-        long templong;
 
-        for (i = 0; i < nwords; i++)
+        for (int i = 0; i < nwords; i++)
         {
-            templong = 0;
-            for (j = 0; j < npieces_word; j++)
-            {  // temp = floor(in*2^{-11}+0.5)
-                temp = (short) (((in[index] & maskq) + (1 << (D - B - 1))) >> (D - B));
-                templong |= ((long) (temp & maskex)) << (B * j);
+            long templong = 0;
+            for (int j = 0; j < npieces_word; j++)
+            {
+                // temp = floor(in*2^{-11}+0.5)
+                short temp = (short)(((in[index] & maskq) + (1 << (D - B - 1))) >> (D - B));
+                templong |= ((long)(temp & maskex)) << (B * j);
                 index++;
             }
-            for (j = 0; j < B; j++)
-                out[i * B + j] = (byte) ((templong >> (8 * j)) & 0xFF);
+            for (int j = 0; j < B; j++)
+            {
+                out[i * B + j] = (byte)((templong >> (8 * j)) & 0xFF);
+            }
         }
         return out;
     }
 
-
-    private short ctverify(short[] a1, short[] a2, short[] b1, short[] b2)
-    {
-        // Compare two arrays in constant time.
-        // Returns 0 if the byte arrays are equal, -1 otherwise.
-        short r = 0;
-
-        for (short i = 0; i < a1.length; i++)
-            r |= a1[i] ^ b1[i];
-
-        for (short i = 0; i < a2.length; i++)
-            r |= a2[i] ^ b2[i];
-
-//        r = (short) ((-(short)(r >> 1) | -(short)(r & 1)) >> (8*2-1));
-        if (r == 0)
-            return 0;
-        return -1;
-    }
-
-    private byte[] ctselect(byte[] a, byte[] b, short selector)
-    {
-        // Select one of the two input arrays to be moved to r
-        // If (selector == 0) then load r with a, else if (selector == -1) load r with b
-        byte[] r = new byte[a.length];
-        for (int i = 0; i < a.length; i++)
-            r[i] = (byte) (((~selector & a[i]) & 0xff) | ((selector & b[i]) & 0xff));
-
-        return r;
-    }
-
     public void kem_dec(byte[] ss, byte[] ct, byte[] sk)
     {
         // Parse ct = c1 || c2
@@ -481,15 +451,10 @@ public void kem_dec(byte[] ss, byte[] ct, byte[] sk)
         byte[] c2 = Arrays.copyOfRange(ct, offset, offset + length);
 
         // Parse sk = (s || seedA || b, S^T, pkh)
-        offset = 0;
-        length = len_s_bytes;
-        byte[] s = Arrays.copyOfRange(sk, offset, offset + length);
+//        byte[] s = Arrays.copyOfRange(sk, 0, len_s_bytes);
+//        byte[] seedA = Arrays.copyOfRange(sk, len_s_bytes, len_s_bytes + len_seedA_bytes);
 
-        offset += length;
-        length = len_seedA_bytes;
-        byte[] seedA = Arrays.copyOfRange(sk, offset, offset + length);
-
-        offset += length;
+        offset = len_s_bytes + len_seedA_bytes;
         length = (D * n * nbar) / 8;
         byte[] b = Arrays.copyOfRange(sk, offset, offset + length);
 
@@ -500,8 +465,12 @@ public void kem_dec(byte[] ss, byte[] ct, byte[] sk)
         short[] Stransposed = new short[nbar * n];
 
         for (int i = 0; i < nbar; i++)
+        {
             for (int j = 0; j < n; j++)
+            {
                 Stransposed[i*n+j] = Pack.littleEndianToShort(Sbytes, i * n * 2 + j * 2);
+            }
+        }
 
         short[] S = matrix_transpose(Stransposed, nbar, n);
 
@@ -516,7 +485,7 @@ public void kem_dec(byte[] ss, byte[] ct, byte[] sk)
         short[] C = unpack(c2, mbar, nbar);
 
         // 3. M = C - B' S
-        short[] BprimeS = matrix_mul(Bprime, mbar, n, S, n, nbar);
+        short[] BprimeS = matrix_mul(Bprime, mbar, n, S, nbar);
         short[] M = matrix_sub(C, BprimeS, mbar, nbar);
 
         // 4. mu' = Frodo.Decode(M)
@@ -530,19 +499,15 @@ public void kem_dec(byte[] ss, byte[] ct, byte[] sk)
         digest.update(muprime, 0, len_mu_bytes);
         digest.doFinal(seedSEprime_kprime, 0, len_seedSE_bytes + len_k_bytes);
 
-        byte[] kprime = Arrays.copyOfRange(seedSEprime_kprime, len_seedSE_bytes, len_seedSE_bytes + len_k_bytes);
+        byte[] K = Arrays.copyOfRange(seedSEprime_kprime, len_seedSE_bytes, len_seedSE_bytes + len_k_bytes);
 
         // 7. r = SHAKE(0x96 || seedSE', 2*mbar*n + mbar*nbar*len_chi) (length in bits)
         byte[] rbytes = new byte[(2 * mbar * n + mbar * mbar) * len_chi_bytes];
         digest.update((byte)0x96);
         digest.update(seedSEprime_kprime, 0, len_seedSE_bytes);
         digest.doFinal(rbytes, 0, rbytes.length);
 
-        short[] r = new short[2 * mbar * n + mbar * nbar];
-        for (int i = 0; i < r.length; i++)
-        {
-            r[i] = Pack.littleEndianToShort(rbytes, i * 2);
-        }
+        short[] r = Pack.littleEndianToShort(rbytes, 0, rbytes.length / 2);
 
         // 8. S' = Frodo.SampleMatrix(r[0 .. mbar*n-1], mbar, n)
         short[] Sprime = sample_matrix(r, 0, mbar, n);
@@ -551,10 +516,10 @@ public void kem_dec(byte[] ss, byte[] ct, byte[] sk)
         short[] Eprime = sample_matrix(r, mbar * n, mbar, n);
 
         // 10. A = Frodo.Gen(seedA)
-        short[] A = gen.genMatrix(seedA);
+        short[] A = gen.genMatrix(sk, len_s_bytes, len_seedA_bytes);
 
         // 11. B'' = S' A + E'
-        short[] Bprimeprime = matrix_add(matrix_mul(Sprime, mbar, n, A, n, n), Eprime, mbar, n);
+        short[] Bprimeprime = matrix_add(matrix_mul(Sprime, mbar, n, A, n), Eprime, mbar, n);
 
         // 12. E'' = Frodo.SampleMatrix(r[2*mbar*n .. 2*mbar*n + mbar*nbar-1], mbar, n)
         short[] Eprimeprime = sample_matrix(r, 2 * mbar * n, mbar, nbar);
@@ -563,7 +528,7 @@ public void kem_dec(byte[] ss, byte[] ct, byte[] sk)
         short[] B = unpack(b, n, nbar);
 
         // 14. V = S' B + E''
-        short[] V = matrix_add(matrix_mul(Sprime, mbar, n, B, n, nbar), Eprimeprime, mbar, nbar);
+        short[] V = matrix_add(matrix_mul(Sprime, mbar, n, B, nbar), Eprimeprime, mbar, nbar);
 
         // 15. C' = V + Frodo.Encode(muprime)
         short[] Cprime = matrix_add(V, encode(muprime), mbar, nbar);
@@ -573,14 +538,27 @@ public void kem_dec(byte[] ss, byte[] ct, byte[] sk)
         // Qian Guo, Thomas Johansson, Alexander Nilsson. A key-recovery timing attack on post-quantum
         // primitives using the Fujisaki-Okamoto transformation and its application on FrodoKEM. In CRYPTO 2020.
         //TODO change it so Bprime and C are in the same array same with B'' and C'
-        short use_kprime = ctverify(Bprime, C, Bprimeprime, Cprime);
-        byte[] kbar = ctselect(kprime, s, use_kprime);
+        int use_kprime = ctverify(Bprime, C, Bprimeprime, Cprime);
+        Bytes.cmov(K.length, ~use_kprime, sk, K);
 
         // 17. ss = SHAKE(c1 || c2 || kbar, len_ss) (length in bits)
         digest.update(c1, 0, c1.length);
         digest.update(c2, 0, c2.length);
-        digest.update(kbar, 0, kbar.length);
+        digest.update(K, 0, K.length);
         digest.doFinal(ss, 0, len_ss_bytes);
     }
 
-}
\ No newline at end of file
+    private static int ctverify(short[] a1, short[] a2, short[] b1, short[] b2)
+    {
+        int r = 0;
+        for (int i = 0; i < a1.length; i++)
+        {
+            r |= a1[i] ^ b1[i];
+        }
+        for (int i = 0; i < a2.length; i++)
+        {
+            r |= a2[i] ^ b2[i];
+        }
+        return Nat.czero(r);
+    }
+}

@@ -9,87 +9,85 @@
 
 abstract class FrodoMatrixGenerator
 {
-    int n;
-    int q;
+    final int n;
+    final int q;
 
-    public FrodoMatrixGenerator(int n, int q)
+    FrodoMatrixGenerator(int n, int q)
     {
         this.n = n;
         this.q = q;
     }
 
-    abstract short[] genMatrix(byte[] seedA);
+    abstract short[] genMatrix(byte[] seed, int seedOff, int seedLen);
 
-    static class Shake128MatrixGenerator
-            extends FrodoMatrixGenerator
+    static class Shake128MatrixGenerator extends FrodoMatrixGenerator
     {
         public Shake128MatrixGenerator(int n, int q)
         {
             super(n, q);
         }
 
-        short[] genMatrix(byte[] seedA)
+        short[] genMatrix(byte[] seed, int seedOff, int seedLen)
         {
-            short[] A = new short[n*n];
-            short i, j;
+            short[] A = new short[n * n];
             byte[] tmp = new byte[(16 * n) / 8];
-            byte[] b = new byte[2 + seedA.length];
-            System.arraycopy(seedA, 0, b, 2, seedA.length);
+            byte[] b = new byte[2 + seedLen];
+            System.arraycopy(seed, seedOff, b, 2, seedLen);
 
-            Xof digest = new SHAKEDigest(128);
+            SHAKEDigest digest = new SHAKEDigest(128);
 
-            for (i = 0; i < n; i++)
+            for (int i = 0; i < n; i++)
             {
-                // 1. b = i || seedA in {0,1}^{16 + len_seedA}, where i is encoded as a 16-bit integer in little-endian byte order
-                Pack.shortToLittleEndian(i, b, 0);
+                // 1. b = i || seedA in {0,1}^{16 + len_seedA}, where i is encoded as 16-bit LE
+                Pack.shortToLittleEndian((short)i, b, 0);
 
-                // 2. c_{i,0} || c_{i,1} || ... || c_{i,n-1} = SHAKE128(b, 16n) (length in bits) where each c_{i,j} is parsed as a 16-bit integer in little-endian byte order format
+                // 2. c_{i,0} || c_{i,1} || ... || c_{i,n-1} = SHAKE128(b, 16n) (length in bits) where each c_{i,j}
+                // is parsed as 16-bit LE
                 digest.update(b, 0, b.length);
                 digest.doFinal(tmp, 0, tmp.length);
-                for (j = 0; j < n; j++)
+
+                for (int j = 0; j < n; j++)
                 {
-                    A[i*n+j] = (short) (Pack.littleEndianToShort(tmp, 2 * j) & (q - 1));
+                    A[i * n + j] = (short)(Pack.littleEndianToShort(tmp, 2 * j) & (q - 1));
                 }
             }
             return A;
         }
     }
 
-    static class Aes128MatrixGenerator
-            extends FrodoMatrixGenerator
+    static class Aes128MatrixGenerator extends FrodoMatrixGenerator
     {
         public Aes128MatrixGenerator(int n, int q)
         {
             super(n, q);
         }
 
-        short[] genMatrix(byte[] seedA)
+        short[] genMatrix(byte[] seed, int seedOff, int seedLen)
         {
             // """Generate matrix A using AES-128 (FrodoKEM specification, Algorithm 7)"""
             // A = [[None for j in range(self.n)] for i in range(self.n)]
-            short[] A = new short[n*n];
+            short[] A = new short[n * n];
             byte[] b = new byte[16];
             byte[] c = new byte[16];
 
-            BlockCipher cipher = new AESEngine();
-            cipher.init(true, new KeyParameter(seedA));
+            BlockCipher cipher = AESEngine.newInstance();
+            cipher.init(true, new KeyParameter(seed, seedOff, seedLen));
 
-            // 1. for i = 0; i < n; i += 1
             for (int i = 0; i < n; i++)
             {
                 Pack.shortToLittleEndian((short)i, b, 0);
-                // 2. for j = 0; j < n; j += 8
-                for (int j = 0; j < n; j+=8)
+
+                for (int j = 0; j < n; j += 8)
                 {
-                    // 3. b = i || j || 0 || ... || 0 in {0,1}^128, where i and j are encoded as 16-bit integers in little-endian byte order
+                    // 3. b = i || j || 0 || ... || 0 in {0,1}^128, where i and j are encoded as 16-bit LE
                     Pack.shortToLittleEndian((short)j, b, 2);
                     // 4. c = AES128(seedA, b)
                     cipher.processBlock(b, 0, c, 0);
-                    // 5. for k = 0; k < 8; k += 1
+
                     for (int k = 0; k < 8; k++)
                     {
-                        // 6. A[i][j+k] = c[k] where c is treated as a sequence of 8 16-bit integers each in little-endian byte order
-                        A[i*n+ j + k] = (short) (Pack.littleEndianToShort(c, 2 * k) & (q - 1));
+                        // 6. A[i][j+k] = c[k] where c is treated as a sequence of 8 16-bit LE
+                        A[i * n + j + k] = (short)(Pack.littleEndianToShort(c, 2 * k) & (q - 1));
                     }
                 }
             }

@@ -100,31 +100,10 @@ public FrodoEngine(int n, int D, int B, short[] cdf_table, Xof digest, FrodoMatr
         this.gen = mGen;
     }
 
-    private short sample(short r)
-    {
-        short t, e;
-        // 1. t = sum_{i=1}^{len_x - 1} r_i * 2^{i-1}
-        t = (short) ((r & 0xffff) >>> 1);
-        e = 0; // 2. e = 0
-        for (int z = 0; z < T_chi.length; z++)
-        {
-            if (t > T_chi[z]) // 4. if t > T_chi(z)
-                e++; // 5. e = e + 1
-        }
-        // 6. e = (-1)^{r_0} * e
-
-        if (((r & 0xffff) % 2) == 1)
-            e = (short) ((e) * (-1) & 0xffff);
-
-        return e;
-    }
-
     private short[] sample_matrix(short[] r, int offset, int n1, int n2)
     {
         short[] E = new short[n1 * n2];
-        for (int i = 0; i < n1; i++)
-            for (int j = 0; j < n2; j++)
-                E[i*n2+j] = sample(r[i * n2 + j + offset]);
+        Noise.sample(T_chi, r, offset, E);
         return E;
     }
 

@@ -0,0 +1,30 @@
+package org.bouncycastle.pqc.crypto.frodo;
+
+abstract class Noise
+{
+    static void sample(short[] cdf, short[] r, int rOff, short[] s) 
+    {
+        // No need to compare with the last value.
+//        assert cdf[cdf.length - 1] == 0x7FFF;
+
+        // Fills 's' with samples from the noise distribution 'cdf' using pseudo-random values 'r[rOff..]' 
+        for (int i = 0, n = s.length; i < n; ++i)
+        {
+            int sample = 0;
+            int r_i = r[rOff + i] & 0xFFFF;
+            int prnd = r_i >>> 1;   // Drop the least significant bit
+            int sign = r_i & 1;     // Pick the least significant bit
+
+            for (int j = 0; j < cdf.length - 1; ++j)
+            {
+                // Constant time comparison: 1 if cdf[j] < prnd, 0 otherwise.
+                sample += (cdf[j] - prnd) >>> 31;
+            }
+
+            // Assuming that sign is either 0 or 1, flips sample iff sign = 1
+            sample = ((-sign) ^ sample) + sign;
+
+            s[i] = (short)sample;
+        }
+    }    
+}