From 89a44f176c5d89c039bf5d00e39bb8dccaae3c66 Mon Sep 17 00:00:00 2001 From: gbrodman Date: Mon, 12 Aug 2019 17:41:29 -0400 Subject: [PATCH] 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 --- .../GenerateAllocationTokensCommand.java | 104 +++++++++++++----- .../GenerateAllocationTokensCommandTest.java | 77 ++++++++++++- 2 files changed, 154 insertions(+), 27 deletions(-) diff --git a/core/src/main/java/google/registry/tools/GenerateAllocationTokensCommand.java b/core/src/main/java/google/registry/tools/GenerateAllocationTokensCommand.java index 08a869f25..d7ba3e622 100644 --- a/core/src/main/java/google/registry/tools/GenerateAllocationTokensCommand.java +++ b/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 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( - (numTokens > 0) ^ (domainNamesFile != null), - "Must specify either --number or --domain_names_file, but not both"); - - 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"); - + verifyInput(); + if (tokenStrings != null) { + numTokens = tokenStrings.size(); + } Deque domainNames; if (domainNamesFile == null) { domainNames = null; @@ -166,8 +162,7 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi { int tokensSaved = 0; do { ImmutableSet tokens = - generateTokens(BATCH_SIZE).stream() - .limit(numTokens - tokensSaved) + getNextTokenBatch(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 existingTokenStrings = + getExistingTokenStrings(ImmutableSet.copyOf(tokenStrings)); + checkArgument( + existingTokenStrings.isEmpty(), + String.format( + "Cannot create specified tokens; the following tokens already exist: %s", + existingTokenStrings)); + } + + private Stream 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 tokens) { Collection savedTokens = @@ -210,14 +261,17 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi { stringGenerator.createStrings(tokenLength, count).stream() .map(s -> prefix + s) .collect(toImmutableSet()); + ImmutableSet existingTokenStrings = getExistingTokenStrings(candidates); + return ImmutableSet.copyOf(difference(candidates, existingTokenStrings)); + } + + private ImmutableSet getExistingTokenStrings(ImmutableSet candidates) { ImmutableSet> existingTokenKeys = candidates.stream() .map(input -> Key.create(AllocationToken.class, input)) .collect(toImmutableSet()); - ImmutableSet existingTokenStrings = - ofy().load().keys(existingTokenKeys).values().stream() - .map(AllocationToken::getToken) - .collect(toImmutableSet()); - return ImmutableSet.copyOf(difference(candidates, existingTokenStrings)); + return ofy().load().keys(existingTokenKeys).values().stream() + .map(AllocationToken::getToken) + .collect(toImmutableSet()); } } diff --git a/core/src/test/java/google/registry/tools/GenerateAllocationTokensCommandTest.java b/core/src/test/java/google/registry/tools/GenerateAllocationTokensCommandTest.java index dba124ad2..0a2cd64a3 100644 --- a/core/src/test/java/google/registry/tools/GenerateAllocationTokensCommandTest.java +++ b/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 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(