Add TINK output prefix to LegacyKmsEnvelopeAeadKey.

PiperOrigin-RevId: 619226820
Change-Id: If7d0741f209f22b6c9d5e95edbefb449823e57a7

src/main/java/com/google/crypto/tink/aead/BUILD.bazel

@@ -249,6 +249,7 @@ java_library(
         "//src/main/java/com/google/crypto/tink:key_template",
         "//src/main/java/com/google/crypto/tink:kms_clients",
         "//src/main/java/com/google/crypto/tink:parameters",
+        "//src/main/java/com/google/crypto/tink/aead/internal:legacy_full_aead",
         "//src/main/java/com/google/crypto/tink/config/internal:tink_fips_util",
         "//src/main/java/com/google/crypto/tink/internal:key_manager_registry",
         "//src/main/java/com/google/crypto/tink/internal:legacy_key_manager_impl",
@@ -680,6 +681,7 @@ android_library(
         "//src/main/java/com/google/crypto/tink:key_template-android",
         "//src/main/java/com/google/crypto/tink:kms_clients-android",
         "//src/main/java/com/google/crypto/tink:parameters-android",
+        "//src/main/java/com/google/crypto/tink/aead/internal:legacy_full_aead-android",
         "//src/main/java/com/google/crypto/tink/config/internal:tink_fips_util-android",
         "//src/main/java/com/google/crypto/tink/internal:key_manager_registry-android",
         "//src/main/java/com/google/crypto/tink/internal:legacy_key_manager_impl-android",
@@ -1125,7 +1127,9 @@ android_library(
         ":aead_key-android",
         ":legacy_kms_envelope_aead_parameters-android",
         "//src/main/java/com/google/crypto/tink:key-android",
+        "//src/main/java/com/google/crypto/tink/internal:output_prefix_util-android",
         "//src/main/java/com/google/crypto/tink/util:bytes-android",
+        "@maven//:com_google_code_findbugs_jsr305",
     ],
 )
 
@@ -1136,7 +1140,9 @@ java_library(
         ":aead_key",
         ":legacy_kms_envelope_aead_parameters",
         "//src/main/java/com/google/crypto/tink:key",
+        "//src/main/java/com/google/crypto/tink/internal:output_prefix_util",
         "//src/main/java/com/google/crypto/tink/util:bytes",
+        "@maven//:com_google_code_findbugs_jsr305",
     ],
 )
 

src/main/java/com/google/crypto/tink/aead/KmsEnvelopeAeadKeyManager.java

@@ -22,6 +22,7 @@
 import com.google.crypto.tink.KeyTemplate;
 import com.google.crypto.tink.KmsClients;
 import com.google.crypto.tink.Parameters;
+import com.google.crypto.tink.aead.internal.LegacyFullAead;
 import com.google.crypto.tink.config.internal.TinkFipsUtil;
 import com.google.crypto.tink.internal.KeyManagerRegistry;
 import com.google.crypto.tink.internal.LegacyKeyManagerImpl;
