From d80f431e211a940d67ea0c677279459a5ef0945f Mon Sep 17 00:00:00 2001 From: mcilwain Date: Wed, 8 Aug 2018 07:59:50 -0700 Subject: [PATCH] Add domain name support to AllocationToken entities The design doc is at [] The next step will be to tie this into the domain create flow, and if the domain name is on a reserved list, allow it to be created if the token is specified that has the given domain name on it. ------------- Created by MOE: https://github.com/google/moe MOE_MIGRATED_REVID=207884521 --- .../model/domain/token/AllocationToken.java | 14 +++++ .../GenerateAllocationTokensCommand.java | 52 ++++++++++++++-- .../domain/token/AllocationTokenTest.java | 14 +++-- .../google/registry/model/testdata/schema.txt | 1 + .../GenerateAllocationTokensCommandTest.java | 60 +++++++++++++++---- 5 files changed, 117 insertions(+), 24 deletions(-) diff --git a/java/google/registry/model/domain/token/AllocationToken.java b/java/google/registry/model/domain/token/AllocationToken.java index 083e814e7..806b9cf02 100644 --- a/java/google/registry/model/domain/token/AllocationToken.java +++ b/java/google/registry/model/domain/token/AllocationToken.java @@ -21,12 +21,14 @@ import static google.registry.util.PreconditionsUtils.checkArgumentNotNull; import com.googlecode.objectify.Key; import com.googlecode.objectify.annotation.Entity; import com.googlecode.objectify.annotation.Id; +import com.googlecode.objectify.annotation.Index; import google.registry.model.BackupGroupRoot; import google.registry.model.Buildable; import google.registry.model.CreateAutoTimestamp; import google.registry.model.annotations.ReportedOn; import google.registry.model.reporting.HistoryEntry; import java.util.Optional; +import javax.annotation.Nullable; import org.joda.time.DateTime; /** An entity representing an allocation token. */ @@ -40,6 +42,9 @@ public class AllocationToken extends BackupGroupRoot implements Buildable { /** The key of the history entry for which the token was used. Null if not yet used. */ Key redemptionHistoryEntry; + /** The fully-qualified domain name that this token is limited to, if any. */ + @Nullable @Index String domainName; + /** When this token was created. */ CreateAutoTimestamp creationTime = CreateAutoTimestamp.create(null); @@ -55,6 +60,10 @@ public class AllocationToken extends BackupGroupRoot implements Buildable { return redemptionHistoryEntry != null; } + public Optional getDomainName() { + return Optional.ofNullable(domainName); + } + public Optional getCreationTime() { return Optional.ofNullable(creationTime.getTimestamp()); } @@ -86,6 +95,11 @@ public class AllocationToken extends BackupGroupRoot implements Buildable { return this; } + public Builder setDomainName(@Nullable String domainName) { + getInstance().domainName = domainName; + return this; + } + public Builder setCreationTime(DateTime creationTime) { checkState( getInstance().creationTime.getTimestamp() == null, "creationTime can only be set once"); diff --git a/java/google/registry/tools/GenerateAllocationTokensCommand.java b/java/google/registry/tools/GenerateAllocationTokensCommand.java index 04129402f..7f7f034e1 100644 --- a/java/google/registry/tools/GenerateAllocationTokensCommand.java +++ b/java/google/registry/tools/GenerateAllocationTokensCommand.java @@ -14,22 +14,31 @@ package google.registry.tools; +import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableSet.toImmutableSet; +import static com.google.common.collect.Queues.newArrayDeque; import static com.google.common.collect.Sets.difference; import static google.registry.model.ofy.ObjectifyService.ofy; +import static java.nio.charset.StandardCharsets.UTF_8; 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.base.Joiner; +import com.google.common.base.Splitter; import com.google.common.collect.ImmutableSet; +import com.google.common.io.Files; import com.googlecode.objectify.Key; import google.registry.model.domain.token.AllocationToken; import google.registry.tools.Command.RemoteApiCommand; import google.registry.util.NonFinalForTesting; import google.registry.util.Retrier; import google.registry.util.StringGenerator; +import java.io.File; +import java.io.IOException; import java.util.Collection; +import java.util.Deque; import javax.inject.Inject; import javax.inject.Named; @@ -50,11 +59,16 @@ public class GenerateAllocationTokensCommand implements RemoteApiCommand { @Parameter( names = {"-n", "--number"}, - description = "The number of tokens to generate", - required = true + description = "The number of tokens to generate" ) private long numTokens; + @Parameter( + names = {"-d", "--domain_names_file"}, + description = "A file with a list of newline-delimited domain names to create tokens for" + ) + private String domainNamesFile; + @Parameter( names = {"-l", "--length"}, description = "The length of each token, exclusive of the prefix (if specified); defaults to 12" @@ -62,7 +76,7 @@ public class GenerateAllocationTokensCommand implements RemoteApiCommand { private int tokenLength = 12; @Parameter( - names = {"-d", "--dry_run"}, + names = {"--dry_run"}, description = "Do not actually persist the tokens; defaults to false") boolean dryRun; @@ -70,16 +84,41 @@ public class GenerateAllocationTokensCommand implements RemoteApiCommand { @Inject Retrier retrier; private static final int BATCH_SIZE = 20; + private static final Joiner SKIP_NULLS = Joiner.on(", ").skipNulls(); @Override - public void run() { + public void run() throws IOException { + checkArgument( + (numTokens > 0) ^ (domainNamesFile != null), + "Must specify either --number or --domain_names_file, but not both"); + + Deque domainNames; + if (domainNamesFile == null) { + domainNames = null; + } else { + domainNames = + newArrayDeque( + Splitter.on('\n') + .omitEmptyStrings() + .trimResults() + .split(Files.asCharSource(new File(domainNamesFile), UTF_8).read())); + numTokens = domainNames.size(); + } + int tokensSaved = 0; do { ImmutableSet tokens = generateTokens(BATCH_SIZE) .stream() .limit(numTokens - tokensSaved) - .map(t -> new AllocationToken.Builder().setToken(t).build()) + .map( + t -> { + AllocationToken.Builder token = new AllocationToken.Builder().setToken(t); + if (domainNames != null) { + token.setDomainName(domainNames.removeFirst()); + } + return token.build(); + }) .collect(toImmutableSet()); // Wrap in a retrier to deal with transient 404 errors (thrown as RemoteApiExceptions). tokensSaved += retrier.callWithRetry(() -> saveTokens(tokens), RemoteApiException.class); @@ -90,7 +129,8 @@ public class GenerateAllocationTokensCommand implements RemoteApiCommand { 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); + savedTokens.forEach( + t -> System.out.println(SKIP_NULLS.join(t.getDomainName().orElse(null), t.getToken()))); return savedTokens.size(); } diff --git a/javatests/google/registry/model/domain/token/AllocationTokenTest.java b/javatests/google/registry/model/domain/token/AllocationTokenTest.java index 041059f37..ae3ff14be 100644 --- a/javatests/google/registry/model/domain/token/AllocationTokenTest.java +++ b/javatests/google/registry/model/domain/token/AllocationTokenTest.java @@ -36,6 +36,7 @@ public class AllocationTokenTest extends EntityTestCase { new AllocationToken.Builder() .setToken("abc123") .setRedemptionHistoryEntry(Key.create(HistoryEntry.class, 1L)) + .setDomainName("foo.example") .setCreationTime(DateTime.parse("2010-11-12T05:00:00Z")) .build()); assertThat(ofy().load().entity(token).now()).isEqualTo(token); @@ -44,11 +45,14 @@ public class AllocationTokenTest extends EntityTestCase { @Test public void testIndexing() throws Exception { verifyIndexing( - new AllocationToken.Builder() - .setToken("abc123") - .setCreationTime(DateTime.parse("2010-11-12T05:00:00Z")) - .build(), - "token"); + persistResource( + new AllocationToken.Builder() + .setToken("abc123") + .setDomainName("blahdomain.fake") + .setCreationTime(DateTime.parse("2010-11-12T05:00:00Z")) + .build()), + "token", + "domainName"); } @Test diff --git a/javatests/google/registry/model/testdata/schema.txt b/javatests/google/registry/model/testdata/schema.txt index 8f02b0328..daa644697 100644 --- a/javatests/google/registry/model/testdata/schema.txt +++ b/javatests/google/registry/model/testdata/schema.txt @@ -307,6 +307,7 @@ class google.registry.model.domain.token.AllocationToken { com.googlecode.objectify.Key redemptionHistoryEntry; google.registry.model.CreateAutoTimestamp creationTime; google.registry.model.UpdateAutoTimestamp updateTimestamp; + java.lang.String domainName; } class google.registry.model.eppcommon.AuthInfo$PasswordAuth { java.lang.String repoId; diff --git a/javatests/google/registry/tools/GenerateAllocationTokensCommandTest.java b/javatests/google/registry/tools/GenerateAllocationTokensCommandTest.java index 0112bc0b5..baa071d7c 100644 --- a/javatests/google/registry/tools/GenerateAllocationTokensCommandTest.java +++ b/javatests/google/registry/tools/GenerateAllocationTokensCommandTest.java @@ -18,12 +18,13 @@ 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.assertThrows; +import static java.nio.charset.StandardCharsets.UTF_8; 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.google.common.io.Files; import com.googlecode.objectify.Key; import google.registry.model.domain.token.AllocationToken; import google.registry.model.reporting.HistoryEntry; @@ -33,6 +34,7 @@ import google.registry.testing.FakeClock; import google.registry.testing.FakeSleeper; import google.registry.util.Retrier; import google.registry.util.StringGenerator.Alphabets; +import java.io.File; import java.util.function.Function; import javax.annotation.Nullable; import org.joda.time.DateTime; @@ -54,7 +56,7 @@ public class GenerateAllocationTokensCommandTest @Test public void testSuccess_oneToken() throws Exception { runCommand("--prefix", "blah", "--number", "1", "--length", "9"); - assertAllocationTokens(createToken("blah123456789", null)); + assertAllocationTokens(createToken("blah123456789", null, null)); assertInStdout("blah123456789"); } @@ -62,16 +64,16 @@ public class GenerateAllocationTokensCommandTest public void testSuccess_threeTokens() throws Exception { runCommand("--prefix", "foo", "--number", "3", "--length", "10"); assertAllocationTokens( - createToken("foo123456789A", null), - createToken("fooBCDEFGHJKL", null), - createToken("fooMNPQRSTUVW", null)); + createToken("foo123456789A", null, null), + createToken("fooBCDEFGHJKL", null, null), + createToken("fooMNPQRSTUVW", null, null)); assertInStdout("foo123456789A\nfooBCDEFGHJKL\nfooMNPQRSTUVW"); } @Test public void testSuccess_defaults() throws Exception { runCommand("--number", "1"); - assertAllocationTokens(createToken("123456789ABC", null)); + assertAllocationTokens(createToken("123456789ABC", null, null)); assertInStdout("123456789ABC"); } @@ -85,7 +87,7 @@ public class GenerateAllocationTokensCommandTest .when(spyCommand) .saveTokens(Mockito.any()); runCommand("--number", "1"); - assertAllocationTokens(createToken("123456789ABC", null)); + assertAllocationTokens(createToken("123456789ABC", null, null)); assertInStdout("123456789ABC"); } @@ -94,7 +96,7 @@ public class GenerateAllocationTokensCommandTest AllocationToken existingToken = persistResource(new AllocationToken.Builder().setToken("DEADBEEF123456789ABC").build()); runCommand("--number", "1", "--prefix", "DEADBEEF"); - assertAllocationTokens(existingToken, createToken("DEADBEEFDEFGHJKLMNPQ", null)); + assertAllocationTokens(existingToken, createToken("DEADBEEFDEFGHJKLMNPQ", null, null)); assertInStdout("DEADBEEFDEFGHJKLMNPQ"); } @@ -116,10 +118,39 @@ public class GenerateAllocationTokensCommandTest } @Test - public void testFailure_mustSpecifyNumberOfTokens() { - ParameterException thrown = - assertThrows(ParameterException.class, () -> runCommand("--prefix", "FEET")); - assertThat(thrown).hasMessageThat().contains("The following option is required: -n, --number"); + public void testSuccess_domainNames() throws Exception { + File domainNamesFile = tmpDir.newFile("domain_names.txt"); + Files.asCharSink(domainNamesFile, UTF_8).write("foo1.tld\nboo2.tld\nbaz9.tld\n"); + runCommand("--domain_names_file", domainNamesFile.getPath()); + assertAllocationTokens( + createToken("123456789ABC", null, "foo1.tld"), + createToken("DEFGHJKLMNPQ", null, "boo2.tld"), + createToken("RSTUVWXYZabc", null, "baz9.tld")); + assertInStdout("foo1.tld, 123456789ABC\nboo2.tld, DEFGHJKLMNPQ\nbaz9.tld, RSTUVWXYZabc"); + } + + @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"); + } + + @Test + public void testFailure_mustNotSpecifyBothNumberOfTokensAndDomainsFile() { + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, + () -> + runCommand( + "--prefix", "FEET", + "--number", "999", + "--domain_names_file", "/path/to/blaaaaah")); + assertThat(thrown) + .hasMessageThat() + .isEqualTo("Must specify either --number or --domain_names_file, but not both"); } private void assertAllocationTokens(AllocationToken... expectedTokens) { @@ -142,11 +173,14 @@ public class GenerateAllocationTokensCommandTest } private AllocationToken createToken( - String token, @Nullable Key redemptionHistoryEntry) { + String token, + @Nullable Key redemptionHistoryEntry, + @Nullable String domainName) { AllocationToken.Builder builder = new AllocationToken.Builder().setToken(token); if (redemptionHistoryEntry != null) { builder.setRedemptionHistoryEntry(redemptionHistoryEntry); } + builder.setDomainName(domainName); return builder.build(); } }