Add a concrete DHKEM implementation

Author: jchambers

src/main/java/com/eatthepath/noise/component/DhkemKeyEncapsulationMechanism.java

@@ -0,0 +1,81 @@
+package com.eatthepath.noise.component;
+
+import javax.crypto.DecapsulateException;
+import javax.crypto.KEM;
+import javax.crypto.SecretKey;
+import javax.crypto.spec.SecretKeySpec;
+import java.security.*;
+
+public class DhkemKeyEncapsulationMechanism implements NoiseKeyEncapsulationMechanism {
+
+  private final KEM kem;
+  private final KeyPairGenerator keyPairGenerator;
+  private final KeyFactory keyFactory;
+
+  public DhkemKeyEncapsulationMechanism() throws NoSuchAlgorithmException {
+    this.keyPairGenerator = KeyPairGenerator.getInstance("X25519");
+    this.keyFactory = KeyFactory.getInstance("X25519");
+    this.kem = KEM.getInstance("DHKEM");
+  }
+
+  @Override
+  public String getName() {
+    return "DHKEM";
+  }
+
+  @Override
+  public KeyPair generateKeyPair() {
+    return keyPairGenerator.generateKeyPair();
+  }
+
+  @Override
+  public KEM.Encapsulated encapsulate(final PublicKey publicKey) {
+    try {
+      return kem.newEncapsulator(publicKey).encapsulate();
+    } catch (final InvalidKeyException e) {
+      throw new IllegalArgumentException("Invalid public key for encapsulation", e);
+    }
+  }
+
+  @Override
+  public byte[] decapsulate(final PrivateKey privateKey, final byte[] encapsulation) {
+    try {
+      return serializeSharedSecret(kem.newDecapsulator(privateKey).decapsulate(encapsulation));
+    } catch (final DecapsulateException e) {
+      throw new IllegalArgumentException("Invalid encapsulation", e);
+    } catch (final InvalidKeyException e) {
+      throw new IllegalArgumentException("Invalid private key for decapsulation", e);
+    }
+  }
+
+  @Override
+  public int getPublicKeyLength() {
+    return 32;
+  }
+
+  @Override
+  public int getEncapsulationLength() {
+    return 32;
+  }
+
+  @Override
+  public byte[] serializePublicKey(final PublicKey publicKey) {
+    return XECUtil.serializePublicKey(publicKey, getPublicKeyLength(), XECUtil.X25519_X509_PREFIX);
+  }
+
+  @Override
+  public PublicKey deserializePublicKey(final byte[] publicKeyBytes) {
+    return XECUtil.deserializePublicKey(publicKeyBytes, getPublicKeyLength(), XECUtil.X25519_X509_PREFIX, keyFactory);
+  }
+
+  @Override
+  public byte[] serializeSharedSecret(final SecretKey sharedSecret) {
+    // For DHKEM, the shared secret has a "raw" encoding
+    return sharedSecret.getEncoded();
+  }
+
+  @Override
+  public SecretKey deserializeSharedSecret(final byte[] sharedSecretBytes) {
+    return new SecretKeySpec(sharedSecretBytes, "Generic");
+  }
+}

src/main/java/com/eatthepath/noise/component/NoiseKeyEncapsulationMechanism.java

@@ -4,13 +4,17 @@
 import javax.crypto.KEM;
 import javax.crypto.SecretKey;
 import java.security.KeyPair;
