Skip to content

Commit

Permalink
Updates to ML-KEM Updated to FIPS203
Browse files Browse the repository at this point in the history 
Added internal functions isolating randomness
changed keygen gen matrix to not be transposed (p||i||j -> p||j||i)
added domain separation in genKeyPair
refactored input order to match specs
  • Loading branch information
royb committed Aug 13, 2024
1 parent 8a50bb7 commit 953ef6b
Show file tree
Hide file tree
Showing 4 changed files with 59 additions and 44 deletions.
78 changes: 46 additions & 32 deletions core/src/main/java/org/bouncycastle/pqc/crypto/crystals/kyber/KyberEngine.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -196,9 +196,10 @@ public void init(SecureRandom random)
this.random = random;
}

public byte[][] generateKemKeyPair()
//Internal functions are deterministic. No randomness is sampled inside them
public byte[][] generateKemKeyPairInternal(byte[] d, byte[] z)
{
byte[][] indCpaKeyPair = indCpa.generateKeyPair();
byte[][] indCpaKeyPair = indCpa.generateKeyPair(d);

byte[] s = new byte[KyberIndCpaSecretKeyBytes];

Expand All @@ -208,41 +209,19 @@ public byte[][] generateKemKeyPair()

symmetric.hash_h(hashedPublicKey, indCpaKeyPair[0], 0);

byte[] z = new byte[KyberSymBytes];
random.nextBytes(z);

byte[] outputPublicKey = new byte[KyberIndCpaPublicKeyBytes];
System.arraycopy(indCpaKeyPair[0], 0, outputPublicKey, 0, KyberIndCpaPublicKeyBytes);
return new byte[][]{ Arrays.copyOfRange(outputPublicKey, 0, outputPublicKey.length - 32), Arrays.copyOfRange(outputPublicKey, outputPublicKey.length - 32, outputPublicKey.length), s, hashedPublicKey, z };
}

public byte[][] kemEncrypt(byte[] publicKeyInput)
public byte[][] kemEncryptInternal(byte[] publicKeyInput, byte[] randBytes)
{
// Input validation (6.2 ML-KEM Encaps)
// Type Check
if (publicKeyInput.length != KyberIndCpaPublicKeyBytes)
{
throw new IllegalArgumentException("Input validation Error: Type check failed for ml-kem encapsulation");
}
// Modulus Check
PolyVec polyVec = new PolyVec(this);
byte[] seed = indCpa.unpackPublicKey(polyVec, publicKeyInput);
byte[] ek = indCpa.packPublicKey(polyVec, seed);
if (!Arrays.areEqual(ek, publicKeyInput))
{
throw new IllegalArgumentException("Input validation: Modulus check failed for ml-kem encapsulation");
}


byte[] outputCipherText;

byte[] buf = new byte[2 * KyberSymBytes];
byte[] kr = new byte[2 * KyberSymBytes];

byte[] randBytes = new byte[KyberSymBytes];

random.nextBytes(randBytes);

System.arraycopy(randBytes, 0, buf, 0, KyberSymBytes);

// SHA3-256 Public Key
Expand All @@ -252,33 +231,32 @@ public byte[][] kemEncrypt(byte[] publicKeyInput)
symmetric.hash_g(kr, buf);

// IndCpa Encryption
outputCipherText = indCpa.encrypt(Arrays.copyOfRange(buf, 0, KyberSymBytes), publicKeyInput, Arrays.copyOfRange(kr, 32, kr.length));
outputCipherText = indCpa.encrypt(publicKeyInput, Arrays.copyOfRange(buf, 0, KyberSymBytes), Arrays.copyOfRange(kr, 32, kr.length));

byte[] outputSharedSecret = new byte[sessionKeyLength];

System.arraycopy(kr, 0, outputSharedSecret, 0, outputSharedSecret.length);

byte[][] outBuf = new byte[2][];
outBuf[0] = outputSharedSecret;
outBuf[1] = outputCipherText;

return outBuf;
}

public byte[] kemDecrypt(byte[] cipherText, byte[] secretKey)
public byte[] kemDecryptInternal(byte[] secretKey, byte[] cipherText)
{
byte[] buf = new byte[2 * KyberSymBytes],
kr = new byte[2 * KyberSymBytes];
kr = new byte[2 * KyberSymBytes];

byte[] publicKey = Arrays.copyOfRange(secretKey, KyberIndCpaSecretKeyBytes, secretKey.length);

System.arraycopy(indCpa.decrypt(cipherText, secretKey), 0, buf, 0, KyberSymBytes);
System.arraycopy(indCpa.decrypt(secretKey, cipherText), 0, buf, 0, KyberSymBytes);

System.arraycopy(secretKey, KyberSecretKeyBytes - 2 * KyberSymBytes, buf, KyberSymBytes, KyberSymBytes);

symmetric.hash_g(kr, buf);

byte[] cmp = indCpa.encrypt(Arrays.copyOfRange(buf, 0, KyberSymBytes), publicKey, Arrays.copyOfRange(kr, KyberSymBytes, kr.length));
byte[] cmp = indCpa.encrypt(publicKey, Arrays.copyOfRange(buf, 0, KyberSymBytes), Arrays.copyOfRange(kr, KyberSymBytes, kr.length));

boolean fail = !(Arrays.constantTimeAreEqual(cipherText, cmp));

Expand All @@ -289,6 +267,42 @@ public byte[] kemDecrypt(byte[] cipherText, byte[] secretKey)
return Arrays.copyOfRange(kr, 0, sessionKeyLength);
}

