From ee1f8b34ea718dea15c79a7e9162af8179cdf83b Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Fri, 19 Nov 2021 12:17:33 -0500
Subject: [PATCH] Add a command for reserving usernames

---
 .../textsecuregcm/WhisperServerService.java   |  2 +
 .../workers/ReserveUsernameCommand.java       | 64 +++++++++++++++++++
 2 files changed, 66 insertions(+)
 create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/workers/ReserveUsernameCommand.java

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index d09394375..51963be92 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -217,6 +217,7 @@ import org.whispersystems.textsecuregcm.websocket.WebSocketAccountAuthenticator;
 import org.whispersystems.textsecuregcm.workers.CertificateCommand;
 import org.whispersystems.textsecuregcm.workers.CheckDynamicConfigurationCommand;
 import org.whispersystems.textsecuregcm.workers.DeleteUserCommand;
+import org.whispersystems.textsecuregcm.workers.ReserveUsernameCommand;
 import org.whispersystems.textsecuregcm.workers.ServerVersionCommand;
 import org.whispersystems.textsecuregcm.workers.SetCrawlerAccelerationTask;
 import org.whispersystems.textsecuregcm.workers.SetRequestLoggingEnabledTask;
@@ -243,6 +244,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     bootstrap.addCommand(new ServerVersionCommand());
     bootstrap.addCommand(new CheckDynamicConfigurationCommand());
     bootstrap.addCommand(new SetUserDiscoverabilityCommand());
+    bootstrap.addCommand(new ReserveUsernameCommand());
 
     bootstrap.addBundle(new NameableMigrationsBundle<WhisperServerConfiguration>("accountdb", "accountsdb.xml") {
       @Override
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/ReserveUsernameCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/ReserveUsernameCommand.java
new file mode 100644
index 000000000..e700a3956
--- /dev/null
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/ReserveUsernameCommand.java
@@ -0,0 +1,64 @@
+/*
+ * Copyright 2013-2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.workers;
+
+import io.dropwizard.cli.ConfiguredCommand;
+import io.dropwizard.setup.Bootstrap;
+import net.sourceforge.argparse4j.inf.Namespace;
+import net.sourceforge.argparse4j.inf.Subparser;
+import org.whispersystems.textsecuregcm.WhisperServerConfiguration;
+import org.whispersystems.textsecuregcm.storage.ReservedUsernames;
+import org.whispersystems.textsecuregcm.util.DynamoDbFromConfig;
+import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
+import java.util.UUID;
+import java.util.regex.Pattern;
+
+public class ReserveUsernameCommand extends ConfiguredCommand<WhisperServerConfiguration> {
+
+  public ReserveUsernameCommand() {
+    super("reserve-username", "Reserve a username pattern for a specific account identifier");
+  }
+
+  @Override
+  public void configure(final Subparser subparser) {
+    super.configure(subparser);
+
+    subparser.addArgument("-p", "--pattern")
+        .dest("pattern")
+        .type(String.class)
+        .required(true);
+
+    subparser.addArgument("-u", "--uuid")
+        .dest("uuid")
+        .type(String.class)
+        .required(true);
+  }
+
+  @Override
+  protected void run(final Bootstrap<WhisperServerConfiguration> bootstrap, final Namespace namespace,
+      final WhisperServerConfiguration config) throws Exception {
+
+    final DynamoDbClient reservedUsernamesDynamoDbClient = DynamoDbFromConfig.client(config.getReservedUsernamesDynamoDbConfiguration(),
+        software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
+
+    final ReservedUsernames reservedUsernames = new ReservedUsernames(reservedUsernamesDynamoDbClient,
+        config.getReservedUsernamesDynamoDbConfiguration().getTableName());
+
+    final String pattern = namespace.getString("pattern").trim();
+
+    try {
+      Pattern.compile(pattern);
+    } catch (final Exception e) {
+      throw new IllegalArgumentException("Bad pattern: " + pattern, e);
+    }
+
+    final UUID aci = UUID.fromString(namespace.getString("uuid").trim());
+
+    reservedUsernames.reserveUsername(pattern, aci);
+
+    System.out.format("Reserved %s for account %s\n", pattern, aci);
+  }
+}