+import java.security.NoSuchAlgorithmException;
 import java.security.PrivateKey;
 import java.security.PublicKey;
 
 public interface NoiseKeyEncapsulationMechanism {
 
-  static NoiseKeyEncapsulationMechanism getInstance(final String name) {
-    throw new IllegalArgumentException("Unrecognized key encapsulation method name: " + name);
+  static NoiseKeyEncapsulationMechanism getInstance(final String name) throws NoSuchAlgorithmException {
+    return switch (name) {
+      case "DHKEM" -> new DhkemKeyEncapsulationMechanism();
+      default -> throw new IllegalArgumentException("Unrecognized key encapsulation method name: " + name);
+    };
   }
 
   String getName();

src/test/java/com/eatthepath/noise/NoiseProtocolIntegrationTest.java

@@ -8,6 +8,7 @@
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
 import org.junit.jupiter.api.Assumptions;
 import org.junit.jupiter.api.Named;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
@@ -18,6 +19,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
 import java.security.*;
 import java.security.spec.NamedParameterSpec;
 import java.util.List;
@@ -665,4 +667,78 @@ public void nextBytes(final byte[] bytes) {
       throw new RuntimeException(e);
     }
   }
+
+  @Test
+  void dhkemHfs() throws NoSuchAlgorithmException, AEADBadTagException {
+    final NoiseHandshake initiatorHandshake = NoiseHandshakeBuilder.forNNHfsInitiator()
+        .setKeyAgreement("25519")
+        .setKeyEncapsulationMechanism("DHKEM")
+        .setCipher("AESGCM")
+        .setHash("SHA256")
+        .build();
+
+    final NoiseHandshake responderHandshake = NoiseHandshakeBuilder.forNNHfsResponder()
+        .setKeyAgreement("25519")
+        .setKeyEncapsulationMechanism("DHKEM")
+        .setCipher("AESGCM")
+        .setHash("SHA256")
+        .build();
+
+    // -> e (with an empty payload)
+    final byte[] initiatorEMessage = initiatorHandshake.writeMessage((byte[]) null);
+    responderHandshake.readMessage(initiatorEMessage);
+
+    // <- e, ee (with an empty payload)
+    final byte[] responderEEeMessage = responderHandshake.writeMessage((byte[]) null);
+    initiatorHandshake.readMessage(responderEEeMessage);
+
+    assertTrue(initiatorHandshake.isDone());
+    assertTrue(responderHandshake.isDone());
+
+    final NoiseTransport initiatorTransport = initiatorHandshake.toTransport();
+    final NoiseTransport responderTransport = responderHandshake.toTransport();
+
+    final byte[] originalPlaintext = "Original payload!".getBytes(StandardCharsets.UTF_8);
+    final byte[] originalCiphertext = initiatorTransport.writeMessage(originalPlaintext);
+    final byte[] decryptedPlaintext = responderTransport.readMessage(originalCiphertext);
+
+    assertArrayEquals(originalPlaintext, decryptedPlaintext);
+  }
+
+  @Test
+  void dhkemHfsByteBuffer() throws NoSuchAlgorithmException, AEADBadTagException {
+    final NoiseHandshake initiatorHandshake = NoiseHandshakeBuilder.forNNHfsInitiator()
+        .setKeyAgreement("25519")
+        .setKeyEncapsulationMechanism("DHKEM")
+        .setCipher("AESGCM")
+        .setHash("SHA256")
+        .build();
+
+    final NoiseHandshake responderHandshake = NoiseHandshakeBuilder.forNNHfsResponder()
+        .setKeyAgreement("25519")
+        .setKeyEncapsulationMechanism("DHKEM")
+        .setCipher("AESGCM")
+        .setHash("SHA256")
+        .build();
+
+    // -> e (with an empty payload)
+    final ByteBuffer initiatorEMessage = initiatorHandshake.writeMessage((ByteBuffer) null);
+    responderHandshake.readMessage(initiatorEMessage);
+
+    // <- e, ee (with an empty payload)
+    final ByteBuffer responderEEeMessage = responderHandshake.writeMessage((ByteBuffer) null);
+    initiatorHandshake.readMessage(responderEEeMessage);
+
+    assertTrue(initiatorHandshake.isDone());
+    assertTrue(responderHandshake.isDone());
+
+    final NoiseTransport initiatorTransport = initiatorHandshake.toTransport();
+    final NoiseTransport responderTransport = responderHandshake.toTransport();
+
+    final ByteBuffer originalPlaintext = ByteBuffer.wrap("Original payload!".getBytes(StandardCharsets.UTF_8));
+    final ByteBuffer originalCiphertext = initiatorTransport.writeMessage(originalPlaintext);
+    final ByteBuffer decryptedPlaintext = responderTransport.readMessage(originalCiphertext);
+
+    assertEquals(originalPlaintext.rewind(), decryptedPlaintext);
+  }
 }