Skip to content

Commit

Permalink
Add domain name support to AllocationToken entities
Browse files Browse the repository at this point in the history 
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
  • Loading branch information
@CydeWeys @jianglai
CydeWeys authored and jianglai committed Aug 10, 2018
1 parent e397702 commit d80f431
Show file tree
Hide file tree
Showing 5 changed files with 117 additions and 24 deletions.
14 changes: 14 additions & 0 deletions java/google/registry/model/domain/token/AllocationToken.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,12 +21,14 @@
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. */
Expand All @@ -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<HistoryEntry> 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);

Expand All @@ -55,6 +60,10 @@ public boolean isRedeemed() {
return redemptionHistoryEntry != null;
}

public Optional<String> getDomainName() {
return Optional.ofNullable(domainName);
}

public Optional<DateTime> getCreationTime() {
return Optional.ofNullable(creationTime.getTimestamp());
}
Expand Down Expand Up @@ -86,6 +95,11 @@ public Builder setRedemptionHistoryEntry(Key<HistoryEntry> redemptionHistoryEntr
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");
Expand Down
52 changes: 46 additions & 6 deletions java/google/registry/tools/GenerateAllocationTokensCommand.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -50,36 +59,66 @@ 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"
)
private int tokenLength = 12;

@Parameter(
names = {"-d", "--dry_run"},
names = {"--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;
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<String> 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<AllocationToken> 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);
Expand All @@ -90,7 +129,8 @@ public void run() {
int saveTokens(final ImmutableSet<AllocationToken> tokens) {
Collection<AllocationToken> 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();
}

Expand Down
14 changes: 9 additions & 5 deletions javatests/google/registry/model/domain/token/AllocationTokenTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ public void testPersistence() {
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);
Expand All @@ -44,11 +45,14 @@ public void testPersistence() {
@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
Expand Down
1 change: 1 addition & 0 deletions javatests/google/registry/model/testdata/schema.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -307,6 +307,7 @@ class google.registry.model.domain.token.AllocationToken {
com.googlecode.objectify.Key<google.registry.model.reporting.HistoryEntry> 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;
Expand Down
60 changes: 47 additions & 13 deletions javatests/google/registry/tools/GenerateAllocationTokensCommandTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,13 @@
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;
Expand All @@ -33,6 +34,7 @@
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;
Expand All @@ -54,24 +56,24 @@ public void init() {
@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");
}

@Test
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");
}

Expand All @@ -85,7 +87,7 @@ public void testSuccess_retry() throws Exception {
.when(spyCommand)
.saveTokens(Mockito.any());
runCommand("--number", "1");
assertAllocationTokens(createToken("123456789ABC", null));
assertAllocationTokens(createToken("123456789ABC", null, null));
assertInStdout("123456789ABC");
}

Expand All @@ -94,7 +96,7 @@ 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));
assertAllocationTokens(existingToken, createToken("DEADBEEFDEFGHJKLMNPQ", null, null));
assertInStdout("DEADBEEFDEFGHJKLMNPQ");
}

Expand All @@ -116,10 +118,39 @@ public void testSuccess_largeNumberOfTokens() throws Exception {
}

@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) {
Expand All @@ -142,11 +173,14 @@ private void assertAllocationTokens(AllocationToken... expectedTokens) {
}

private AllocationToken createToken(
String token, @Nullable Key<HistoryEntry> redemptionHistoryEntry) {
String token,
@Nullable Key<HistoryEntry> redemptionHistoryEntry,
@Nullable String domainName) {
AllocationToken.Builder builder = new AllocationToken.Builder().setToken(token);
if (redemptionHistoryEntry != null) {
builder.setRedemptionHistoryEntry(redemptionHistoryEntry);
}
builder.setDomainName(domainName);
return builder.build();
}
}

0 comments on commit d80f431

Please sign in to comment.