From ac1153c7cf76a568e8c07bbe84e08338780c044b Mon Sep 17 00:00:00 2001 From: Moxie Marlinspike Date: Sat, 14 Mar 2020 17:59:37 -0700 Subject: [PATCH] Additional limits --- .../configuration/RateLimitsConfiguration.java | 7 +++++++ .../controllers/DirectoryController.java | 12 +++++++++++- .../textsecuregcm/limits/RateLimiters.java | 9 +++++++++ .../tests/controllers/DirectoryControllerTest.java | 14 +++++++++----- 4 files changed, 36 insertions(+), 6 deletions(-) diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/RateLimitsConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/RateLimitsConfiguration.java index c259fb5ae..75866c41f 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/RateLimitsConfiguration.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/RateLimitsConfiguration.java @@ -50,6 +50,9 @@ public class RateLimitsConfiguration { @JsonProperty private RateLimitConfiguration contactQueries = new RateLimitConfiguration(50000, 50000); + @JsonProperty + private RateLimitConfiguration contactIpQueries = new RateLimitConfiguration(200, (100.0 / 60.0)); + @JsonProperty private RateLimitConfiguration prekeys = new RateLimitConfiguration(3, 1.0 / 10.0); @@ -101,6 +104,10 @@ public class RateLimitsConfiguration { return contactQueries; } + public RateLimitConfiguration getContactIpQueries() { + return contactIpQueries; + } + public RateLimitConfiguration getAttachments() { return attachments; } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/DirectoryController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/DirectoryController.java index e37642fa5..fb92c3a9c 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/DirectoryController.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/DirectoryController.java @@ -38,6 +38,7 @@ import org.whispersystems.textsecuregcm.util.Constants; import javax.validation.Valid; import javax.ws.rs.Consumes; import javax.ws.rs.GET; +import javax.ws.rs.HeaderParam; import javax.ws.rs.PUT; import javax.ws.rs.Path; import javax.ws.rs.PathParam; @@ -47,6 +48,7 @@ import javax.ws.rs.core.MediaType; import javax.ws.rs.core.Response; import javax.ws.rs.core.Response.Status; import java.io.IOException; +import java.util.Arrays; import java.util.HashMap; import java.util.LinkedList; import java.util.List; @@ -163,9 +165,17 @@ public class DirectoryController { @Path("/tokens") @Produces(MediaType.APPLICATION_JSON) @Consumes(MediaType.APPLICATION_JSON) - public ClientContacts getContactIntersection(@Auth Account account, @Valid ClientContactTokens contacts) + public ClientContacts getContactIntersection(@Auth Account account, + @HeaderParam("X-Forwarded-For") String forwardedFor, + @Valid ClientContactTokens contacts) throws RateLimitExceededException { + String requester = Arrays.stream(forwardedFor.split(",")) + .map(String::trim) + .reduce((a, b) -> b) + .orElseThrow(); + + rateLimiters.getContactsIpLimiter().validate(requester); rateLimiters.getContactsLimiter().validate(account.getNumber(), contacts.getContacts().size()); contactsHistogram.update(contacts.getContacts().size()); diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiters.java b/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiters.java index cd22cc82f..f58db72b4 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiters.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiters.java @@ -33,6 +33,7 @@ public class RateLimiters { private final RateLimiter attachmentLimiter; private final RateLimiter contactsLimiter; + private final RateLimiter contactsIpLimiter; private final RateLimiter preKeysLimiter; private final RateLimiter messagesLimiter; @@ -87,6 +88,10 @@ public class RateLimiters { config.getContactQueries().getBucketSize(), config.getContactQueries().getLeakRatePerMinute()); + this.contactsIpLimiter = new RateLimiter(cacheClient, "contactsIpQuery", + config.getContactIpQueries().getBucketSize(), + config.getContactIpQueries().getLeakRatePerMinute()); + this.preKeysLimiter = new RateLimiter(cacheClient, "prekeys", config.getPreKeys().getBucketSize(), config.getPreKeys().getLeakRatePerMinute()); @@ -144,6 +149,10 @@ public class RateLimiters { return contactsLimiter; } + public RateLimiter getContactsIpLimiter() { + return contactsIpLimiter; + } + public RateLimiter getAttachmentLimiter() { return this.attachmentLimiter; } diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/DirectoryControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/DirectoryControllerTest.java index 033587de6..8cbcb282e 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/DirectoryControllerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/DirectoryControllerTest.java @@ -34,14 +34,14 @@ import io.dropwizard.testing.junit.ResourceTestRule; import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.ArgumentMatchers.anyListOf; import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; +import static org.mockito.Mockito.*; public class DirectoryControllerTest { - private final RateLimiters rateLimiters = mock(RateLimiters.class); - private final RateLimiter rateLimiter = mock(RateLimiter.class); - private final DirectoryManager directoryManager = mock(DirectoryManager.class); + private final RateLimiters rateLimiters = mock(RateLimiters.class ); + private final RateLimiter rateLimiter = mock(RateLimiter.class ); + private final RateLimiter ipLimiter = mock(RateLimiter.class ); + private final DirectoryManager directoryManager = mock(DirectoryManager.class ); private final ExternalServiceCredentialGenerator directoryCredentialsGenerator = mock(ExternalServiceCredentialGenerator.class); private final ExternalServiceCredentials validCredentials = new ExternalServiceCredentials("username", "password"); @@ -60,6 +60,7 @@ public class DirectoryControllerTest { @Before public void setup() throws Exception { when(rateLimiters.getContactsLimiter()).thenReturn(rateLimiter); + when(rateLimiters.getContactsIpLimiter()).thenReturn(ipLimiter); when(directoryManager.get(anyListOf(byte[].class))).thenAnswer(new Answer>() { @Override public List answer(InvocationOnMock invocationOnMock) throws Throwable { @@ -180,10 +181,13 @@ public class DirectoryControllerTest { .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_NUMBER, AuthHelper.VALID_PASSWORD)) + .header("X-Forwarded-For", "192.168.1.1, 1.1.1.1") .put(Entity.entity(new ClientContactTokens(tokens), MediaType.APPLICATION_JSON_TYPE)); assertThat(response.getStatus()).isEqualTo(200); assertThat(response.readEntity(ClientContactTokens.class).getContacts()).isEqualTo(expectedResponse); + + verify(ipLimiter).validate("1.1.1.1"); } }