Add TINK output prefix to KmsAeadKey.

PiperOrigin-RevId: 618775859
Change-Id: Ie50a5f211871e5339216d01ca29aa7c0a04f5d69

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

@@ -144,6 +144,7 @@ java_library(
         "//src/main/java/com/google/crypto/tink:key_manager",
         "//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/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",
@@ -574,6 +575,7 @@ android_library(
         "//src/main/java/com/google/crypto/tink:key_manager-android",
         "//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/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",
@@ -1037,13 +1039,19 @@ java_library(
 java_library(
     name = "legacy_kms_aead_parameters",
     srcs = ["LegacyKmsAeadParameters.java"],
-    deps = [":aead_parameters"],
+    deps = [
+        ":aead_parameters",
+        "@maven//:com_google_errorprone_error_prone_annotations",
+    ],
 )
 
 android_library(
     name = "legacy_kms_aead_parameters-android",
     srcs = ["LegacyKmsAeadParameters.java"],
-    deps = [":aead_parameters-android"],
+    deps = [
+        ":aead_parameters-android",
+        "@maven//:com_google_errorprone_error_prone_annotations",
+    ],
 )
 
 android_library(
@@ -1054,6 +1062,7 @@ android_library(
         ":legacy_kms_aead_parameters-android",
         "//src/main/java/com/google/crypto/tink:key-android",
         "//src/main/java/com/google/crypto/tink/util:bytes-android",
+        "@maven//:com_google_code_findbugs_jsr305",
     ],
 )
 
@@ -1065,6 +1074,7 @@ java_library(
         ":legacy_kms_aead_parameters",
         "//src/main/java/com/google/crypto/tink:key",
         "//src/main/java/com/google/crypto/tink/util:bytes",
+        "@maven//:com_google_code_findbugs_jsr305",
     ],
 )
 

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

@@ -21,6 +21,7 @@
 import com.google.crypto.tink.KeyManager;
 import com.google.crypto.tink.KeyTemplate;
 import com.google.crypto.tink.KmsClients;
+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;
@@ -38,7 +39,10 @@
  */
 public final class KmsAeadKeyManager {
   private static Aead create(LegacyKmsAeadKey key) throws GeneralSecurityException {
-    return KmsClients.get(key.getParameters().keyUri()).getAead(key.getParameters().keyUri());
+
+    Aead rawAead =
+        KmsClients.get(key.getParameters().keyUri()).getAead(key.getParameters().keyUri());
+    return LegacyFullAead.create(rawAead, key.getOutputPrefix());
   }
 
   private static final PrimitiveConstructor<LegacyKmsAeadKey, Aead>
@@ -66,11 +70,7 @@ private static Aead create(LegacyKmsAeadKey key) throws GeneralSecurityException
   private static LegacyKmsAeadKey newKey(
       LegacyKmsAeadParameters parameters, @Nullable Integer idRequirement)
       throws GeneralSecurityException {
-    if (idRequirement != null) {
-      throw new GeneralSecurityException(
-          "Id Requirement is not supported for LegacyKmsEnvelopeAeadKey");
-    }
-    return LegacyKmsAeadKey.create(parameters);
+    return LegacyKmsAeadKey.create(parameters, idRequirement);
   }
 
   @SuppressWarnings("InlineLambdaConstant") // We need a correct Object#equals in registration.

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

@@ -18,7 +18,10 @@
 
 import com.google.crypto.tink.Key;
 import com.google.crypto.tink.util.Bytes;
+import java.nio.ByteBuffer;
 import java.security.GeneralSecurityException;
+import java.util.Objects;
+import javax.annotation.Nullable;
 
 /**
  * Describes an Aead backed by a KMS.
@@ -31,19 +34,47 @@
  */
 public class LegacyKmsAeadKey extends AeadKey {
   private final LegacyKmsAeadParameters parameters;
+  private final Bytes outputPrefix;
+  @Nullable private final Integer idRequirement;
 
-  private LegacyKmsAeadKey(LegacyKmsAeadParameters parameters) {
+  private LegacyKmsAeadKey(
+      LegacyKmsAeadParameters parameters, Bytes outputPrefix, @Nullable Integer idRequirement) {
     this.parameters = parameters;
+    this.outputPrefix = outputPrefix;
+    this.idRequirement = idRequirement;
+  }
+
+  public static LegacyKmsAeadKey create(
+      LegacyKmsAeadParameters parameters, @Nullable Integer idRequirement)
+      throws GeneralSecurityException {
+    Bytes outputPrefix;
+    if (parameters.variant() == LegacyKmsAeadParameters.Variant.TINK) {
+      if (idRequirement == null) {
+        throw new GeneralSecurityException(
+            "For given Variant TINK the value of idRequirement must be non-null");
+      }
+      outputPrefix =
+          Bytes.copyFrom(ByteBuffer.allocate(5).put((byte) 1).putInt(idRequirement).array());
+    } else if (parameters.variant() == LegacyKmsAeadParameters.Variant.NO_PREFIX) {
+      if (idRequirement != null) {
+        throw new GeneralSecurityException(
+            "For given Variant NO_PREFIX the value of idRequirement must be null");
+      }
+      outputPrefix = Bytes.copyFrom(new byte[] {});
+    } else {
+      throw new GeneralSecurityException("Unknown Variant: " + parameters.variant());
+    }
+    return new LegacyKmsAeadKey(parameters, outputPrefix, idRequirement);
   }
 
   public static LegacyKmsAeadKey create(LegacyKmsAeadParameters parameters)
       throws GeneralSecurityException {
-    return new LegacyKmsAeadKey(parameters);
+    return create(parameters, null);
   }
 
   @Override
   public Bytes getOutputPrefix() {
-    return Bytes.copyFrom(new byte[] {});
+    return outputPrefix;
   }
 
   @Override
@@ -53,7 +84,7 @@ public LegacyKmsAeadParameters getParameters() {
 
   @Override
   public Integer getIdRequirementOrNull() {
-    return null;
+    return idRequirement;
   }
 
   @Override
@@ -62,6 +93,6 @@ public boolean equalsKey(Key o) {
       return false;
     }
     LegacyKmsAeadKey that = (LegacyKmsAeadKey) o;
-    return that.parameters.equals(parameters);
+    return that.parameters.equals(parameters) && Objects.equals(that.idRequirement, idRequirement);
   }
 }

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

@@ -16,29 +16,61 @@
 
 package com.google.crypto.tink.aead;
 
+import com.google.errorprone.annotations.Immutable;
 import java.security.GeneralSecurityException;
 import java.util.Objects;
 
 /** Describes the parameters of an {@link KmsAeadKey} */
 public final class LegacyKmsAeadParameters 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;
+    }
+  }
+
   private final String keyUri;
+  private final Variant variant;
 
-  private LegacyKmsAeadParameters(String keyUri) {
+  private LegacyKmsAeadParameters(String keyUri, Variant variant) {
     this.keyUri = keyUri;
+    this.variant = variant;
   }
 
   public static LegacyKmsAeadParameters create(String keyUri) throws GeneralSecurityException {
-    return new LegacyKmsAeadParameters(keyUri);
+    return new LegacyKmsAeadParameters(keyUri, Variant.NO_PREFIX);
+  }
+
+  public static LegacyKmsAeadParameters create(String keyUri, Variant variant) {
+    return new LegacyKmsAeadParameters(keyUri, variant);
   }
 
   public String keyUri() {
     return keyUri;
   }
 
+  public Variant variant() {
+    return variant;
+  }
+
   @Override
   public boolean hasIdRequirement() {
-    return false;
+    return variant != Variant.NO_PREFIX;
   }
 
   @Override
@@ -47,16 +79,16 @@ public boolean equals(Object o) {
       return false;
     }
     LegacyKmsAeadParameters that = (LegacyKmsAeadParameters) o;
-    return that.keyUri.equals(keyUri);
+    return that.keyUri.equals(keyUri) && that.variant.equals(variant);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(LegacyKmsAeadParameters.class, keyUri);
+    return Objects.hash(LegacyKmsAeadParameters.class, keyUri, variant);
   }
 
   @Override
   public String toString() {
-    return "LegacyKmsAead Parameters (keyUri: " + keyUri + ")";
+    return "LegacyKmsAead Parameters (keyUri: " + keyUri + ", variant: " + variant + ")";
   }
 }

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

@@ -64,14 +64,38 @@ final class LegacyKmsAeadProtoSerialization {
       KeyParser.create(
           LegacyKmsAeadProtoSerialization::parseKey, TYPE_URL_BYTES, ProtoKeySerialization.class);
 
+  private static OutputPrefixType toProtoOutputPrefixType(LegacyKmsAeadParameters.Variant variant)
+      throws GeneralSecurityException {
+    if (LegacyKmsAeadParameters.Variant.TINK.equals(variant)) {
+      return OutputPrefixType.TINK;
+    }
+    if (LegacyKmsAeadParameters.Variant.NO_PREFIX.equals(variant)) {
+      return OutputPrefixType.RAW;
+    }
+    throw new GeneralSecurityException("Unable to serialize variant: " + variant);
+  }
+
+  private static LegacyKmsAeadParameters.Variant toVariant(OutputPrefixType outputPrefixType)
+      throws GeneralSecurityException {
+    switch (outputPrefixType) {
+      case TINK:
+        return LegacyKmsAeadParameters.Variant.TINK;
+      case RAW:
+        return LegacyKmsAeadParameters.Variant.NO_PREFIX;
+      default:
+        throw new GeneralSecurityException(
+            "Unable to parse OutputPrefixType: " + outputPrefixType.getNumber());
+    }
+  }
+
   private static ProtoParametersSerialization serializeParameters(
       LegacyKmsAeadParameters parameters) throws GeneralSecurityException {
     return ProtoParametersSerialization.create(
         KeyTemplate.newBuilder()
             .setTypeUrl(TYPE_URL)
             .setValue(
                 KmsAeadKeyFormat.newBuilder().setKeyUri(parameters.keyUri()).build().toByteString())
-            .setOutputPrefixType(OutputPrefixType.RAW)
+            .setOutputPrefixType(toProtoOutputPrefixType(parameters.variant()))
             .build());
   }
 
@@ -90,14 +114,8 @@ private static LegacyKmsAeadParameters parseParameters(ProtoParametersSerializat
     } catch (InvalidProtocolBufferException e) {
       throw new GeneralSecurityException("Parsing KmsAeadKeyFormat failed: ", e);
     }
-    if (serialization.getKeyTemplate().getOutputPrefixType() != OutputPrefixType.RAW) {
-      throw new GeneralSecurityException(
-          "Only key templates with RAW are accepted, but got "
-              + serialization.getKeyTemplate().getOutputPrefixType()
-              + " with format "
-              + format);
-    }
-    return LegacyKmsAeadParameters.create(format.getKeyUri());
+    return LegacyKmsAeadParameters.create(
+        format.getKeyUri(), toVariant(serialization.getKeyTemplate().getOutputPrefixType()));
   }
 
   private static ProtoKeySerialization serializeKey(
@@ -110,7 +128,7 @@ private static ProtoKeySerialization serializeKey(
             .build()
             .toByteString(),
         KeyMaterialType.REMOTE,
-        OutputPrefixType.RAW,
+        toProtoOutputPrefixType(key.getParameters().variant()),
         key.getIdRequirementOrNull());
   }
 
@@ -121,10 +139,6 @@ private static LegacyKmsAeadKey parseKey(
       throw new IllegalArgumentException(
           "Wrong type URL in call to LegacyKmsAeadProtoSerialization.parseKey");
     }
-    if (serialization.getOutputPrefixType() != OutputPrefixType.RAW) {
-      throw new GeneralSecurityException(
-          "KmsAeadKey are only accepted with RAW, got " + serialization.getOutputPrefixType());
-    }
     try {
       com.google.crypto.tink.proto.KmsAeadKey protoKey =
           com.google.crypto.tink.proto.KmsAeadKey.parseFrom(
@@ -134,8 +148,9 @@ private static LegacyKmsAeadKey parseKey(
             "KmsAeadKey are only accepted with version 0, got " + protoKey);
       }
       LegacyKmsAeadParameters parameters =
-          LegacyKmsAeadParameters.create(protoKey.getParams().getKeyUri());
-      return LegacyKmsAeadKey.create(parameters);
+          LegacyKmsAeadParameters.create(
+              protoKey.getParams().getKeyUri(), toVariant(serialization.getOutputPrefixType()));
+      return LegacyKmsAeadKey.create(parameters, serialization.getIdRequirementOrNull());
     } catch (InvalidProtocolBufferException e) {
       throw new GeneralSecurityException("Parsing KmsAeadKey failed: ", e);
     }

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

@@ -241,12 +241,14 @@ java_test(
         "//src/main/java/com/google/crypto/tink:aead",
         "//src/main/java/com/google/crypto/tink:kms_clients",
         "//src/main/java/com/google/crypto/tink:registry_cluster",
+        "//src/main/java/com/google/crypto/tink:registry_configuration",
         "//src/main/java/com/google/crypto/tink:tink_proto_keyset_format",
         "//src/main/java/com/google/crypto/tink/aead:aead_config",
         "//src/main/java/com/google/crypto/tink/aead:kms_aead_key_manager",
         "//src/main/java/com/google/crypto/tink/aead:legacy_kms_aead_key",
         "//src/main/java/com/google/crypto/tink/aead:legacy_kms_aead_parameters",
         "//src/main/java/com/google/crypto/tink/internal:key_manager_registry",
+        "//src/main/java/com/google/crypto/tink/internal:util",
         "//src/main/java/com/google/crypto/tink/subtle:random",
         "//src/main/java/com/google/crypto/tink/testing:fake_kms_client",
         "//src/main/java/com/google/crypto/tink/testing:test_util",
@@ -558,6 +560,7 @@ java_test(
         "//src/main/java/com/google/crypto/tink/aead:legacy_kms_aead_key",
         "//src/main/java/com/google/crypto/tink/aead:legacy_kms_aead_parameters",
         "//src/main/java/com/google/crypto/tink/aead:x_cha_cha20_poly1305_key",
+        "//src/main/java/com/google/crypto/tink/util:bytes",
         "//src/main/java/com/google/crypto/tink/util:secret_bytes",
         "@maven//:com_google_truth_truth",
         "@maven//:junit_junit",