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

@@ -1,6 +1,24 @@
package org.whispersystems.textsecuregcm.websocket;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import com.google.protobuf.InvalidProtocolBufferException;
import java.time.Duration;
import java.util.Optional;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;
@@ -9,23 +27,20 @@ import org.whispersystems.textsecuregcm.push.ProvisioningManager;
import org.whispersystems.websocket.WebSocketClient;
import org.whispersystems.websocket.session.WebSocketSessionContext;
import java.util.Optional;
import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
class ProvisioningConnectListenerTest {
private ProvisioningManager provisioningManager;
private ProvisioningConnectListener provisioningConnectListener;
private ScheduledExecutorService scheduledExecutorService;
private static Duration TIMEOUT = Duration.ofSeconds(5);
@BeforeEach
void setUp() {
provisioningManager = mock(ProvisioningManager.class);
provisioningConnectListener = new ProvisioningConnectListener(provisioningManager);
scheduledExecutorService = mock(ScheduledExecutorService.class);
provisioningConnectListener =
new ProvisioningConnectListener(provisioningManager, scheduledExecutorService, TIMEOUT);
}
@Test
@@ -60,4 +75,49 @@ class ProvisioningConnectListenerTest {
assertEquals(addListenerProvisioningAddressCaptor.getValue(), removeListenerProvisioningAddressCaptor.getValue());
assertEquals(addListenerProvisioningAddressCaptor.getValue(), sentProvisioningAddress);
}
@Test
void schedulesTimeout() {
final WebSocketClient webSocketClient = mock(WebSocketClient.class);
final WebSocketSessionContext context = new WebSocketSessionContext(webSocketClient);
when(webSocketClient.supportsProvisioningSocketTimeouts()).thenReturn(true);
final ScheduledFuture<?> scheduledFuture = mock(ScheduledFuture.class);
doReturn(scheduledFuture).when(scheduledExecutorService).schedule(any(Runnable.class), anyLong(), any());
final ArgumentCaptor<Runnable> scheduleCaptor = ArgumentCaptor.forClass(Runnable.class);
provisioningConnectListener.onWebSocketConnect(context);
verify(scheduledExecutorService).schedule(scheduleCaptor.capture(), eq(TIMEOUT.getSeconds()), eq(TimeUnit.SECONDS));
verify(webSocketClient, never()).close(anyInt(), any());
scheduleCaptor.getValue().run();
verify(webSocketClient, times(1)).close(eq(1000), anyString());
}
@Test
void cancelsTimeout() {
final WebSocketClient webSocketClient = mock(WebSocketClient.class);
final WebSocketSessionContext context = new WebSocketSessionContext(webSocketClient);
when(webSocketClient.supportsProvisioningSocketTimeouts()).thenReturn(true);
final ScheduledFuture<?> scheduledFuture = mock(ScheduledFuture.class);
doReturn(scheduledFuture).when(scheduledExecutorService).schedule(any(Runnable.class), anyLong(), any());
provisioningConnectListener.onWebSocketConnect(context);
verify(scheduledExecutorService).schedule(any(Runnable.class), eq(TIMEOUT.getSeconds()), eq(TimeUnit.SECONDS));
context.notifyClosed(1000, "Test");
verify(scheduledFuture).cancel(false);
verify(webSocketClient, never()).close(anyInt(), any());
}
@Test
void skipsTimeoutIfUnsupported() {
final WebSocketClient webSocketClient = mock(WebSocketClient.class);
final WebSocketSessionContext context = new WebSocketSessionContext(webSocketClient);
provisioningConnectListener.onWebSocketConnect(context);
verify(scheduledExecutorService, never())
.schedule(any(Runnable.class), eq(TIMEOUT.getSeconds()), eq(TimeUnit.SECONDS));
}
}