Additional limits

This commit is contained in:
Moxie Marlinspike 2020-03-14 17:59:37 -07:00
parent d4c4220299
commit ac1153c7cf
4 changed files with 36 additions and 6 deletions

View File

@ -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;
}

View File

@ -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());

View File

@ -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;
}

View File

@ -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<List<byte[]>>() {
@Override
public List<byte[]> 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");
}
}