From bbde93a3c74ba2811d1771d998a9bfe2ea257eb0 Mon Sep 17 00:00:00 2001
From: Ehren Kret <ehren@signal.org>
Date: Wed, 13 Oct 2021 11:46:20 -0500
Subject: [PATCH] Enable unwrapping of CompletionStage

---
 .../textsecuregcm/WhisperServerService.java   |  7 +-
 .../SubscriptionControllerTest.java           | 97 +++++++++++++++++++
 2 files changed, 102 insertions(+), 2 deletions(-)
 create mode 100644 service/src/test/java/org/whispersystems/textsecuregcm/controllers/SubscriptionControllerTest.java

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index f1fb4bde3..a9c1415e9 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -53,6 +53,7 @@ import javax.servlet.DispatcherType;
 import javax.servlet.FilterRegistration;
 import javax.servlet.ServletRegistration;
 import org.eclipse.jetty.servlets.CrossOriginFilter;
+import org.glassfish.jersey.server.ServerProperties;
 import org.jdbi.v3.core.Jdbi;
 import org.signal.zkgroup.ServerSecretParams;
 import org.signal.zkgroup.auth.ServerZkAuthOperations;
@@ -675,11 +676,13 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     provisioning.addMapping("/v1/websocket/provisioning/");
     provisioning.setAsyncSupported(true);
 
+    environment.jersey().property(ServerProperties.UNWRAP_COMPLETION_STAGE_IN_WRITER_ENABLE, Boolean.TRUE);
+    webSocketEnvironment.jersey().property(ServerProperties.UNWRAP_COMPLETION_STAGE_IN_WRITER_ENABLE, Boolean.TRUE);
+    provisioningEnvironment.jersey().property(ServerProperties.UNWRAP_COMPLETION_STAGE_IN_WRITER_ENABLE, Boolean.TRUE);
+
     environment.admin().addTask(new SetRequestLoggingEnabledTask());
     environment.admin().addTask(new SetCrawlerAccelerationTask(accountDatabaseCrawlerCache));
 
