Add opt-in timeouts to provisioning websocket

This commit is contained in:
ravi-signal
2024-12-18 18:45:53 -06:00
committed by GitHub
parent 6460327372
commit 68f27be7cd
10 changed files with 310 additions and 37 deletions

View File

@@ -569,6 +569,8 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
.scheduledExecutorService(name(getClass(), "cloudflareTurnRetry-%d")).threads(1).build();
ScheduledExecutorService messagePollExecutor = environment.lifecycle()
.scheduledExecutorService(name(getClass(), "messagePollExecutor-%d")).threads(1).build();
ScheduledExecutorService provisioningWebsocketTimeoutExecutor = environment.lifecycle()
.scheduledExecutorService(name(getClass(), "provisioningWebsocketTimeout-%d")).threads(1).build();
final ManagedNioEventLoopGroup dnsResolutionEventLoopGroup = new ManagedNioEventLoopGroup();
final DnsNameResolver cloudflareDnsResolver = new DnsNameResolverBuilder(dnsResolutionEventLoopGroup.next())
@@ -1171,7 +1173,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
webSocketEnvironment.getRequestLog(), Duration.ofMillis(60000));
provisioningEnvironment.jersey().register(new WebsocketRefreshApplicationEventListener(accountsManager,
disconnectionRequestManager));
provisioningEnvironment.setConnectListener(new ProvisioningConnectListener(provisioningManager));
provisioningEnvironment.setConnectListener(new ProvisioningConnectListener(provisioningManager, provisioningWebsocketTimeoutExecutor, Duration.ofSeconds(90)));
provisioningEnvironment.jersey().register(new MetricsApplicationEventListener(TrafficSource.WEBSOCKET, clientReleaseManager));
provisioningEnvironment.jersey().register(new KeepAliveController(webSocketConnectionEventManager));
provisioningEnvironment.jersey().register(new TimestampResponseFilter());

View File

@@ -130,7 +130,7 @@ import org.whispersystems.textsecuregcm.util.ExceptionUtils;
import org.whispersystems.textsecuregcm.util.HeaderUtils;
import org.whispersystems.textsecuregcm.util.Util;
import org.whispersystems.textsecuregcm.websocket.WebSocketConnection;
import org.whispersystems.websocket.Stories;
import org.whispersystems.websocket.WebsocketHeaders;
import org.whispersystems.websocket.auth.ReadOnly;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
@@ -749,10 +749,10 @@ public class MessageController {
@GET
@Produces(MediaType.APPLICATION_JSON)
public CompletableFuture<OutgoingMessageEntityList> getPendingMessages(@ReadOnly @Auth AuthenticatedDevice auth,
@HeaderParam(Stories.X_SIGNAL_RECEIVE_STORIES) String receiveStoriesHeader,
@HeaderParam(WebsocketHeaders.X_SIGNAL_RECEIVE_STORIES) String receiveStoriesHeader,
@HeaderParam(HttpHeaders.USER_AGENT) String userAgent) {
boolean shouldReceiveStories = Stories.parseReceiveStoriesHeader(receiveStoriesHeader);
boolean shouldReceiveStories = WebsocketHeaders.parseReceiveStoriesHeader(receiveStoriesHeader);
pushNotificationManager.handleMessagesRetrieved(auth.getAccount(), auth.getAuthenticatedDevice(), userAgent);

View File

@@ -7,10 +7,15 @@ package org.whispersystems.textsecuregcm.websocket;
import com.google.common.annotations.VisibleForTesting;
import java.security.SecureRandom;
import java.time.Duration;
import java.util.Base64;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import org.whispersystems.textsecuregcm.auth.AuthenticatedDevice;
import org.whispersystems.textsecuregcm.controllers.ProvisioningController;
import org.whispersystems.textsecuregcm.entities.MessageProtos;
import org.whispersystems.textsecuregcm.entities.ProvisioningMessage;
import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
@@ -31,7 +36,7 @@ import org.whispersystems.websocket.setup.WebSocketConnectListener;
* a random, temporary "provisioning address," which it transmits via the newly-opened WebSocket. From there, the new
* device generally displays the provisioning address (and a public key) as a QR code. After that, the primary device
* will scan the QR code and send an encrypted provisioning message to the new device via
* {@link org.whispersystems.textsecuregcm.controllers.ProvisioningController#sendProvisioningMessage(AuthenticatedDevice, String, ProvisioningMessage, String)}.
* {@link ProvisioningController#sendProvisioningMessage(AuthenticatedDevice, String, ProvisioningMessage, String)}.
* Once the server receives the message from the primary device, it sends the message to the new device via the open
* WebSocket, then closes the WebSocket connection.
*/
@@ -39,9 +44,15 @@ public class ProvisioningConnectListener implements WebSocketConnectListener {
private final ProvisioningManager provisioningManager;
private final OpenWebSocketCounter openWebSocketCounter;
private final ScheduledExecutorService timeoutExecutor;
private final Duration timeout;
public ProvisioningConnectListener(final ProvisioningManager provisioningManager) {
public ProvisioningConnectListener(final ProvisioningManager provisioningManager,
final ScheduledExecutorService timeoutExecutor,
final Duration timeout) {
this.provisioningManager = provisioningManager;
this.timeoutExecutor = timeoutExecutor;
this.timeout = timeout;
this.openWebSocketCounter = new OpenWebSocketCounter(MetricsUtil.name(getClass(), "openWebsockets"),
MetricsUtil.name(getClass(), "sessionDuration"));
}
@@ -50,8 +61,17 @@ public class ProvisioningConnectListener implements WebSocketConnectListener {
public void onWebSocketConnect(WebSocketSessionContext context) {
openWebSocketCounter.countOpenWebSocket(context);
final Optional<ScheduledFuture<?>> maybeTimeoutFuture = context.getClient().supportsProvisioningSocketTimeouts()
? Optional.of(timeoutExecutor.schedule(() ->
context.getClient().close(1000, "Timeout"), timeout.toSeconds(), TimeUnit.SECONDS))
: Optional.empty();
final String provisioningAddress = generateProvisioningAddress();
context.addWebsocketClosedListener((context1, statusCode, reason) -> provisioningManager.removeListener(provisioningAddress));
context.addWebsocketClosedListener((context1, statusCode, reason) -> {
provisioningManager.removeListener(provisioningAddress);
maybeTimeoutFuture.ifPresent(future -> future.cancel(false));
});
provisioningManager.addListener(provisioningAddress, message -> {
assert message.getType() == PubSubProtos.PubSubMessage.Type.DELIVER;