Clean up token generation (#205)

* Clean up token generation

- Allow tokenLength of 0
- If specifying a token length of 0, throw an error if numTokens > 1

* Allow generation of 0-length strings

* Allow for --tokens option to generate specific tokens

* Revert String generators and disallow 0 'length' param

* Add verifyInput method and batch the listed tokens

* Check the number of tokens created

core/src/main/java/google/registry/tools/GenerateAllocationTokensCommand.java

@@ -35,6 +35,8 @@ import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Streams;
 import com.google.common.io.Files;
 import com.googlecode.objectify.Key;
 import google.registry.model.domain.token.AllocationToken;
@@ -51,6 +53,7 @@ import java.util.Collection;
 import java.util.Deque;
 import java.util.List;
 import java.util.Optional;
+import java.util.stream.Stream;
 import javax.inject.Inject;
 import javax.inject.Named;
 import org.joda.time.DateTime;
@@ -64,6 +67,13 @@ import org.joda.time.DateTime;
 @NonFinalForTesting
 class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
 
+  @Parameter(
+      names = {"--tokens"},
+      description =
+          "Comma-separated list of exact tokens to generate, otherwise use --prefix or "
+              + "--domain_names_file to generate tokens if the exact strings do not matter")
+  private List<String> tokenStrings;
+
   @Parameter(
       names = {"-p", "--prefix"},
       description = "Allocation token prefix; defaults to blank")
@@ -132,24 +142,10 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
 
   @Override
   public void run() throws IOException {
-    checkArgument(
+    verifyInput();
-        (numTokens > 0) ^ (domainNamesFile != null),
+    if (tokenStrings != null) {
-        "Must specify either --number or --domain_names_file, but not both");
+      numTokens = tokenStrings.size();
-
+    }
-    checkArgument(
-        !(UNLIMITED_USE.equals(tokenType) && CollectionUtils.isNullOrEmpty(tokenStatusTransitions)),
-        "For UNLIMITED_USE tokens, must specify --token_status_transitions");
-
-    // A list consisting solely of the empty string means user error when formatting parameters
-    checkArgument(
-        !ImmutableList.of("").equals(allowedClientIds),
-        "Either omit --allowed_client_ids if all registrars are allowed, or include a"
-            + " comma-separated list");
-
-    checkArgument(
-        !ImmutableList.of("").equals(allowedTlds),
-        "Either omit --allowed_tlds if all TLDs are allowed, or include a comma-separated list");
-
     Deque<String> domainNames;
     if (domainNamesFile == null) {
       domainNames = null;
@@ -166,8 +162,7 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
     int tokensSaved = 0;
     do {
       ImmutableSet<AllocationToken> tokens =
-          generateTokens(BATCH_SIZE).stream()
+          getNextTokenBatch(tokensSaved)
-              .limit(numTokens - tokensSaved)
               .map(
                   t -> {
                     AllocationToken.Builder token =
@@ -189,6 +184,62 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
     } while (tokensSaved < numTokens);
   }
 
+  private void verifyInput() {
+    int inputMethods = 0;
+    if (numTokens > 0) {
+      inputMethods++;
+    }
+    if (domainNamesFile != null) {
+      inputMethods++;
+    }
+    if (tokenStrings != null && !tokenStrings.isEmpty()) {
+      inputMethods++;
+    }
+    checkArgument(
+        inputMethods == 1,
+        "Must specify exactly one of '--number', '--domain_names_file', and '--tokens'");
+
+    checkArgument(
+        tokenLength > 0,
+        "Token length should not be 0. To generate exact tokens, use the --tokens parameter.");
+
+    checkArgument(
+        !(UNLIMITED_USE.equals(tokenType) && CollectionUtils.isNullOrEmpty(tokenStatusTransitions)),
+        "For UNLIMITED_USE tokens, must specify --token_status_transitions");
+
+    // A list consisting solely of the empty string means user error when formatting parameters
+    checkArgument(
+        !ImmutableList.of("").equals(allowedClientIds),
+        "Either omit --allowed_client_ids if all registrars are allowed, or include a"
+            + " comma-separated list");
+
+    checkArgument(
+        !ImmutableList.of("").equals(allowedTlds),
+        "Either omit --allowed_tlds if all TLDs are allowed, or include a comma-separated list");
+
+    if (tokenStrings != null) {
+      verifyTokenStringsDoNotExist();
+    }
+  }
+
+  private void verifyTokenStringsDoNotExist() {
+    ImmutableSet<String> existingTokenStrings =
+        getExistingTokenStrings(ImmutableSet.copyOf(tokenStrings));
+    checkArgument(
+        existingTokenStrings.isEmpty(),
+        String.format(
+            "Cannot create specified tokens; the following tokens already exist: %s",
+            existingTokenStrings));
+  }
+
+  private Stream<String> getNextTokenBatch(int tokensSaved) {
+    if (tokenStrings != null) {
+      return Streams.stream(Iterables.limit(Iterables.skip(tokenStrings, tokensSaved), BATCH_SIZE));
+    } else {
+      return generateTokens(BATCH_SIZE).stream().limit(numTokens - tokensSaved);
+    }
+  }
+
   @VisibleForTesting
   int saveTokens(final ImmutableSet<AllocationToken> tokens) {
     Collection<AllocationToken> savedTokens =
@@ -210,14 +261,17 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
         stringGenerator.createStrings(tokenLength, count).stream()
             .map(s -> prefix + s)
             .collect(toImmutableSet());
+    ImmutableSet<String> existingTokenStrings = getExistingTokenStrings(candidates);
+    return ImmutableSet.copyOf(difference(candidates, existingTokenStrings));
+  }
+
+  private ImmutableSet<String> getExistingTokenStrings(ImmutableSet<String> candidates) {
     ImmutableSet<Key<AllocationToken>> existingTokenKeys =
         candidates.stream()
             .map(input -> Key.create(AllocationToken.class, input))
             .collect(toImmutableSet());
-    ImmutableSet<String> existingTokenStrings =
+    return ofy().load().keys(existingTokenKeys).values().stream()
-        ofy().load().keys(existingTokenKeys).values().stream()
         .map(AllocationToken::getToken)
         .collect(toImmutableSet());
-    return ImmutableSet.copyOf(difference(candidates, existingTokenStrings));
   }
 }

core/src/test/java/google/registry/tools/GenerateAllocationTokensCommandTest.java

@@ -28,9 +28,11 @@ import static org.mockito.Mockito.spy;
 
 import com.beust.jcommander.ParameterException;
 import com.google.appengine.tools.remoteapi.RemoteApiException;
+import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
+import com.google.common.collect.Iterables;
 import com.google.common.io.Files;
 import com.googlecode.objectify.Key;
 import google.registry.model.domain.token.AllocationToken;
@@ -43,6 +45,7 @@ import google.registry.testing.FakeSleeper;
 import google.registry.util.Retrier;
 import google.registry.util.StringGenerator.Alphabets;
 import java.io.File;
+import java.util.Collection;
 import java.util.function.Function;
 import javax.annotation.Nullable;
 import org.joda.time.DateTime;
@@ -172,13 +175,30 @@ public class GenerateAllocationTokensCommandTest
             .build());
   }
 
+  @Test
+  public void testSuccess_specifyTokens() throws Exception {
+    runCommand("--tokens", "foobar,foobaz");
+    assertAllocationTokens(createToken("foobar", null, null), createToken("foobaz", null, null));
+    assertInStdout("foobar", "foobaz");
+  }
+
+  @Test
+  public void testSuccess_specifyManyTokens() throws Exception {
+    command.stringGenerator =
+        new DeterministicStringGenerator(Alphabets.BASE_58, Rule.PREPEND_COUNTER);
+    Collection<String> sampleTokens = command.stringGenerator.createStrings(13, 100);
+    runCommand("--tokens", Joiner.on(",").join(sampleTokens));
+    assertInStdout(Iterables.toArray(sampleTokens, String.class));
+    assertThat(ofy().load().type(AllocationToken.class).count()).isEqualTo(100);
+  }
+
   @Test
   public void testFailure_mustSpecifyNumberOfTokensOrDomainsFile() {
     IllegalArgumentException thrown =
         assertThrows(IllegalArgumentException.class, () -> runCommand("--prefix", "FEET"));
     assertThat(thrown)
         .hasMessageThat()
-        .isEqualTo("Must specify either --number or --domain_names_file, but not both");
+        .isEqualTo("Must specify exactly one of '--number', '--domain_names_file', and '--tokens'");
   }
 
   @Test
@@ -193,7 +213,48 @@ public class GenerateAllocationTokensCommandTest
                     "--domain_names_file", "/path/to/blaaaaah"));
     assertThat(thrown)
         .hasMessageThat()
