Add new Registry-based key template APIs.

PiperOrigin-RevId: 368164487

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

@@ -247,6 +247,7 @@ java_library(
     name = "key_type_manager",
     srcs = ["KeyTypeManager.java"],
     deps = [
+        ":key_template",
         "//proto:tink_java_proto",
         "//src/main/java/com/google/crypto/tink/annotations:alpha",
         "@com_google_protobuf//:protobuf_javalite",
@@ -257,6 +258,7 @@ android_library(
     name = "key_type_manager-android",
     srcs = ["KeyTypeManager.java"],
     deps = [
+        ":key_template-android",
         "//proto:tink_java_proto_lite",
         "//src/main/java/com/google/crypto/tink/annotations:alpha",
         "@com_google_protobuf//:protobuf_javalite",

java_src/src/main/java/com/google/crypto/tink/KeyTypeManager.java

@@ -180,6 +180,20 @@ public KeyFactory(Class<KeyFormatProtoT> clazz) {
       this.clazz = clazz;
     }
 
+    /**
+     * A container that contains key format and other information that form key templates supported
+     * by this factory.
+     */
+    public static final class KeyFormat<KeyFormatProtoT> {
+      public KeyFormatProtoT keyFormat;
+      public KeyTemplate.OutputPrefixType outputPrefixType;
+
+      public KeyFormat(KeyFormatProtoT keyFormat, KeyTemplate.OutputPrefixType outputPrefixType) {
+        this.keyFormat = keyFormat;
+        this.outputPrefixType = outputPrefixType;
+      }
+    }
+
     /**
      * Returns the class corresponding to the key format protobuffer.
      */
@@ -208,23 +222,27 @@ public abstract KeyFormatProtoT parseKeyFormat(ByteString byteString)
     public abstract KeyT createKey(KeyFormatProtoT keyFormat) throws GeneralSecurityException;
 
     /**
-     * Derives a new key from a given format, using the given {@param pseudoRandomness}.
+     * Derives a new key from a given format, using the given {@code pseudoRandomness}.
      *
-     * <p>Implementations need to note that the given paramter {@param pseudoRandomness} may only
+     * <p>Implementations need to note that the given paramter {@code pseudoRandomness} may only
      * produce a finite amount of randomness. Hence, proper implementations will first obtain all
      * the pseudorandom bytes needed; and only after produce the key.
      *
      * <p>While {@link validateKeyFormat} is called before this method will be called,
-     * implementations must check the version of the given {@param keyFormat}, as {@link
+     * implementations must check the version of the given {@code keyFormat}, as {@link
      * validateKeyFormat} is also called from {@link createKey}.
      *
      * <p>Not every KeyTypeManager needs to implement this; if not implemented a {@link
      * GeneralSecurityException} will be thrown.
      */
     public KeyT deriveKey(KeyFormatProtoT keyFormat, InputStream pseudoRandomness)
         throws GeneralSecurityException {
-      throw new GeneralSecurityException(
-          "deriveKey not implemented for key of type " + clazz.toString());
+      throw new GeneralSecurityException("deriveKey not implemented for key of type " + clazz);
+    }
+
+    /** Returns supported key formats and their names. */
+    public Map<String, KeyFormat<KeyFormatProtoT>> keyFormats() {
+      return Collections.emptyMap();
     }
   }
 

java_src/src/main/java/com/google/crypto/tink/Registry.java

@@ -17,14 +17,14 @@
 package com.google.crypto.tink;
 
 import com.google.crypto.tink.proto.KeyData;
-import com.google.crypto.tink.proto.KeyTemplate;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.InvalidProtocolBufferException;
 import com.google.protobuf.MessageLite;
 import java.io.InputStream;
 import java.security.GeneralSecurityException;
 import java.util.Collections;
 import java.util.Locale;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
@@ -88,6 +88,9 @@ public final class Registry {
   private static final ConcurrentMap<Class<?>, PrimitiveWrapper<?, ?>> primitiveWrapperMap =
       new ConcurrentHashMap<>();
 
+  private static final ConcurrentMap<String, KeyTemplate> keyTemplateMap =
+      new ConcurrentHashMap<>(); // name -> KeyTemplate mapping
+
   /**
    * A container which either is constructed from a {@link KeyTypeManager} or from a {@link
    * KeyManager}.
@@ -328,6 +331,7 @@ static synchronized void reset() {
     newKeyAllowedMap.clear();
     catalogueMap.clear();
     primitiveWrapperMap.clear();
+    keyTemplateMap.clear();
   }
 
   /**
@@ -428,28 +432,66 @@ public static synchronized <P> void registerKeyManager(final KeyManager<P> manag
   }
 
   /**
-   * Throws a general security exception iff there is already a key manager registered for typeURL,
-   * and at least one of the following is true:
-   *   * The class implementing the existing key manager differs from the given one.
-   *   * The value of newKeyAllowed currently registered is false, but the input parameter is true.
+   * Throws a general security exception if one of these conditions holds:
+   *
+   * <ul>
+   *   <li>There is already a key manager registered for {@code typeURL}, and at least one of the
+   *       following is true:
+   *       <ul>
+   *         <li>The class implementing the existing key manager differs from the given one.
+   *         <li>The value of {@code newKeyAllowed} currently registered is false, but the input
+   *             parameter is true.
+   *       </ul>
+   *   <li>The {@code newKeyAllowed} flag is true, and at least one of the following is true:
+   *       <ul>
+   *         <li>The key manager was already registered, but it contains new key templates.
+   *         <li>The key manager is new, but it contains existing key templates.
    */
-  private static synchronized void ensureKeyManagerInsertable(
-      String typeUrl, Class<?> implementingClass, boolean newKeyAllowed)
-      throws GeneralSecurityException {
-    if (!keyManagerMap.containsKey(typeUrl)) {
-      return;
-    }
+  private static synchronized <KeyProtoT extends MessageLite, KeyFormatProtoT extends MessageLite>
+      void ensureKeyManagerInsertable(
+          String typeUrl,
+          Class<?> implementingClass,
+          Map<String, KeyTypeManager.KeyFactory.KeyFormat<KeyFormatProtoT>> keyFormats,
+          boolean newKeyAllowed)
+          throws GeneralSecurityException {
     KeyManagerContainer container = keyManagerMap.get(typeUrl);
-    if (!container.getImplementingClass().equals(implementingClass)) {
-        logger.warning("Attempted overwrite of a registered key manager for key type " + typeUrl);
+    if (container != null && !container.getImplementingClass().equals(implementingClass)) {
+      logger.warning("Attempted overwrite of a registered key manager for key type " + typeUrl);
       throw new GeneralSecurityException(
           String.format(
               "typeUrl (%s) is already registered with %s, cannot be re-registered with %s",
               typeUrl, container.getImplementingClass().getName(), implementingClass.getName()));
     }
-    if (newKeyAllowed && !newKeyAllowedMap.get(typeUrl)) {
+    if (newKeyAllowed && newKeyAllowedMap.containsKey(typeUrl) && !newKeyAllowedMap.get(typeUrl)) {
       throw new GeneralSecurityException("New keys are already disallowed for key type " + typeUrl);
     }
+
+    if (newKeyAllowed) {
+      if (keyManagerMap.containsKey(typeUrl)) {
+        // When re-inserting an already present KeyTypeManager, no new key templates should be
+        // present.
+        for (Map.Entry<String, KeyTypeManager.KeyFactory.KeyFormat<KeyFormatProtoT>> entry :
+            keyFormats.entrySet()) {
+          if (!keyTemplateMap.containsKey(entry.getKey())) {
+            throw new GeneralSecurityException(
+                "Attempted to register a new key template "
+                    + entry.getKey()
+                    + " from an existing key manager of type "
+                    + typeUrl);
+          }
+        }
+      } else {
+        // Check that new key managers can't overwrite existing key templates.
+        for (Map.Entry<String, KeyTypeManager.KeyFactory.KeyFormat<KeyFormatProtoT>> entry :
+            keyFormats.entrySet()) {
+
+          if (keyTemplateMap.containsKey(entry.getKey())) {
+            throw new GeneralSecurityException(
+                "Attempted overwrite of a registered key template " + entry.getKey());
+          }
+        }
+      }
+    }
   }
 
   /**
@@ -470,22 +512,47 @@ public static synchronized <P> void registerKeyManager(
       throw new IllegalArgumentException("key manager must be non-null.");
     }
     String typeUrl = manager.getKeyType();
-    ensureKeyManagerInsertable(typeUrl, manager.getClass(), newKeyAllowed);
+    // Use an empty key format because old-style key managers don't export their key formats
+    ensureKeyManagerInsertable(typeUrl, manager.getClass(), Collections.emptyMap(), newKeyAllowed);
     keyManagerMap.putIfAbsent(typeUrl, createContainerFor(manager));
     newKeyAllowedMap.put(typeUrl, Boolean.valueOf(newKeyAllowed));
   }
 
+  /**
+   * Tries to register {@code manager} for {@code manager.getKeyType()}. If {@code newKeyAllowed} is
+   * true, users can generate new keys with this manager using the {@link Registry#newKey} methods.
+   *
+   * <p>If there is an existing key manager, throws an exception if {@code manager} and the existing
+   * key manager aren't instances of the same class, or if {@code newKeyAllowed} is true while the
+   * existing key manager could not create new keys. Otherwise registration succeeds.
+   *
+   * <p>If {@code newKeyAllowed} is true, also tries to register the key templates supported by
+   * {@code manager}.
+   *
+   * @throws GeneralSecurityException if there's an existing key manager is not an instance of the
+   *     class of {@code manager}, or the registration tries to re-enable the generation of new
+   *     keys.
+   * @throws GeneralSecurityException if there's an existing key template.
+   */
   public static synchronized <KeyProtoT extends MessageLite> void registerKeyManager(
       final KeyTypeManager<KeyProtoT> manager, boolean newKeyAllowed)
       throws GeneralSecurityException {
     if (manager == null) {
       throw new IllegalArgumentException("key manager must be non-null.");
     }
     String typeUrl = manager.getKeyType();
-    ensureKeyManagerInsertable(typeUrl, manager.getClass(), newKeyAllowed);
+    ensureKeyManagerInsertable(
+        typeUrl,
+        manager.getClass(),
+        newKeyAllowed ? manager.keyFactory().keyFormats() : Collections.emptyMap(),
+        newKeyAllowed);
+
     if (!keyManagerMap.containsKey(typeUrl)) {
       keyManagerMap.put(typeUrl, createContainerFor(manager));
       keyDeriverMap.put(typeUrl, createDeriverFor(manager));
+      if (newKeyAllowed) {
+        registerKeyTemplates(typeUrl, manager.keyFactory().keyFormats());
+      }
     }
     newKeyAllowedMap.put(typeUrl, Boolean.valueOf(newKeyAllowed));
   }
@@ -494,13 +561,17 @@ public static synchronized <KeyProtoT extends MessageLite> void registerKeyManag
    * Tries to register {@code manager} for {@code manager.getKeyType()}. If {@code newKeyAllowed} is
    * true, users can generate new keys with this manager using the {@link Registry#newKey} methods.
    *
+   * <p>If {@code newKeyAllowed} is true, also tries to register the key templates supported by
+   * {@code manager}.
+   *
    * <p>If there is an existing key manager, throws an exception if {@code manager} and the existing
    * key manager aren't instances of the same class, or if {@code newKeyAllowed} is true while the
    * existing key manager could not create new keys. Otherwise registration succeeds.
    *
    * @throws GeneralSecurityException if there's an existing key manager is not an instance of the
    *     class of {@code manager}, or the registration tries to re-enable the generation of new
    *     keys.
+   * @throws GeneralSecurityException if there's an existing key template.
    */
   public static synchronized <KeyProtoT extends MessageLite, PublicKeyProtoT extends MessageLite>
       void registerAsymmetricKeyManagers(
@@ -513,8 +584,14 @@ void registerAsymmetricKeyManagers(
     }
     String privateTypeUrl = privateKeyTypeManager.getKeyType();
     String publicTypeUrl = publicKeyTypeManager.getKeyType();
-    ensureKeyManagerInsertable(privateTypeUrl, privateKeyTypeManager.getClass(), newKeyAllowed);
-    ensureKeyManagerInsertable(publicTypeUrl, publicKeyTypeManager.getClass(), false);
+    ensureKeyManagerInsertable(
+        privateTypeUrl,
+        privateKeyTypeManager.getClass(),
+        newKeyAllowed ? privateKeyTypeManager.keyFactory().keyFormats() : Collections.emptyMap(),
+        newKeyAllowed);
+    // No key format because a public key manager cannot create new keys
+    ensureKeyManagerInsertable(
+        publicTypeUrl, publicKeyTypeManager.getClass(), Collections.emptyMap(), false);
     if (privateTypeUrl.equals(publicTypeUrl)) {
       throw new GeneralSecurityException("Private and public key type must be different.");
     }
@@ -548,6 +625,10 @@ void registerAsymmetricKeyManagers(
           privateTypeUrl,
           createPrivateKeyContainerFor(privateKeyTypeManager, publicKeyTypeManager));
       keyDeriverMap.put(privateTypeUrl, createDeriverFor(privateKeyTypeManager));
+      if (newKeyAllowed) {
+        registerKeyTemplates(
+            privateKeyTypeManager.getKeyType(), privateKeyTypeManager.keyFactory().keyFormats());
+      }
     }
     newKeyAllowedMap.put(privateTypeUrl, newKeyAllowed);
     if (!keyManagerMap.containsKey(publicTypeUrl)) {
@@ -558,6 +639,20 @@ void registerAsymmetricKeyManagers(
     newKeyAllowedMap.put(publicTypeUrl, false);
   }
 
+  private static <KeyFormatProtoT extends MessageLite> void registerKeyTemplates(
+      String typeUrl,
+      Map<String, KeyTypeManager.KeyFactory.KeyFormat<KeyFormatProtoT>> keyFormats) {
+    for (Map.Entry<String, KeyTypeManager.KeyFactory.KeyFormat<KeyFormatProtoT>> entry :
+        keyFormats.entrySet()) {
+      keyTemplateMap.put(
+          entry.getKey(),
+          KeyTemplate.create(
+              typeUrl,
+              entry.getValue().keyFormat.toByteArray(),
+              entry.getValue().outputPrefixType));
+    }
+  }
+
   /**
    * Tries to register {@code manager} for the given {@code typeUrl}. Users can generate new keys
    * with this manager using the {@link Registry#newKey} methods.
@@ -739,8 +834,8 @@ public static synchronized KeyData newKeyData(com.google.crypto.tink.KeyTemplate
    *
    * @return a new key
    */
-  public static synchronized MessageLite newKey(KeyTemplate keyTemplate)
-      throws GeneralSecurityException {
+  public static synchronized MessageLite newKey(
+      com.google.crypto.tink.proto.KeyTemplate keyTemplate) throws GeneralSecurityException {
     KeyManager<?> manager = getUntypedKeyManager(keyTemplate.getTypeUrl());
     if (newKeyAllowedMap.get(keyTemplate.getTypeUrl()).booleanValue()) {
       return manager.newKey(keyTemplate.getValue());
@@ -778,7 +873,8 @@ public static synchronized MessageLite newKey(String typeUrl, MessageLite format
    *
    * <p>This functions ignores {@code keyTemplate.getOutputPrefix()}.
    */
-  static synchronized KeyData deriveKey(KeyTemplate keyTemplate, InputStream randomStream)
+  static synchronized KeyData deriveKey(
+      com.google.crypto.tink.proto.KeyTemplate keyTemplate, InputStream randomStream)
       throws GeneralSecurityException {
     String typeUrl = keyTemplate.getTypeUrl();
     if (!keyDeriverMap.containsKey(typeUrl)) {
@@ -965,6 +1061,14 @@ public static <P> P wrap(PrimitiveSet<P> primitiveSet)
     return wrap(primitiveSet, primitiveSet.getPrimitiveClass());
   }
 
+  /**
+   * Returns an immutable map of key templates and their names supported by registered key managers
+   * that are allowed to generate new keys.
+   */
+  public static synchronized Map<String, KeyTemplate> keyTemplates() {
+    return Collections.unmodifiableMap(keyTemplateMap);
+  }
+
   /**
    * Returns the input primitive required when creating a {@code wrappedPrimitive}.
    *

java_src/src/main/java/com/google/crypto/tink/aead/AesGcmKeyManager.java

@@ -32,6 +32,9 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.security.GeneralSecurityException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 
 /**
  * This key manager generates new {@code AesGcmKey} keys and produces new instances of {@code
@@ -116,6 +119,16 @@ public AesGcmKey deriveKey(AesGcmKeyFormat format, InputStream inputStream)
           throw new GeneralSecurityException("Reading pseudorandomness failed", e);
         }
       }
+
+      @Override
+      public Map<String, KeyFactory.KeyFormat<AesGcmKeyFormat>> keyFormats() {
+        Map<String, KeyFactory.KeyFormat<AesGcmKeyFormat>> result = new HashMap<>();
+        result.put("AES128_GCM", createKeyFormat(16, KeyTemplate.OutputPrefixType.TINK));
+        result.put("AES128_GCM_RAW", createKeyFormat(16, KeyTemplate.OutputPrefixType.RAW));
+        result.put("AES256_GCM", createKeyFormat(32, KeyTemplate.OutputPrefixType.TINK));
+        result.put("AES256_GCM_RAW", createKeyFormat(32, KeyTemplate.OutputPrefixType.RAW));
+        return Collections.unmodifiableMap(result);
+      }
     };
   }
 
@@ -197,4 +210,10 @@ private static KeyTemplate createKeyTemplate(
     return KeyTemplate.create(
         new AesGcmKeyManager().getKeyType(), format.toByteArray(), prefixType);
   }
+
+  private static KeyFactory.KeyFormat<AesGcmKeyFormat> createKeyFormat(
+      int keySize, KeyTemplate.OutputPrefixType prefixType) {
+    AesGcmKeyFormat format = AesGcmKeyFormat.newBuilder().setKeySize(keySize).build();
+    return new KeyFactory.KeyFormat<>(format, prefixType);
+  }
 }