@@ -63,11 +64,7 @@ public class KmsEnvelopeAeadKeyManager {
   private static LegacyKmsEnvelopeAeadKey newKey(
       LegacyKmsEnvelopeAeadParameters parameters, @Nullable Integer idRequirement)
       throws GeneralSecurityException {
-    if (idRequirement != null) {
-      throw new GeneralSecurityException(
-          "Id Requirement is not supported for LegacyKmsEnvelopeAeadKey");
-    }
-    return LegacyKmsEnvelopeAeadKey.create(parameters);
+    return LegacyKmsEnvelopeAeadKey.create(parameters, idRequirement);
   }
 
   @SuppressWarnings("InlineLambdaConstant") // We need a correct Object#equals in registration.
@@ -77,8 +74,11 @@ private static LegacyKmsEnvelopeAeadKey newKey(
   @AccessesPartialKey
   private static Aead create(LegacyKmsEnvelopeAeadKey key) throws GeneralSecurityException {
     String kekUri = key.getParameters().getKekUri();
-    return KmsEnvelopeAead.create(
-        key.getParameters().getDekParametersForNewKeys(), KmsClients.get(kekUri).getAead(kekUri));
+    Aead rawAead =
+        KmsEnvelopeAead.create(
+            key.getParameters().getDekParametersForNewKeys(),
+            KmsClients.get(kekUri).getAead(kekUri));
+    return LegacyFullAead.create(rawAead, key.getOutputPrefix());
   }
 
   private static final PrimitiveConstructor<LegacyKmsEnvelopeAeadKey, Aead>

src/main/java/com/google/crypto/tink/aead/LegacyKmsEnvelopeAeadKey.java

@@ -17,8 +17,11 @@
 package com.google.crypto.tink.aead;
 
 import com.google.crypto.tink.Key;
+import com.google.crypto.tink.internal.OutputPrefixUtil;
 import com.google.crypto.tink.util.Bytes;
 import java.security.GeneralSecurityException;
+import java.util.Objects;
+import javax.annotation.Nullable;
 
 /**
  * Describes an EnvelopeAead backed by a KMS.
@@ -37,19 +40,48 @@
  */
 public class LegacyKmsEnvelopeAeadKey extends AeadKey {
   private final LegacyKmsEnvelopeAeadParameters parameters;
+  private final Bytes outputPrefix;
+  @Nullable private final Integer idRequirement;
 
-  private LegacyKmsEnvelopeAeadKey(LegacyKmsEnvelopeAeadParameters parameters) {
+  private LegacyKmsEnvelopeAeadKey(
+      LegacyKmsEnvelopeAeadParameters parameters,
+      Bytes outputPrefix,
+      @Nullable Integer idRequirement) {
     this.parameters = parameters;
+    this.outputPrefix = outputPrefix;
+    this.idRequirement = idRequirement;
+  }
+
+  public static LegacyKmsEnvelopeAeadKey create(
+      LegacyKmsEnvelopeAeadParameters parameters, @Nullable Integer idRequirement)
+      throws GeneralSecurityException {
+    Bytes outputPrefix;
+    if (parameters.getVariant() == LegacyKmsEnvelopeAeadParameters.Variant.NO_PREFIX) {
+      if (idRequirement != null) {
+        throw new GeneralSecurityException(
+            "For given Variant NO_PREFIX the value of idRequirement must be null");
+      }
+      outputPrefix = OutputPrefixUtil.EMPTY_PREFIX;
+    } else if (parameters.getVariant() == LegacyKmsEnvelopeAeadParameters.Variant.TINK) {
+      if (idRequirement == null) {
+        throw new GeneralSecurityException(
+            "For given Variant TINK the value of idRequirement must be non-null");
+      }
+      outputPrefix = OutputPrefixUtil.getTinkOutputPrefix(idRequirement);
+    } else {
+      throw new GeneralSecurityException("Unknown Variant: " + parameters.getVariant());
+    }
+    return new LegacyKmsEnvelopeAeadKey(parameters, outputPrefix, idRequirement);
   }
 
   public static LegacyKmsEnvelopeAeadKey create(LegacyKmsEnvelopeAeadParameters parameters)
       throws GeneralSecurityException {
-    return new LegacyKmsEnvelopeAeadKey(parameters);
+    return create(parameters, null);
   }
 
   @Override
   public Bytes getOutputPrefix() {
-    return Bytes.copyFrom(new byte[] {});
+    return this.outputPrefix;
   }
 
   @Override
@@ -59,7 +91,7 @@ public LegacyKmsEnvelopeAeadParameters getParameters() {
 
   @Override
   public Integer getIdRequirementOrNull() {
-    return null;
+    return idRequirement;
   }
 
   @Override
@@ -68,6 +100,6 @@ public boolean equalsKey(Key o) {
       return false;
     }
     LegacyKmsEnvelopeAeadKey that = (LegacyKmsEnvelopeAeadKey) o;
-    return that.parameters.equals(parameters);
+    return that.parameters.equals(parameters) && Objects.equals(that.idRequirement, idRequirement);
   }
 }

src/main/java/com/google/crypto/tink/aead/LegacyKmsEnvelopeAeadParameters.java

@@ -17,6 +17,7 @@
 package com.google.crypto.tink.aead;
 
 import com.google.crypto.tink.AccessesPartialKey;
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.Immutable;
 import com.google.errorprone.annotations.RestrictedApi;
 import java.security.GeneralSecurityException;
@@ -64,6 +65,27 @@
  */
 public final class LegacyKmsEnvelopeAeadParameters extends AeadParameters {
 
+  /**
+   * Describes how the prefix is computed. There are two main possibilities: NO_PREFIX (empty
+   * prefix) and TINK (prefix the ciphertext with 0x01 followed by a 4-byte key id in big endian.
+   */
+  @Immutable
+  public static final class Variant {
+    public static final Variant TINK = new Variant("TINK");
+    public static final Variant NO_PREFIX = new Variant("NO_PREFIX");
+
+    private final String name;
+
+    private Variant(String name) {
+      this.name = name;
+    }
+
+    @Override
+    public String toString() {
+      return name;
+    }
+  }
+
   /**
    * Specifies how the DEK in received ciphertexts are parsed.
    *
@@ -107,42 +129,55 @@ public String toString() {
     }
   }
 
+  private final Variant variant;
   private final String kekUri;
   private final DekParsingStrategy dekParsingStrategy;
   private final AeadParameters dekParametersForNewKeys;
 
   private LegacyKmsEnvelopeAeadParameters(
+      Variant variant,
       String kekUri,
       DekParsingStrategy dekParsingStrategy,
       AeadParameters dekParametersForNewKeys) {
+    this.variant = variant;
     this.kekUri = kekUri;
     this.dekParsingStrategy = dekParsingStrategy;
     this.dekParametersForNewKeys = dekParametersForNewKeys;
   }
 
   /** Builder for {@link LegacyKmsEnvelopeAeadParameters}. */
   public static class Builder {
+    @Nullable private Variant variant;
     @Nullable private String kekUri;
     @Nullable private DekParsingStrategy dekParsingStrategy;
     @Nullable private AeadParameters dekParametersForNewKeys;
 
     private Builder() {}
 
+    @CanIgnoreReturnValue
+    public Builder setVariant(Variant variant) {
+      this.variant = variant;
+      return this;
+    }
+
     /**
      * Sets the URI of the KMS to be used.
      *
      * <p>The KMS will be used to encrypt the DEK key as an AEAD.
      */
+    @CanIgnoreReturnValue
     public Builder setKekUri(String kekUri) {
       this.kekUri = kekUri;
       return this;
     }
 
+    @CanIgnoreReturnValue
     public Builder setDekParsingStrategy(DekParsingStrategy dekParsingStrategy) {
       this.dekParsingStrategy = dekParsingStrategy;
       return this;
     }
 
+    @CanIgnoreReturnValue
     public Builder setDekParametersForNewKeys(AeadParameters aeadParameters) {
       this.dekParametersForNewKeys = aeadParameters;
       return this;
@@ -179,6 +214,10 @@ private static boolean parsingStrategyAllowed(
 
     /** Builds the LegacyKmsEnvelopeAeadParameters. */
     public LegacyKmsEnvelopeAeadParameters build() throws GeneralSecurityException {
+      if (variant == null) {
+        // Use NO_PREFIX as default prefix.
+        variant = Variant.NO_PREFIX;
+      }
       if (kekUri == null) {
         throw new GeneralSecurityException("kekUri must be set");
       }
@@ -189,8 +228,7 @@ public LegacyKmsEnvelopeAeadParameters build() throws GeneralSecurityException {
         throw new GeneralSecurityException("dekParametersForNewKeys must be set");
       }
       if (dekParametersForNewKeys.hasIdRequirement()) {
-        throw new GeneralSecurityException(
-            "dekParametersForNewKeys must note have ID Requirements");
+        throw new GeneralSecurityException("dekParametersForNewKeys must not have ID Requirements");
       }
       if (!parsingStrategyAllowed(dekParsingStrategy, dekParametersForNewKeys)) {
         throw new GeneralSecurityException(
@@ -202,7 +240,7 @@ public LegacyKmsEnvelopeAeadParameters build() throws GeneralSecurityException {
       }
 
       return new LegacyKmsEnvelopeAeadParameters(
-          kekUri, dekParsingStrategy, dekParametersForNewKeys);
+          variant, kekUri, dekParsingStrategy, dekParametersForNewKeys);
     }
   }
 
@@ -225,9 +263,13 @@ public String getKekUri() {
     return kekUri;
   }
 
+  public Variant getVariant() {
+    return variant;
+  }
+
   @Override
   public boolean hasIdRequirement() {
-    return false;
+    return variant != Variant.NO_PREFIX;
   }
 
   /**
@@ -252,13 +294,18 @@ public boolean equals(Object o) {
     LegacyKmsEnvelopeAeadParameters that = (LegacyKmsEnvelopeAeadParameters) o;
     return that.dekParsingStrategy.equals(dekParsingStrategy)
         && that.dekParametersForNewKeys.equals(dekParametersForNewKeys)
-        && that.kekUri.equals(kekUri);
+        && that.kekUri.equals(kekUri)
+        && that.variant.equals(variant);
   }
 
   @Override
   public int hashCode() {
     return Objects.hash(
-        LegacyKmsEnvelopeAeadParameters.class, kekUri, dekParsingStrategy, dekParametersForNewKeys);
+        LegacyKmsEnvelopeAeadParameters.class,
+        kekUri,
+        dekParsingStrategy,
+        dekParametersForNewKeys,
+        variant);
   }
 
   @Override
@@ -271,6 +318,9 @@ public String toString() {
         + ", "
         + "dekParametersForNewKeys: "
         + dekParametersForNewKeys
+        + ", "
+        + "variant: "
+        + variant
         + ")";
   }
 }