From 646dcecd7ed959a477a9407dabab5ea059d0c05a Mon Sep 17 00:00:00 2001 From: mcilwain Date: Tue, 9 Jan 2018 18:05:46 -0800 Subject: [PATCH] Create GenerateAllocationTokens nomulus tool command This creates a specified number of tokens of a given schema, with a dryrun option to not persist them. ------------- Created by MOE: https://github.com/google/moe MOE_MIGRATED_REVID=181403775 --- .../imports/XjcToDomainResourceConverter.java | 7 +- .../GenerateAllocationTokensCommand.java | 126 +++++++++++++++ java/google/registry/tools/RegistryTool.java | 1 + .../registry/tools/RegistryToolComponent.java | 1 + .../registry/tools/RegistryToolModule.java | 23 ++- .../registry/util/RandomStringGenerator.java | 6 +- .../testing/DeterministicStringGenerator.java | 4 +- .../GenerateAllocationTokensCommandTest.java | 153 ++++++++++++++++++ 8 files changed, 307 insertions(+), 14 deletions(-) create mode 100644 java/google/registry/tools/GenerateAllocationTokensCommand.java create mode 100644 javatests/google/registry/tools/GenerateAllocationTokensCommandTest.java diff --git a/java/google/registry/rde/imports/XjcToDomainResourceConverter.java b/java/google/registry/rde/imports/XjcToDomainResourceConverter.java index 0904f6364..563dbddae 100644 --- a/java/google/registry/rde/imports/XjcToDomainResourceConverter.java +++ b/java/google/registry/rde/imports/XjcToDomainResourceConverter.java @@ -57,7 +57,6 @@ import google.registry.xjc.secdns.XjcSecdnsDsDataType; import java.security.NoSuchAlgorithmException; import java.security.ProviderException; import java.security.SecureRandom; -import java.util.Random; import java.util.function.Function; import org.joda.time.DateTime; @@ -65,10 +64,10 @@ import org.joda.time.DateTime; final class XjcToDomainResourceConverter extends XjcToEppResourceConverter { @NonFinalForTesting - static StringGenerator stringGenerator = new RandomStringGenerator( - StringGenerator.Alphabets.BASE_64, getRandom()); + static StringGenerator stringGenerator = + new RandomStringGenerator(StringGenerator.Alphabets.BASE_64, getRandom()); - static Random getRandom() { + static SecureRandom getRandom() { try { return SecureRandom.getInstance("NativePRNG"); } catch (NoSuchAlgorithmException e) { diff --git a/java/google/registry/tools/GenerateAllocationTokensCommand.java b/java/google/registry/tools/GenerateAllocationTokensCommand.java new file mode 100644 index 000000000..ff3b2db02 --- /dev/null +++ b/java/google/registry/tools/GenerateAllocationTokensCommand.java @@ -0,0 +1,126 @@ +// Copyright 2017 The Nomulus Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package google.registry.tools; + +import static com.google.common.collect.ImmutableSet.toImmutableSet; +import static com.google.common.collect.Sets.difference; +import static google.registry.model.ofy.ObjectifyService.ofy; + +import com.beust.jcommander.Parameter; +import com.beust.jcommander.Parameters; +import com.google.appengine.tools.remoteapi.RemoteApiException; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableSet; +import com.googlecode.objectify.Key; +import google.registry.model.domain.AllocationToken; +import google.registry.tools.Command.RemoteApiCommand; +import google.registry.util.NonFinalForTesting; +import google.registry.util.Retrier; +import google.registry.util.StringGenerator; +import java.util.Collection; +import javax.inject.Inject; +import javax.inject.Named; + +/** Command to generate and persist {@link AllocationToken}s. */ +@NonFinalForTesting +@Parameters( + separators = " =", + commandDescription = + "Generates and persists the given number of AllocationTokens, printing each token to stdout." +) +public class GenerateAllocationTokensCommand implements RemoteApiCommand { + + @Parameter( + names = {"-p", "--prefix"}, + description = "Allocation token prefix; defaults to blank" + ) + private String prefix = ""; + + @Parameter( + names = {"-n", "--number"}, + description = "The number of tokens to generate", + required = true + ) + private long numTokens; + + @Parameter( + names = {"-l", "--length"}, + description = "The length of each token, exclusive of the prefix (if specified); defaults to 12" + ) + private int tokenLength = 12; + + @Parameter( + names = {"-d", "--dry_run"}, + description = "Do not actually persist the tokens; defaults to false") + boolean dryRun; + + @Inject @Named("base58StringGenerator") StringGenerator stringGenerator; + @Inject Retrier retrier; + + private static final int BATCH_SIZE = 20; + + @Override + public void run() throws Exception { + int tokensSaved = 0; + do { + ImmutableSet tokens = + generateTokens(BATCH_SIZE) + .stream() + .limit(numTokens - tokensSaved) + .map(t -> new AllocationToken.Builder().setToken(t).build()) + .collect(toImmutableSet()); + // Wrap in a retrier to deal with transient 404 errors (thrown as RemoteApiExceptions). + tokensSaved += retrier.callWithRetry(() -> saveTokens(tokens), RemoteApiException.class); + } while (tokensSaved < numTokens); + } + + @VisibleForTesting + int saveTokens(final ImmutableSet tokens) { + Collection savedTokens = + dryRun ? tokens : ofy().transact(() -> ofy().save().entities(tokens).now().values()); + savedTokens.stream().map(AllocationToken::getToken).forEach(System.out::println); + return savedTokens.size(); + } + + /** + * This function generates at MOST {@code count} tokens, filtering out already-existing token + * strings. + * + *

