mirror of
https://github.com/signalapp/Signal-Server
synced 2026-04-23 06:38:03 +01:00
Add opt-in timeouts to provisioning websocket
This commit is contained in:
@@ -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));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user