Add authEncapsulate and authDecapsulate to HpkeKem, NistCurvesHpkeKem and X25519HpkeKem.

PiperOrigin-RevId: 554812615

6 files changed, 455 insertions, 25 deletions

diff --git a/java_src/src/main/java/com/google/crypto/tink/hybrid/internal/HpkeKem.java b/java_src/src/main/java/com/google/crypto/tink/hybrid/internal/HpkeKem.java
index 2a6a33b3c..e733d4739 100644
--- a/java_src/src/main/java/com/google/crypto/tink/hybrid/internal/HpkeKem.java
+++ b/java_src/src/main/java/com/google/crypto/tink/hybrid/internal/HpkeKem.java
@@ -48,6 +48,31 @@ interface HpkeKem {
       throws GeneralSecurityException;
 
   /**
+   * Similar to {@link encapsulate}, but the output additionally encodes an assurance that the KEM
+   * shared secret was generated by the holder of {@code senderPrivateKey}. Returns a {@link
+   * com.google.crypto.tink.hybrid.internal.HpkeKemEncapOutput} object that contains the raw shared
+   * secret and the encapsulated key. The HPKE RFC refers to this method as AuthEncap(), which is
+   * used by the sender.
+   *
+   * @throws GeneralSecurityException when either the shared secret cannot be generated or the
+   *     shared secret cannot be encapsulated.
+   */
+  HpkeKemEncapOutput authEncapsulate(byte[] recipientPublicKey, HpkeKemPrivateKey senderPrivateKey)
+      throws GeneralSecurityException;
+
+  /**
+   * Extracts the shared secret from {@code encapsulatedKey} using {@code recipientPrivateKey}. The
+   * recipient is assured that the KEM shared secret was generated by the holder of the private key
+   * corresponding to {@code senderPublicKey}. Returns the raw shared secret. The HPKE RFC refers to
+   * this method as AuthDecap(), which is used by the recipient.
+   *
+   * @throws GeneralSecurityException if the shared secret cannot be extracted.
+   */
+  byte[] authDecapsulate(
+      byte[] encapsulatedKey, HpkeKemPrivateKey recipientPrivateKey, byte[] senderPublicKey)
+      throws GeneralSecurityException;
+
+  /**
    * Returns the HPKE KEM algorithm identifier for the underlying KEM implementation.
    *
    * <p>More details at
diff --git a/java_src/src/main/java/com/google/crypto/tink/hybrid/internal/NistCurvesHpkeKem.java b/java_src/src/main/java/com/google/crypto/tink/hybrid/internal/NistCurvesHpkeKem.java
index 201cdcd80..b8074b18a 100644
--- a/java_src/src/main/java/com/google/crypto/tink/hybrid/internal/NistCurvesHpkeKem.java
+++ b/java_src/src/main/java/com/google/crypto/tink/hybrid/internal/NistCurvesHpkeKem.java
@@ -51,12 +51,27 @@ final class NistCurvesHpkeKem implements HpkeKem {
   }
 
   private byte[] deriveKemSharedSecret(
-      byte[] dhSharedSecret, byte[] senderPublicKey, byte[] recipientPublicKey)
+      byte[] dhSharedSecret, byte[] senderEphemeralPublicKey, byte[] recipientPublicKey)
+      throws GeneralSecurityException {
+    byte[] kemContext = Bytes.concat(senderEphemeralPublicKey, recipientPublicKey);
+    return extractAndExpand(dhSharedSecret, kemContext);
+  }
+
+  private byte[] deriveKemSharedSecret(
+      byte[] dhSharedSecret,
+      byte[] senderEphemeralPublicKey,
+      byte[] recipientPublicKey,
+      byte[] senderPublicKey)
+      throws GeneralSecurityException {
+    byte[] kemContext = Bytes.concat(senderEphemeralPublicKey, recipientPublicKey, senderPublicKey);
+    return extractAndExpand(dhSharedSecret, kemContext);
+  }
+
+  private byte[] extractAndExpand(byte[] dhSharedSecret, byte[] kemContext)
       throws GeneralSecurityException {
-    byte[] kemContext = Bytes.concat(senderPublicKey, recipientPublicKey);
     byte[] kemSuiteID = HpkeUtil.kemSuiteId(getKemId());
     return hkdf.extractAndExpand(
-        /*salt=*/ null,
+        /* salt= */ null,
         dhSharedSecret,
         "eae_prk",
         kemContext,
@@ -66,18 +81,18 @@ final class NistCurvesHpkeKem implements HpkeKem {
   }
 
   /** Helper function factored out to facilitate unit testing. */
-  HpkeKemEncapOutput encapsulate(byte[] recipientPublicKey, KeyPair senderKeyPair)
+  HpkeKemEncapOutput encapsulate(byte[] recipientPublicKey, KeyPair senderEphemeralKeyPair)
       throws GeneralSecurityException {
     ECPublicKey recipientECPublicKey =
         EllipticCurves.getEcPublicKey(curve, PointFormatType.UNCOMPRESSED, recipientPublicKey);
     byte[] dhSharedSecret =
         EllipticCurves.computeSharedSecret(
-            (ECPrivateKey) senderKeyPair.getPrivate(), recipientECPublicKey);
+            (ECPrivateKey) senderEphemeralKeyPair.getPrivate(), recipientECPublicKey);
     byte[] senderPublicKey =
         EllipticCurves.pointEncode(
             curve,
             PointFormatType.UNCOMPRESSED,
-            ((ECPublicKey) senderKeyPair.getPublic()).getW());
+            ((ECPublicKey) senderEphemeralKeyPair.getPublic()).getW());
     byte[] kemSharedSecret =
         deriveKemSharedSecret(dhSharedSecret, senderPublicKey, recipientPublicKey);
     return new HpkeKemEncapOutput(kemSharedSecret, senderPublicKey);