-///
-
     environment.healthChecks().register("cacheCluster", new RedisClusterHealthCheck(cacheCluster));
 
     environment.metrics().register(name(CpuUsageGauge.class, "cpu"), new CpuUsageGauge(3, TimeUnit.SECONDS));
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/SubscriptionControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/SubscriptionControllerTest.java
new file mode 100644
index 000000000..0ec90e3f7
--- /dev/null
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/SubscriptionControllerTest.java
@@ -0,0 +1,97 @@
+/*
+ * Copyright 2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.controllers;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.when;
+
+import io.dropwizard.auth.PolymorphicAuthValueFactoryProvider;
+import io.dropwizard.testing.junit5.DropwizardExtensionsSupport;
+import io.dropwizard.testing.junit5.ResourceExtension;
+import java.math.BigDecimal;
+import java.time.Clock;
+import java.util.Map;
+import java.util.Set;
+import org.glassfish.jersey.server.ServerProperties;
+import org.glassfish.jersey.test.grizzly.GrizzlyWebTestContainerFactory;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.signal.zkgroup.receipts.ServerZkReceiptOperations;
+import org.whispersystems.textsecuregcm.auth.AuthenticatedAccount;
+import org.whispersystems.textsecuregcm.auth.DisabledPermittedAuthenticatedAccount;
+import org.whispersystems.textsecuregcm.configuration.SubscriptionConfiguration;
+import org.whispersystems.textsecuregcm.configuration.SubscriptionLevelConfiguration;
+import org.whispersystems.textsecuregcm.configuration.SubscriptionPriceConfiguration;
+import org.whispersystems.textsecuregcm.controllers.SubscriptionController.GetLevelsResponse;
+import org.whispersystems.textsecuregcm.storage.IssuedReceiptsManager;
+import org.whispersystems.textsecuregcm.storage.SubscriptionManager;
+import org.whispersystems.textsecuregcm.stripe.StripeManager;
+import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
+import org.whispersystems.textsecuregcm.util.SystemMapper;
+
+@ExtendWith(DropwizardExtensionsSupport.class)
+class SubscriptionControllerTest {
+
+  private static final Clock CLOCK = mock(Clock.class);
+  private static final SubscriptionConfiguration SUBSCRIPTION_CONFIG = mock(SubscriptionConfiguration.class);
+  private static final SubscriptionManager SUBSCRIPTION_MANAGER = mock(SubscriptionManager.class);
+  private static final StripeManager STRIPE_MANAGER = mock(StripeManager.class);
+  private static final ServerZkReceiptOperations ZK_OPS = mock(ServerZkReceiptOperations.class);
+  private static final IssuedReceiptsManager ISSUED_RECEIPTS_MANAGER = mock(IssuedReceiptsManager.class);
+  private static final SubscriptionController SUBSCRIPTION_CONTROLLER = new SubscriptionController(
+      CLOCK, SUBSCRIPTION_CONFIG, SUBSCRIPTION_MANAGER, STRIPE_MANAGER, ZK_OPS, ISSUED_RECEIPTS_MANAGER);
+  private static final ResourceExtension RESOURCE_EXTENSION = ResourceExtension.builder()
+      .addProperty(ServerProperties.UNWRAP_COMPLETION_STAGE_IN_WRITER_ENABLE, Boolean.TRUE)
+      .addProvider(AuthHelper.getAuthFilter())
+      .addProvider(new PolymorphicAuthValueFactoryProvider.Binder<>(Set.of(
+          AuthenticatedAccount.class, DisabledPermittedAuthenticatedAccount.class)))
+      .setMapper(SystemMapper.getMapper())
+      .setTestContainerFactory(new GrizzlyWebTestContainerFactory())
+      .addResource(SUBSCRIPTION_CONTROLLER)
+      .build();
+
+  @AfterEach
+  void tearDown() {
+    reset(CLOCK, SUBSCRIPTION_CONFIG, SUBSCRIPTION_MANAGER, STRIPE_MANAGER, ZK_OPS, ISSUED_RECEIPTS_MANAGER);
+  }
+
+  @Test
+  void getLevels() {
+    when(SUBSCRIPTION_CONFIG.getLevels()).thenReturn(Map.of(
+        1L, new SubscriptionLevelConfiguration("B1", "P1", Map.of("USD", new SubscriptionPriceConfiguration("R1", BigDecimal.valueOf(100)))),
+        2L, new SubscriptionLevelConfiguration("B2", "P2", Map.of("USD", new SubscriptionPriceConfiguration("R2", BigDecimal.valueOf(200)))),
+        3L, new SubscriptionLevelConfiguration("B3", "P3", Map.of("USD", new SubscriptionPriceConfiguration("R3", BigDecimal.valueOf(300))))
+    ));
+
+    GetLevelsResponse response = RESOURCE_EXTENSION.target("/v1/subscription/levels")
+        .request()
+        .get(GetLevelsResponse.class);
+
+    assertThat(response.getLevels()).containsKeys(1L, 2L, 3L).satisfies(longLevelMap -> {
+      assertThat(longLevelMap).extractingByKey(1L).satisfies(level -> {
+        assertThat(level.getBadgeId()).isEqualTo("B1");
+        assertThat(level.getCurrencies()).containsKeys("USD").extractingByKey("USD").satisfies(price -> {
+          assertThat(price.getAmount()).isEqualTo("100");
+        });
+      });
+      assertThat(longLevelMap).extractingByKey(2L).satisfies(level -> {
+        assertThat(level.getBadgeId()).isEqualTo("B2");
+        assertThat(level.getCurrencies()).containsKeys("USD").extractingByKey("USD").satisfies(price -> {
+          assertThat(price.getAmount()).isEqualTo("200");
+        });
+      });
+      assertThat(longLevelMap).extractingByKey(3L).satisfies(level -> {
+        assertThat(level.getBadgeId()).isEqualTo("B3");
+        assertThat(level.getCurrencies()).containsKeys("USD").extractingByKey("USD").satisfies(price -> {
+          assertThat(price.getAmount()).isEqualTo("300");
+        });
+      });
+    });
+  }
+}