public byte[][] generateKemKeyPair()
{
byte[] d = new byte[KyberSymBytes];
byte[] z = new byte[KyberSymBytes];
random.nextBytes(d);
random.nextBytes(z);

return generateKemKeyPairInternal(d, z);
}

public byte[][] kemEncrypt(byte[] publicKeyInput, byte[] randBytes)
{
//TODO: do input validation elsewhere?
// Input validation (6.2 ML-KEM Encaps)
// Type Check
if (publicKeyInput.length != KyberIndCpaPublicKeyBytes)
{
throw new IllegalArgumentException("Input validation Error: Type check failed for ml-kem encapsulation");
}
// Modulus Check
PolyVec polyVec = new PolyVec(this);
byte[] seed = indCpa.unpackPublicKey(polyVec, publicKeyInput);
byte[] ek = indCpa.packPublicKey(polyVec, seed);
if (!Arrays.areEqual(ek, publicKeyInput))
{
throw new IllegalArgumentException("Input validation: Modulus check failed for ml-kem encapsulation");
}

return kemEncryptInternal(publicKeyInput, randBytes);
}
public byte[] kemDecrypt(byte[] secretKey, byte[] cipherText)
{
//TODO: do input validation
return kemDecryptInternal(secretKey, cipherText);
}

private void cmov(byte[] r, byte[] x, int xlen, boolean b)
{
if (b)
Expand Down
17 changes: 7 additions & 10 deletions core/src/main/java/org/bouncycastle/pqc/crypto/crystals/kyber/KyberIndCpa.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import org.bouncycastle.crypto.digests.SHAKEDigest;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.Pack;

class KyberIndCpa
{
Expand Down Expand Up @@ -45,20 +46,16 @@ public KyberIndCpa(KyberEngine engine)
*
* @return KeyPair where each key is represented as bytes
*/
byte[][] generateKeyPair()
byte[][] generateKeyPair(byte[] d)
{
PolyVec secretKey = new PolyVec(engine),
publicKey = new PolyVec(engine),
e = new PolyVec(engine);

byte[] d = new byte[32];

// (p, sigma) <- G(d)

engine.getRandomBytes(d);
// (p, sigma) <- G(d || k)

byte[] buf = new byte[64];
symmetric.hash_g(buf, d);
symmetric.hash_g(buf, Arrays.concatenate(d, Pack.intToLittleEndian(kyberK)));

byte[] publicSeed = new byte[32]; // p in docs
byte[] noiseSeed = new byte[32]; // sigma in docs
Expand Down Expand Up @@ -141,7 +138,7 @@ byte[][] generateKeyPair()
return new byte[][]{packPublicKey(publicKey, publicSeed), packSecretKey(secretKey)};
}

public byte[] encrypt(byte[] msg, byte[] publicKeyInput, byte[] coins)
public byte[] encrypt(byte[] publicKeyInput, byte[] msg, byte[] coins)
{
int i;
byte[] seed;
Expand Down Expand Up @@ -180,7 +177,7 @@ public byte[] encrypt(byte[] msg, byte[] publicKeyInput, byte[] coins)
aMatrixTranspose[i] = new PolyVec(engine);
}

generateMatrix(aMatrixTranspose, seed, true);
generateMatrix(aMatrixTranspose, seed, false);

// System.out.print("matrix transposed = ");
// for (i = 0; i < kyberK; i++) {
Expand Down Expand Up @@ -383,7 +380,7 @@ private static int rejectionSampling(Poly outputBuffer, int coeffOff, int len, b

}

public byte[] decrypt(byte[] cipherText, byte[] secretKey)
public byte[] decrypt(byte[] secretKey, byte[] cipherText)
{
int i;
byte[] outputMessage = new byte[KyberEngine.getKyberIndCpaMsgBytes()];
Expand Down
2 changes: 1 addition & 1 deletion core/src/main/java/org/bouncycastle/pqc/crypto/crystals/kyber/KyberKEMExtractor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ private void initCipher(AsymmetricKeyParameter recipientKey)
public byte[] extractSecret(byte[] encapsulation)
{
// Decryption
byte[] sharedSecret = engine.kemDecrypt(encapsulation, key.getEncoded());
byte[] sharedSecret = engine.kemDecrypt(key.getEncoded(), encapsulation);
return sharedSecret;
}

Expand Down
6 changes: 5 additions & 1 deletion core/src/main/java/org/bouncycastle/pqc/crypto/crystals/kyber/KyberKEMGenerator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,11 @@ public SecretWithEncapsulation generateEncapsulated(AsymmetricKeyParameter recip
KyberPublicKeyParameters key = (KyberPublicKeyParameters)recipientKey;
KyberEngine engine = key.getParameters().getEngine();
engine.init(sr);
byte[][] kemEncrypt = engine.kemEncrypt(key.getEncoded());

byte[] randBytes = new byte[32];
engine.getRandomBytes(randBytes);

byte[][] kemEncrypt = engine.kemEncrypt(key.getEncoded(), randBytes);
return new SecretWithEncapsulationImpl(kemEncrypt[0], kemEncrypt[1]);
}
}

0 comments on commit 953ef6b

Please sign in to comment.