-        .isEqualTo("Must specify either --number or --domain_names_file, but not both");
+        .isEqualTo("Must specify exactly one of '--number', '--domain_names_file', and '--tokens'");
+  }
+
+  @Test
+  public void testFailure_mustNotSpecifyBothNumberOfTokensAndTokenStrings() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                runCommand(
+                    "--prefix", "FEET",
+                    "--number", "999",
+                    "--tokens", "token1,token2"));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Must specify exactly one of '--number', '--domain_names_file', and '--tokens'");
+  }
+
+  @Test
+  public void testFailure_mustNotSpecifyBothTokenStringsAndDomainsFile() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                runCommand(
+                    "--prefix", "FEET",
+                    "--tokens", "token1,token2",
+                    "--domain_names_file", "/path/to/blaaaaah"));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Must specify exactly one of '--number', '--domain_names_file', and '--tokens'");
+  }
+
+  @Test
+  public void testFailure_specifiesAlreadyExistingToken() throws Exception {
+    runCommand("--tokens", "foobar");
+    beforeCommandTestCase(); // reset the command variables
+    IllegalArgumentException thrown =
+        assertThrows(IllegalArgumentException.class, () -> runCommand("--tokens", "foobar,foobaz"));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Cannot create specified tokens; the following tokens already exist: [foobar]");
   }
 
   @Test
@@ -222,6 +283,18 @@ public class GenerateAllocationTokensCommandTest
         .isInstanceOf(IllegalArgumentException.class);
   }
 
+  @Test
+  public void testFailure_lengthOfZero() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> runCommand("--prefix", "somePrefix", "--number", "1", "--length", "0"));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo(
+            "Token length should not be 0. To generate exact tokens, use the --tokens parameter.");
+  }
+
   @Test
   public void testFailure_unlimitedUseMustHaveTransitions() {
     assertThat(