Skip to content
Permalink
Browse files

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
  • Loading branch information...
gbrodman committed Aug 12, 2019
1 parent d2319b1 commit 89a44f176c5d89c039bf5d00e39bb8dccaae3c66
@@ -35,6 +35,8 @@
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.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 @@
@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 @@

@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<String> domainNames;
if (domainNamesFile == null) {
domainNames = null;
@@ -166,8 +162,7 @@ public void run() throws IOException {
int tokensSaved = 0;
do {
ImmutableSet<AllocationToken> tokens =
generateTokens(BATCH_SIZE).stream()
.limit(numTokens - tokensSaved)
getNextTokenBatch(tokensSaved)
.map(
t -> {
AllocationToken.Builder token =
@@ -189,6 +184,62 @@ public void run() throws IOException {
} 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 @@ int saveTokens(final ImmutableSet<AllocationToken> tokens) {
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 =
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());
}
}
@@ -28,9 +28,11 @@

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.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 void testSuccess_promotionToken() throws Exception {
.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 void testFailure_mustNotSpecifyBothNumberOfTokensAndDomainsFile() {
"--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 void testFailure_invalidTokenStatusTransition() {
.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(

0 comments on commit 89a44f1

Please sign in to comment.
You can’t perform that action at this time.