@@ -89,6 +104,43 @@ final class NistCurvesHpkeKem implements HpkeKem {
     return encapsulate(recipientPublicKey, keyPair);
   }
 
+  /** Helper function factored out to facilitate unit testing. */
+  HpkeKemEncapOutput authEncapsulate(
+      byte[] recipientPublicKey, KeyPair senderEphemeralKeyPair, HpkeKemPrivateKey senderPrivateKey)
+      throws GeneralSecurityException {
+    ECPublicKey recipientECPublicKey =
+        EllipticCurves.getEcPublicKey(curve, PointFormatType.UNCOMPRESSED, recipientPublicKey);
+    ECPrivateKey privateKey =
+        EllipticCurves.getEcPrivateKey(
+            curve, senderPrivateKey.getSerializedPrivate().toByteArray());
+    byte[] dhSharedSecret =
+        Bytes.concat(
+            EllipticCurves.computeSharedSecret(
+                (ECPrivateKey) senderEphemeralKeyPair.getPrivate(), recipientECPublicKey),
+            EllipticCurves.computeSharedSecret(privateKey, recipientECPublicKey));
+    byte[] senderEphemeralPublicKey =
+        EllipticCurves.pointEncode(
+            curve,
+            PointFormatType.UNCOMPRESSED,
+            ((ECPublicKey) senderEphemeralKeyPair.getPublic()).getW());
+
+    byte[] kemSharedSecret =
+        deriveKemSharedSecret(
+            dhSharedSecret,
+            senderEphemeralPublicKey,
+            recipientPublicKey,
+            senderPrivateKey.getSerializedPublic().toByteArray());
+    return new HpkeKemEncapOutput(kemSharedSecret, senderEphemeralPublicKey);
+  }
+
+  @Override
+  public HpkeKemEncapOutput authEncapsulate(
+      byte[] recipientPublicKey, HpkeKemPrivateKey senderPrivateKey)
+      throws GeneralSecurityException {
+    KeyPair keyPair = EllipticCurves.generateKeyPair(curve);
+    return authEncapsulate(recipientPublicKey, keyPair, senderPrivateKey);
+  }
+
   @Override
   public byte[] decapsulate(byte[] encapsulatedKey, HpkeKemPrivateKey recipientPrivateKey)
       throws GeneralSecurityException {
@@ -103,6 +155,30 @@ final class NistCurvesHpkeKem implements HpkeKem {
   }
 
   @Override
+  public byte[] authDecapsulate(
+      byte[] encapsulatedKey, HpkeKemPrivateKey recipientPrivateKey, byte[] senderPublicKey)
+      throws GeneralSecurityException {
+    ECPrivateKey privateKey =
+        EllipticCurves.getEcPrivateKey(
+            curve, recipientPrivateKey.getSerializedPrivate().toByteArray());
+    ECPublicKey senderEphemeralPublicKey =
+        EllipticCurves.getEcPublicKey(curve, PointFormatType.UNCOMPRESSED, encapsulatedKey);
+
+    byte[] dhSharedSecret =
+        Bytes.concat(
+            EllipticCurves.computeSharedSecret(privateKey, senderEphemeralPublicKey),
+            EllipticCurves.computeSharedSecret(
+                privateKey,
+                EllipticCurves.getEcPublicKey(
+                    curve, PointFormatType.UNCOMPRESSED, senderPublicKey)));
+    return deriveKemSharedSecret(
+        dhSharedSecret,
+        encapsulatedKey,
+        recipientPrivateKey.getSerializedPublic().toByteArray(),
+        senderPublicKey);
+  }
+
+  @Override
   public byte[] getKemId() throws GeneralSecurityException {
     switch (curve) {
       case NIST_P256:
diff --git a/java_src/src/main/java/com/google/crypto/tink/hybrid/internal/X25519HpkeKem.java b/java_src/src/main/java/com/google/crypto/tink/hybrid/internal/X25519HpkeKem.java
index 717faae8f..9bbdb54de 100644
--- a/java_src/src/main/java/com/google/crypto/tink/hybrid/internal/X25519HpkeKem.java
+++ b/java_src/src/main/java/com/google/crypto/tink/hybrid/internal/X25519HpkeKem.java
@@ -33,12 +33,27 @@ final class X25519HpkeKem implements HpkeKem {
   }
 
   private byte[] deriveKemSharedSecret(
-      byte[] dhSharedSecret, byte[] senderPublicKey, byte[] recipientPublicKey)
+      byte[] dhSharedSecret, byte[] senderEphemeralPublicKey, byte[] recipientPublicKey)
+      throws GeneralSecurityException {
+    byte[] kemContext = Bytes.concat(senderEphemeralPublicKey, recipientPublicKey);
+    return extractAndExpand(dhSharedSecret, kemContext);
+  }
+
+  private byte[] deriveKemSharedSecret(
+      byte[] dhSharedSecret,
+      byte[] senderEphemeralPublicKey,
+      byte[] recipientPublicKey,
+      byte[] senderPublicKey)
+      throws GeneralSecurityException {
+    byte[] kemContext = Bytes.concat(senderEphemeralPublicKey, recipientPublicKey, senderPublicKey);
+    return extractAndExpand(dhSharedSecret, kemContext);
+  }
+
+  private byte[] extractAndExpand(byte[] dhSharedSecret, byte[] kemContext)
       throws GeneralSecurityException {
-    byte[] kemContext = Bytes.concat(senderPublicKey, recipientPublicKey);
     byte[] kemSuiteId = HpkeUtil.kemSuiteId(HpkeUtil.X25519_HKDF_SHA256_KEM_ID);
     return hkdf.extractAndExpand(
-        /*salt=*/ null,
+        /* salt= */ null,
         dhSharedSecret,
         "eae_prk",
         kemContext,
@@ -62,6 +77,33 @@ final class X25519HpkeKem implements HpkeKem {
     return encapsulate(recipientPublicKey, X25519.generatePrivateKey());
   }
 
+  /** Helper function factored out to facilitate unit testing. */
+  HpkeKemEncapOutput authEncapsulate(
+      byte[] recipientPublicKey,
+      byte[] senderEphemeralPrivateKey,
+      HpkeKemPrivateKey senderPrivateKey)
+      throws GeneralSecurityException {
+    byte[] dhSharedSecret =
+        Bytes.concat(
+            X25519.computeSharedSecret(senderEphemeralPrivateKey, recipientPublicKey),
+            X25519.computeSharedSecret(
+                senderPrivateKey.getSerializedPrivate().toByteArray(), recipientPublicKey));
+    byte[] senderEphemeralPublicKey = X25519.publicFromPrivate(senderEphemeralPrivateKey);
+    byte[] senderPublicKey =
+        X25519.publicFromPrivate(senderPrivateKey.getSerializedPrivate().toByteArray());
+    byte[] kemSharedSecret =
+        deriveKemSharedSecret(
+            dhSharedSecret, senderEphemeralPublicKey, recipientPublicKey, senderPublicKey);
+    return new HpkeKemEncapOutput(kemSharedSecret, senderEphemeralPublicKey);
+  }
+
+  @Override
+  public HpkeKemEncapOutput authEncapsulate(
+      byte[] recipientPublicKey, HpkeKemPrivateKey senderPrivateKey)
+      throws GeneralSecurityException {
+    return authEncapsulate(recipientPublicKey, X25519.generatePrivateKey(), senderPrivateKey);
+  }
+
   @Override
   public byte[] decapsulate(byte[] encapsulatedKey, HpkeKemPrivateKey recipientPrivateKey)
       throws GeneralSecurityException {
@@ -73,6 +115,20 @@ final class X25519HpkeKem implements HpkeKem {
   }
 
   @Override
+  public byte[] authDecapsulate(
+      byte[] encapsulatedKey, HpkeKemPrivateKey recipientPrivateKey, byte[] senderPublicKey)
+      throws GeneralSecurityException {
+    byte[] privateKey = recipientPrivateKey.getSerializedPrivate().toByteArray();
+    byte[] dhSharedSecret =
+        Bytes.concat(
+            X25519.computeSharedSecret(privateKey, encapsulatedKey),
+            X25519.computeSharedSecret(privateKey, senderPublicKey));
+    byte[] recipientPublicKey = X25519.publicFromPrivate(privateKey);
+    return deriveKemSharedSecret(
+        dhSharedSecret, encapsulatedKey, recipientPublicKey, senderPublicKey);
+  }
+
+  @Override
   public byte[] getKemId() throws GeneralSecurityException {
     if (Arrays.equals(hkdf.getKdfId(), HpkeUtil.HKDF_SHA256_KDF_ID)) {
       return HpkeUtil.X25519_HKDF_SHA256_KEM_ID;
diff --git a/java_src/src/main/java/com/google/crypto/tink/testing/HpkeTestUtil.java b/java_src/src/main/java/com/google/crypto/tink/testing/HpkeTestUtil.java
index 790884aa9..17c293dd7 100644
--- a/java_src/src/main/java/com/google/crypto/tink/testing/HpkeTestUtil.java
+++ b/java_src/src/main/java/com/google/crypto/tink/testing/HpkeTestUtil.java
@@ -53,8 +53,9 @@ public final class HpkeTestUtil {
               testObject.get("kdf_id").getAsInt(),
               testObject.get("aead_id").getAsInt());
       // Filter out test vectors for unsupported modes and/or KEMs.
-      if (Arrays.equals(testId.mode, HpkeUtil.BASE_MODE)) {
-        HpkeTestSetup testSetup =
+      if (Arrays.equals(testId.mode, HpkeUtil.BASE_MODE)
+          || Arrays.equals(testId.mode, HpkeUtil.AUTH_MODE)) {
+        HpkeTestSetup.Builder testSetupBuilder =
             HpkeTestSetup.builder()
                 .setInfo(testObject.get("info").getAsString())
                 .setSenderEphemeralPublicKey(testObject.get("pkEm").getAsString())
@@ -66,8 +67,14 @@ public final class HpkeTestUtil {
                 .setKeyScheduleContext(testObject.get("key_schedule_context").getAsString())
                 .setSecret(testObject.get("secret").getAsString())
                 .setKey(testObject.get("key").getAsString())
-                .setBaseNonce(testObject.get("base_nonce").getAsString())
-                .build();
+                .setBaseNonce(testObject.get("base_nonce").getAsString());
+        if (Arrays.equals(testId.mode, HpkeUtil.AUTH_MODE)) {
+          testSetupBuilder =
+              testSetupBuilder
+                  .setSenderPublicKey(testObject.get("pkSm").getAsString())
+                  .setSenderPrivateKey(testObject.get("skSm").getAsString());
+        }
+        HpkeTestSetup testSetup = testSetupBuilder.build();
         JsonArray encryptionsArray = testObject.get("encryptions").getAsJsonArray();
         List<HpkeTestEncryption> testEncryptions = new ArrayList<>();
         for (JsonElement encryptionElement : encryptionsArray) {
diff --git a/java_src/src/test/java/com/google/crypto/tink/hybrid/internal/NistCurvesHpkeKemTest.java b/java_src/src/test/java/com/google/crypto/tink/hybrid/internal/NistCurvesHpkeKemTest.java
index 53b27c266..3c2d63e49 100644
--- a/java_src/src/test/java/com/google/crypto/tink/hybrid/internal/NistCurvesHpkeKemTest.java
+++ b/java_src/src/test/java/com/google/crypto/tink/hybrid/internal/NistCurvesHpkeKemTest.java
@@ -157,6 +157,34 @@ public final class NistCurvesHpkeKemTest {
   }
 
   @Theory
+  public void authEncapsulate_succeeds(
+      @FromDataPoints("hpkeKemParams") HpkeKemParams hpkeNistKemParams)
+      throws GeneralSecurityException {
+    HpkeTestId testId =
+        new HpkeTestId(
+            HpkeUtil.AUTH_MODE,
+            hpkeNistKemParams.kemId,
+            hpkeNistKemParams.hkdfId,
+            hpkeNistKemParams.aeadId);
+    HpkeTestSetup testSetup = testVectors.get(testId).getTestSetup();
+    EllipticCurves.CurveType curve = curveTypeFromKemId(hpkeNistKemParams.kemId);
+    ECPrivateKey privateKey =
+        EllipticCurves.getEcPrivateKey(curve, testSetup.senderEphemeralPrivateKey);
+    ECPublicKey publicKey =
+        EllipticCurves.getEcPublicKey(
+            curve, PointFormatType.UNCOMPRESSED, testSetup.senderEphemeralPublicKey);
+    NistCurvesHpkeKem kem = NistCurvesHpkeKem.fromCurve(curve);
+    HpkeKemEncapOutput result =
+        kem.authEncapsulate(
+            testSetup.recipientPublicKey,
+            new KeyPair(publicKey, privateKey),
+            NistCurvesHpkeKemPrivateKey.fromBytes(
+                testSetup.senderPrivateKey, testSetup.senderPublicKey, curve));
+    expect.that(result.getSharedSecret()).isEqualTo(testSetup.sharedSecret);
+    expect.that(result.getEncapsulatedKey()).isEqualTo(testSetup.encapsulatedKey);
+  }
+
+  @Theory
   public void decapsulate_succeeds(@FromDataPoints("hpkeKemParams") HpkeKemParams hpkeNistKemParams)
       throws GeneralSecurityException {
     HpkeTestId testId =
@@ -186,6 +214,37 @@ public final class NistCurvesHpkeKemTest {
   }
 
   @Theory
+  public void authDecapsulate_succeeds(
+      @FromDataPoints("hpkeKemParams") HpkeKemParams hpkeNistKemParams)
+      throws GeneralSecurityException {
+    HpkeTestId testId =
+        new HpkeTestId(
+            HpkeUtil.AUTH_MODE,
+            hpkeNistKemParams.kemId,
+            hpkeNistKemParams.hkdfId,
+            hpkeNistKemParams.aeadId);
+    HpkeTestSetup testSetup = testVectors.get(testId).getTestSetup();
+    HpkeKemPrivateKey recipientKeyPair =
+        HpkeKemKeyFactory.createPrivate(
+            HpkePrivateKey.newBuilder()
+                .setPrivateKey(ByteString.copyFrom(testSetup.recipientPrivateKey))
+                .setPublicKey(
+                    HpkePublicKey.newBuilder()
+                        .setPublicKey(ByteString.copyFrom(testSetup.recipientPublicKey))
+                        .setParams(
+                            HpkeParams.newBuilder()
+                                .setKem(kemIdToKemProtoParam(hpkeNistKemParams.kemId))
+                                .build())
+                        .build())
+                .build());
+    NistCurvesHpkeKem kem =
+        NistCurvesHpkeKem.fromCurve(curveTypeFromKemId(hpkeNistKemParams.kemId));
+    byte[] result =
+        kem.authDecapsulate(testSetup.encapsulatedKey, recipientKeyPair, testSetup.senderPublicKey);
+    expect.that(result).isEqualTo(testSetup.sharedSecret);
+  }
+
+  @Theory
   public void encapsulate_failsWithInvalidRecipientPublicKey(
       @FromDataPoints("hpkeKemParams") HpkeKemParams hpkeNistKemParams)
       throws GeneralSecurityException {
@@ -204,6 +263,31 @@ public final class NistCurvesHpkeKemTest {
   }
 
   @Theory
+  public void authEncapsulate_failsWithInvalidRecipientPublicKey(
+      @FromDataPoints("hpkeKemParams") HpkeKemParams hpkeNistKemParams)
+      throws GeneralSecurityException {
+    HpkeTestId testId =
+        new HpkeTestId(
+            HpkeUtil.AUTH_MODE,
+            hpkeNistKemParams.kemId,
+            hpkeNistKemParams.hkdfId,
+            hpkeNistKemParams.aeadId);
+    HpkeTestSetup testSetup = testVectors.get(testId).getTestSetup();
+    EllipticCurves.CurveType curve = curveTypeFromKemId(hpkeNistKemParams.kemId);
+    NistCurvesHpkeKem kem =
+        NistCurvesHpkeKem.fromCurve(curveTypeFromKemId(hpkeNistKemParams.kemId));
+    byte[] invalidRecipientPublicKey =
+        Arrays.copyOf(testSetup.recipientPublicKey, testSetup.recipientPublicKey.length + 2);
+    assertThrows(
+        GeneralSecurityException.class,
+        () ->
+            kem.authEncapsulate(
+                invalidRecipientPublicKey,
+                NistCurvesHpkeKemPrivateKey.fromBytes(
+                    testSetup.senderPrivateKey, testSetup.senderPublicKey, curve)));
+  }
+
+  @Theory
   public void decapsulate_failsWithInvalidEncapsulatedPublicKey(
       @FromDataPoints("hpkeKemParams") HpkeKemParams hpkeNistKemParams)
       throws GeneralSecurityException {
@@ -225,4 +309,29 @@ public final class NistCurvesHpkeKemTest {
         GeneralSecurityException.class,
         () -> kem.decapsulate(invalidEncapsulatedKey, validRecipientPrivateKey));
   }
+
+  @Theory
+  public void authDecapsulate_failsWithInvalidEncapsulatedPublicKey(
+      @FromDataPoints("hpkeKemParams") HpkeKemParams hpkeNistKemParams)
+      throws GeneralSecurityException {
+    HpkeTestId testId =
+        new HpkeTestId(
+            HpkeUtil.AUTH_MODE,
+            hpkeNistKemParams.kemId,
+            hpkeNistKemParams.hkdfId,
+            hpkeNistKemParams.aeadId);
+    HpkeTestSetup testSetup = testVectors.get(testId).getTestSetup();
+    byte[] invalidEncapsulatedKey =
+        Arrays.copyOf(testSetup.encapsulatedKey, testSetup.encapsulatedKey.length + 2);
+    EllipticCurves.CurveType curve = curveTypeFromKemId(hpkeNistKemParams.kemId);
+    NistCurvesHpkeKem kem = NistCurvesHpkeKem.fromCurve(curve);
+    HpkeKemPrivateKey validRecipientPrivateKey =
+        NistCurvesHpkeKemPrivateKey.fromBytes(
+            testSetup.recipientPrivateKey, testSetup.recipientPublicKey, curve);
+    assertThrows(
+        GeneralSecurityException.class,
+        () ->
+            kem.authDecapsulate(
+                invalidEncapsulatedKey, validRecipientPrivateKey, testSetup.senderPublicKey));
+  }
 }
diff --git a/java_src/src/test/java/com/google/crypto/tink/hybrid/internal/X25519HpkeKemTest.java b/java_src/src/test/java/com/google/crypto/tink/hybrid/internal/X25519HpkeKemTest.java
index 69d8df002..d76f922a5 100644
--- a/java_src/src/test/java/com/google/crypto/tink/hybrid/internal/X25519HpkeKemTest.java
+++ b/java_src/src/test/java/com/google/crypto/tink/hybrid/internal/X25519HpkeKemTest.java
@@ -58,9 +58,9 @@ public final class X25519HpkeKemTest {
     testVectors = HpkeTestUtil.parseTestVectors(Files.newReader(new File(path), UTF_8));
   }
 
-  private HpkeTestId getDefaultTestId() {
+  private HpkeTestId getDefaultTestId(byte[] mode) {
     return new HpkeTestId(
-        HpkeUtil.BASE_MODE,
+        mode,
         HpkeUtil.X25519_HKDF_SHA256_KEM_ID,
         HpkeUtil.HKDF_SHA256_KDF_ID,
         HpkeUtil.AES_128_GCM_AEAD_ID);
@@ -72,8 +72,18 @@ public final class X25519HpkeKemTest {
     HpkeTestSetup testSetup = testVectors.get(testId).getTestSetup();
 
     X25519HpkeKem kem = new X25519HpkeKem(new HkdfHpkeKdf(MAC_ALGORITHM));
-    HpkeKemEncapOutput result =
-        kem.encapsulate(testSetup.recipientPublicKey, testSetup.senderEphemeralPrivateKey);
+    HpkeKemEncapOutput result;
+    if (mode == HpkeUtil.BASE_MODE) {
+      result = kem.encapsulate(testSetup.recipientPublicKey, testSetup.senderEphemeralPrivateKey);
+    } else if (mode == HpkeUtil.AUTH_MODE) {
+      result =
+          kem.authEncapsulate(
+              testSetup.recipientPublicKey,
+              testSetup.senderEphemeralPrivateKey,
+              X25519HpkeKemPrivateKey.fromBytes(testSetup.senderPrivateKey));
+    } else {
+      throw new IllegalArgumentException("Unsupported mode: " + mode[0]);
+    }
     expect.that(result.getSharedSecret()).isEqualTo(testSetup.sharedSecret);
     expect.that(result.getEncapsulatedKey()).isEqualTo(testSetup.encapsulatedKey);
   }
@@ -84,10 +94,21 @@ public final class X25519HpkeKemTest {
     HpkeTestSetup testSetup = testVectors.get(testId).getTestSetup();
 
     X25519HpkeKem kem = new X25519HpkeKem(new HkdfHpkeKdf(MAC_ALGORITHM));
-    byte[] result =
-        kem.decapsulate(
-            testSetup.encapsulatedKey,
-            X25519HpkeKemPrivateKey.fromBytes(testSetup.recipientPrivateKey));
+    byte[] result;
+    if (mode == HpkeUtil.BASE_MODE) {
+      result =
+          kem.decapsulate(
+              testSetup.encapsulatedKey,
+              X25519HpkeKemPrivateKey.fromBytes(testSetup.recipientPrivateKey));
+    } else if (mode == HpkeUtil.AUTH_MODE) {
+      result =
+          kem.authDecapsulate(
+              testSetup.encapsulatedKey,
+              X25519HpkeKemPrivateKey.fromBytes(testSetup.recipientPrivateKey),
+              testSetup.senderPublicKey);
+    } else {
+      throw new IllegalArgumentException("Unsupported mode: " + mode[0]);
+    }
     expect.that(result).isEqualTo(testSetup.sharedSecret);
   }
 
@@ -101,6 +122,16 @@ public final class X25519HpkeKemTest {
   }
 
   @Test
+  public void authEncapsulate_succeedsWithX25519HkdfSha256Aes128Gcm()
+      throws GeneralSecurityException {
+    encapsulate(
+        HpkeUtil.AUTH_MODE,
+        HpkeUtil.X25519_HKDF_SHA256_KEM_ID,
+        HpkeUtil.HKDF_SHA256_KDF_ID,
+        HpkeUtil.AES_128_GCM_AEAD_ID);
+  }
+
+  @Test
   public void encapsulate_succeedsWithX25519HkdfSha256Aes256Gcm() throws GeneralSecurityException {
     encapsulate(
         HpkeUtil.BASE_MODE,
@@ -110,6 +141,16 @@ public final class X25519HpkeKemTest {
   }
 
   @Test
+  public void authEncapsulate_succeedsWithX25519HkdfSha256Aes256Gcm()
+      throws GeneralSecurityException {
+    encapsulate(
+        HpkeUtil.AUTH_MODE,
+        HpkeUtil.X25519_HKDF_SHA256_KEM_ID,
+        HpkeUtil.HKDF_SHA256_KDF_ID,
+        HpkeUtil.AES_256_GCM_AEAD_ID);
+  }
+
+  @Test
   public void encapsulate_succeedsWithX25519HkdfSha256ChaCha20Poly1305()
       throws GeneralSecurityException {
     encapsulate(
@@ -120,6 +161,16 @@ public final class X25519HpkeKemTest {
   }
 
   @Test
+  public void authEncapsulate_succeedsWithX25519HkdfSha256ChaCha20Poly1305()
+      throws GeneralSecurityException {
+    encapsulate(
+        HpkeUtil.AUTH_MODE,
+        HpkeUtil.X25519_HKDF_SHA256_KEM_ID,
+        HpkeUtil.HKDF_SHA256_KDF_ID,
+        HpkeUtil.CHACHA20_POLY1305_AEAD_ID);
+  }
+
+  @Test
   public void encapsulate_succeedsWithX25519HkdfSha256ExportOnlyAead()
       throws GeneralSecurityException {
     encapsulate(
@@ -130,23 +181,60 @@ public final class X25519HpkeKemTest {
   }
 
   @Test
+  public void authEncapsulate_succeedsWithX25519HkdfSha256ExportOnlyAead()
+      throws GeneralSecurityException {
+    encapsulate(
+        HpkeUtil.AUTH_MODE,
+        HpkeUtil.X25519_HKDF_SHA256_KEM_ID,
+        HpkeUtil.HKDF_SHA256_KDF_ID,
+        EXPORT_ONLY_AEAD_ID);
+  }
+
+  @Test
   public void encapsulate_failsWithInvalidMacAlgorithm() {
     X25519HpkeKem kem = new X25519HpkeKem(new HkdfHpkeKdf("BadMac"));
-    HpkeTestSetup testSetup = testVectors.get(getDefaultTestId()).getTestSetup();
+    HpkeTestSetup testSetup = testVectors.get(getDefaultTestId(HpkeUtil.BASE_MODE)).getTestSetup();
     byte[] validRecipientPublicKey = testSetup.recipientPublicKey;
     assertThrows(NoSuchAlgorithmException.class, () -> kem.encapsulate(validRecipientPublicKey));
   }
 
   @Test
+  public void authEncapsulate_failsWithInvalidMacAlgorithm() {
+    X25519HpkeKem kem = new X25519HpkeKem(new HkdfHpkeKdf("BadMac"));
+    HpkeTestSetup testSetup = testVectors.get(getDefaultTestId(HpkeUtil.AUTH_MODE)).getTestSetup();
+    byte[] validRecipientPublicKey = testSetup.recipientPublicKey;
+    byte[] senderPrivateKey = testSetup.senderPrivateKey;
+    assertThrows(
+        NoSuchAlgorithmException.class,
+        () ->
+            kem.authEncapsulate(
+                validRecipientPublicKey, X25519HpkeKemPrivateKey.fromBytes(senderPrivateKey)));
+  }
+
+  @Test
   public void encapsulate_failsWithInvalidRecipientPublicKey() {
     X25519HpkeKem kem = new X25519HpkeKem(new HkdfHpkeKdf(MAC_ALGORITHM));
-    HpkeTestSetup testSetup = testVectors.get(getDefaultTestId()).getTestSetup();
+    HpkeTestSetup testSetup = testVectors.get(getDefaultTestId(HpkeUtil.BASE_MODE)).getTestSetup();
     byte[] invalidRecipientPublicKey =
         Arrays.copyOf(testSetup.recipientPublicKey, testSetup.recipientPublicKey.length + 2);
     assertThrows(InvalidKeyException.class, () -> kem.encapsulate(invalidRecipientPublicKey));
   }
 
   @Test
+  public void authEncapsulate_failsWithInvalidRecipientPublicKey() {
+    X25519HpkeKem kem = new X25519HpkeKem(new HkdfHpkeKdf(MAC_ALGORITHM));
+    HpkeTestSetup testSetup = testVectors.get(getDefaultTestId(HpkeUtil.AUTH_MODE)).getTestSetup();
+    byte[] invalidRecipientPublicKey =
+        Arrays.copyOf(testSetup.recipientPublicKey, testSetup.recipientPublicKey.length + 2);
+    byte[] senderPrivateKey = testSetup.senderPrivateKey;
+    assertThrows(
+        InvalidKeyException.class,
+        () ->
+            kem.authEncapsulate(
+                invalidRecipientPublicKey, X25519HpkeKemPrivateKey.fromBytes(senderPrivateKey)));
+  }
+
+  @Test
   public void decapsulate_succeedsWithX25519HkdfSha256Aes128Gcm() throws GeneralSecurityException {
     decapsulate(
         HpkeUtil.BASE_MODE,
@@ -156,6 +244,16 @@ public final class X25519HpkeKemTest {
   }
 
   @Test
+  public void authDecapsulate_succeedsWithX25519HkdfSha256Aes128Gcm()
+      throws GeneralSecurityException {
+    decapsulate(
+        HpkeUtil.AUTH_MODE,
+        HpkeUtil.X25519_HKDF_SHA256_KEM_ID,
+        HpkeUtil.HKDF_SHA256_KDF_ID,
+        HpkeUtil.AES_128_GCM_AEAD_ID);
+  }
+
+  @Test
   public void decapsulate_succeedsWithX25519HkdfSha256Aes256Gcm() throws GeneralSecurityException {
     decapsulate(
         HpkeUtil.BASE_MODE,
@@ -165,6 +263,16 @@ public final class X25519HpkeKemTest {
   }
 
   @Test
+  public void authDecapsulate_succeedsWithX25519HkdfSha256Aes256Gcm()
+      throws GeneralSecurityException {
+    decapsulate(
+        HpkeUtil.AUTH_MODE,
+        HpkeUtil.X25519_HKDF_SHA256_KEM_ID,
+        HpkeUtil.HKDF_SHA256_KDF_ID,
+        HpkeUtil.AES_256_GCM_AEAD_ID);
+  }
+
+  @Test
   public void decapsulate_succeedsWithX25519HkdfSha256ChaCha20Poly1305()
       throws GeneralSecurityException {
     decapsulate(
@@ -175,6 +283,16 @@ public final class X25519HpkeKemTest {
   }
 
   @Test
+  public void authDecapsulate_succeedsWithX25519HkdfSha256ChaCha20Poly1305()
+      throws GeneralSecurityException {
+    decapsulate(
+        HpkeUtil.AUTH_MODE,
+        HpkeUtil.X25519_HKDF_SHA256_KEM_ID,
+        HpkeUtil.HKDF_SHA256_KDF_ID,
+        HpkeUtil.CHACHA20_POLY1305_AEAD_ID);
+  }
+
+  @Test
   public void decapsulate_succeedsWithX25519HkdfSha256ExportOnlyAead()
       throws GeneralSecurityException {
     decapsulate(
@@ -185,9 +303,19 @@ public final class X25519HpkeKemTest {
   }
 
   @Test
+  public void authDecapsulate_succeedsWithX25519HkdfSha256ExportOnlyAead()
+      throws GeneralSecurityException {
+    decapsulate(
+        HpkeUtil.AUTH_MODE,
+        HpkeUtil.X25519_HKDF_SHA256_KEM_ID,
+        HpkeUtil.HKDF_SHA256_KDF_ID,
+        EXPORT_ONLY_AEAD_ID);
+  }
+
+  @Test
   public void decapsulate_failsWithInvalidMacAlgorithm() throws GeneralSecurityException {
     X25519HpkeKem kem = new X25519HpkeKem(new HkdfHpkeKdf("BadMac"));
-    HpkeTestSetup testSetup = testVectors.get(getDefaultTestId()).getTestSetup();
+    HpkeTestSetup testSetup = testVectors.get(getDefaultTestId(HpkeUtil.BASE_MODE)).getTestSetup();
     byte[] validEncapsulatedKey = testSetup.encapsulatedKey;
     HpkeKemPrivateKey validRecipientPrivateKey =
         X25519HpkeKemPrivateKey.fromBytes(testSetup.recipientPrivateKey);
@@ -197,9 +325,22 @@ public final class X25519HpkeKemTest {
   }
 
   @Test
+  public void authDecapsulate_failsWithInvalidMacAlgorithm() throws GeneralSecurityException {
+    X25519HpkeKem kem = new X25519HpkeKem(new HkdfHpkeKdf("BadMac"));
+    HpkeTestSetup testSetup = testVectors.get(getDefaultTestId(HpkeUtil.AUTH_MODE)).getTestSetup();
+    byte[] validEncapsulatedKey = testSetup.encapsulatedKey;
+    byte[] senderPublicKey = testSetup.senderPublicKey;
+    HpkeKemPrivateKey validRecipientPrivateKey =
+        X25519HpkeKemPrivateKey.fromBytes(testSetup.recipientPrivateKey);
+    assertThrows(
+        NoSuchAlgorithmException.class,
+        () -> kem.authDecapsulate(validEncapsulatedKey, validRecipientPrivateKey, senderPublicKey));
+  }
+
+  @Test
   public void decapsulate_failsWithInvalidEncapsulatedPublicKey() throws GeneralSecurityException {
     X25519HpkeKem kem = new X25519HpkeKem(new HkdfHpkeKdf(MAC_ALGORITHM));
-    HpkeTestSetup testSetup = testVectors.get(getDefaultTestId()).getTestSetup();
+    HpkeTestSetup testSetup = testVectors.get(getDefaultTestId(HpkeUtil.BASE_MODE)).getTestSetup();
     byte[] invalidEncapsulatedKey =
         Arrays.copyOf(testSetup.encapsulatedKey, testSetup.encapsulatedKey.length + 2);
     HpkeKemPrivateKey validRecipientPrivateKey =
@@ -210,6 +351,22 @@ public final class X25519HpkeKemTest {
   }
 
   @Test
+  public void authDecapsulate_failsWithInvalidEncapsulatedPublicKey()
+      throws GeneralSecurityException {
+    X25519HpkeKem kem = new X25519HpkeKem(new HkdfHpkeKdf(MAC_ALGORITHM));
+    HpkeTestSetup testSetup = testVectors.get(getDefaultTestId(HpkeUtil.AUTH_MODE)).getTestSetup();
+    byte[] invalidEncapsulatedKey =
+        Arrays.copyOf(testSetup.encapsulatedKey, testSetup.encapsulatedKey.length + 2);
+    HpkeKemPrivateKey validRecipientPrivateKey =
+        X25519HpkeKemPrivateKey.fromBytes(testSetup.recipientPrivateKey);
+    byte[] senderPublicKey = testSetup.senderPublicKey;
+    assertThrows(
+        InvalidKeyException.class,
+        () ->
+            kem.authDecapsulate(invalidEncapsulatedKey, validRecipientPrivateKey, senderPublicKey));
+  }
+
+  @Test
   public void getKemId_succeeds() throws GeneralSecurityException {
     X25519HpkeKem kem = new X25519HpkeKem(new HkdfHpkeKdf(MAC_ALGORITHM));
     expect.that(kem.getKemId()).isEqualTo(HpkeUtil.X25519_HKDF_SHA256_KEM_ID);