Note that in the incredibly rare case that all generated tokens already exist, this function + * may return an empty set. + */ + private ImmutableSet generateTokens(int count) { + ImmutableSet candidates = + stringGenerator + .createStrings(tokenLength, count) + .stream() + .map(s -> prefix + s) + .collect(toImmutableSet()); + 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)); + } +} diff --git a/java/google/registry/tools/RegistryTool.java b/java/google/registry/tools/RegistryTool.java index caa2114ba..9d93df395 100644 --- a/java/google/registry/tools/RegistryTool.java +++ b/java/google/registry/tools/RegistryTool.java @@ -61,6 +61,7 @@ public final class RegistryTool { .put("domain_check_fee", DomainCheckFeeCommand.class) .put("encrypt_escrow_deposit", EncryptEscrowDepositCommand.class) .put("execute_epp", ExecuteEppCommand.class) + .put("generate_allocation_tokens", GenerateAllocationTokensCommand.class) .put("generate_applications_report", GenerateApplicationsReportCommand.class) .put("generate_auction_data", GenerateAuctionDataCommand.class) .put("generate_dns_report", GenerateDnsReportCommand.class) diff --git a/java/google/registry/tools/RegistryToolComponent.java b/java/google/registry/tools/RegistryToolComponent.java index 5a244a487..504584760 100644 --- a/java/google/registry/tools/RegistryToolComponent.java +++ b/java/google/registry/tools/RegistryToolComponent.java @@ -83,6 +83,7 @@ interface RegistryToolComponent { void inject(CreateTldCommand command); void inject(DeployInvoicingPipelineCommand command); void inject(EncryptEscrowDepositCommand command); + void inject(GenerateAllocationTokensCommand command); void inject(GenerateApplicationsReportCommand command); void inject(GenerateDnsReportCommand command); void inject(GenerateEscrowDepositCommand command); diff --git a/java/google/registry/tools/RegistryToolModule.java b/java/google/registry/tools/RegistryToolModule.java index 3e4cbf870..eb7cefac8 100644 --- a/java/google/registry/tools/RegistryToolModule.java +++ b/java/google/registry/tools/RegistryToolModule.java @@ -19,10 +19,10 @@ import dagger.Module; import dagger.Provides; import google.registry.util.RandomStringGenerator; import google.registry.util.StringGenerator; +import google.registry.util.StringGenerator.Alphabets; import java.security.NoSuchAlgorithmException; import java.security.ProviderException; import java.security.SecureRandom; -import java.util.Random; import javax.inject.Named; /** Dagger module for Registry Tool. */ @@ -38,7 +38,7 @@ abstract class RegistryToolModule { abstract StringGenerator provideStringGenerator(RandomStringGenerator stringGenerator); @Provides - static Random provideRandom() { + static SecureRandom provideSecureRandom() { try { return SecureRandom.getInstance("NativePRNG"); } catch (NoSuchAlgorithmException e) { @@ -47,8 +47,21 @@ abstract class RegistryToolModule { } @Provides - @Named("alphabet") - static String provideAlphabet() { - return StringGenerator.Alphabets.BASE_64; + @Named("alphabetBase64") + static String provideAlphabetBase64() { + return Alphabets.BASE_64; + } + + @Provides + @Named("alphabetBase58") + static String provideAlphabetBase58() { + return Alphabets.BASE_58; + } + + @Provides + @Named("base58StringGenerator") + static StringGenerator provideBase58StringGenerator( + @Named("alphabetBase58") String alphabet, SecureRandom random) { + return new RandomStringGenerator(alphabet, random); } } diff --git a/java/google/registry/util/RandomStringGenerator.java b/java/google/registry/util/RandomStringGenerator.java index 4d8c1dcc0..625448d4e 100644 --- a/java/google/registry/util/RandomStringGenerator.java +++ b/java/google/registry/util/RandomStringGenerator.java @@ -16,17 +16,17 @@ package google.registry.util; import static com.google.common.base.Preconditions.checkArgument; -import java.util.Random; +import java.security.SecureRandom; import javax.inject.Inject; import javax.inject.Named; /** Random string generator. */ public class RandomStringGenerator extends StringGenerator { - private final Random random; + private final SecureRandom random; @Inject - public RandomStringGenerator(@Named("alphabet") String alphabet, Random random) { + public RandomStringGenerator(@Named("alphabetBase64") String alphabet, SecureRandom random) { super(alphabet); this.random = random; } diff --git a/javatests/google/registry/testing/DeterministicStringGenerator.java b/javatests/google/registry/testing/DeterministicStringGenerator.java index 5fd8d4021..a91930d62 100644 --- a/javatests/google/registry/testing/DeterministicStringGenerator.java +++ b/javatests/google/registry/testing/DeterministicStringGenerator.java @@ -73,13 +73,13 @@ public class DeterministicStringGenerator extends StringGenerator { } } - public DeterministicStringGenerator(@Named("alphabet") String alphabet, Rule rule) { + public DeterministicStringGenerator(@Named("alphabetBase64") String alphabet, Rule rule) { super(alphabet); iterator = Iterators.cycle(charactersOf(alphabet)); this.rule = rule; } - public DeterministicStringGenerator(@Named("alphabet") String alphabet) { + public DeterministicStringGenerator(@Named("alphabetBase64") String alphabet) { this(alphabet, Rule.DEFAULT); } } diff --git a/javatests/google/registry/tools/GenerateAllocationTokensCommandTest.java b/javatests/google/registry/tools/GenerateAllocationTokensCommandTest.java new file mode 100644 index 000000000..134a10f76 --- /dev/null +++ b/javatests/google/registry/tools/GenerateAllocationTokensCommandTest.java @@ -0,0 +1,153 @@ +// Copyright 2017 The Nomulus Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package google.registry.tools; + +import static com.google.common.truth.Truth.assertThat; +import static google.registry.model.ofy.ObjectifyService.ofy; +import static google.registry.testing.DatastoreHelper.persistResource; +import static google.registry.testing.JUnitBackports.expectThrows; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.spy; + +import com.beust.jcommander.ParameterException; +import com.google.appengine.tools.remoteapi.RemoteApiException; +import com.google.common.collect.ImmutableMap; +import com.googlecode.objectify.Key; +import google.registry.model.domain.AllocationToken; +import google.registry.model.reporting.HistoryEntry; +import google.registry.testing.DeterministicStringGenerator; +import google.registry.testing.DeterministicStringGenerator.Rule; +import google.registry.testing.FakeClock; +import google.registry.testing.FakeSleeper; +import google.registry.util.Retrier; +import google.registry.util.StringGenerator.Alphabets; +import java.io.IOException; +import java.util.function.Function; +import javax.annotation.Nullable; +import org.joda.time.DateTime; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mockito; + +/** Unit tests for {@link GenerateAllocationTokensCommand}. */ +public class GenerateAllocationTokensCommandTest + extends CommandTestCase { + + @Before + public void init() throws IOException { + command.stringGenerator = new DeterministicStringGenerator(Alphabets.BASE_58); + command.retrier = + new Retrier(new FakeSleeper(new FakeClock(DateTime.parse("2000-01-01TZ"))), 3); + } + + @Test + public void testSuccess_oneToken() throws Exception { + runCommand("--prefix", "blah", "--number", "1", "--length", "9"); + assertAllocationTokens(createToken("blah123456789", null)); + assertInStdout("blah123456789"); + } + + @Test + public void testSuccess_threeTokens() throws Exception { + runCommand("--prefix", "foo", "--number", "3", "--length", "10"); + assertAllocationTokens( + createToken("foo123456789A", null), + createToken("fooBCDEFGHJKL", null), + createToken("fooMNPQRSTUVW", null)); + assertInStdout("foo123456789A\nfooBCDEFGHJKL\nfooMNPQRSTUVW"); + } + + @Test + public void testSuccess_defaults() throws Exception { + runCommand("--number", "1"); + assertAllocationTokens(createToken("123456789ABC", null)); + assertInStdout("123456789ABC"); + } + + @Test + public void testSuccess_retry() throws Exception { + GenerateAllocationTokensCommand spyCommand = spy(command); + RemoteApiException fakeException = new RemoteApiException("foo", "foo", "foo", new Exception()); + doThrow(fakeException) + .doThrow(fakeException) + .doCallRealMethod() + .when(spyCommand) + .saveTokens(Mockito.any()); + runCommand("--number", "1"); + assertAllocationTokens(createToken("123456789ABC", null)); + assertInStdout("123456789ABC"); + } + + @Test + public void testSuccess_tokenCollision() throws Exception { + AllocationToken existingToken = + persistResource(new AllocationToken.Builder().setToken("DEADBEEF123456789ABC").build()); + runCommand("--number", "1", "--prefix", "DEADBEEF"); + assertAllocationTokens(existingToken, createToken("DEADBEEFDEFGHJKLMNPQ", null)); + assertInStdout("DEADBEEFDEFGHJKLMNPQ"); + } + + @Test + public void testSuccess_dryRun_outputsButDoesntSave() throws Exception { + runCommand("--prefix", "foo", "--number", "2", "--length", "10", "--dry_run"); + assertAllocationTokens(); + assertInStdout("foo123456789A\nfooBCDEFGHJKL"); + } + + @Test + public void testSuccess_largeNumberOfTokens() throws Exception { + command.stringGenerator = + new DeterministicStringGenerator(Alphabets.BASE_58, Rule.PREPEND_COUNTER); + runCommand("--prefix", "ooo", "--number", "100", "--length", "16"); + // The deterministic string generator makes it too much hassle to assert about each token, so + // just assert total number. + assertThat(ofy().load().type(AllocationToken.class).count()).isEqualTo(100); + } + + @Test + public void testFailure_mustSpecifyNumberOfTokens() throws Exception { + ParameterException thrown = + expectThrows(ParameterException.class, () -> runCommand("--prefix", "FEET")); + assertThat(thrown).hasMessageThat().contains("The following option is required: -n, --number"); + } + + private void assertAllocationTokens(AllocationToken... expectedTokens) throws Exception { + // Using ImmutableObject comparison here is tricky because the creation/updated timestamps are + // neither easy nor valuable to test here. + ImmutableMap actualTokens = + ofy() + .load() + .type(AllocationToken.class) + .list() + .stream() + .collect(ImmutableMap.toImmutableMap(AllocationToken::getToken, Function.identity())); + assertThat(actualTokens).hasSize(expectedTokens.length); + for (AllocationToken expectedToken : expectedTokens) { + AllocationToken match = actualTokens.get(expectedToken.getToken()); + assertThat(match).isNotNull(); + assertThat(match.getRedemptionHistoryEntry()) + .isEqualTo(expectedToken.getRedemptionHistoryEntry()); + } + } + + private AllocationToken createToken( + String token, @Nullable Key redemptionHistoryEntry) { + AllocationToken.Builder builder = new AllocationToken.Builder().setToken(token); + if (redemptionHistoryEntry != null) { + builder.setRedemptionHistoryEntry(redemptionHistoryEntry); + } + return builder.build(); + } +}