Additional limits
This commit is contained in:
parent
d4c4220299
commit
ac1153c7cf
|
@ -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;
|
||||
}
|
||||
|
|
|
@ -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());
|
||||
|
||||
|
|
|
@ -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;
|
||||
}
|
||||
|
|
|
@ -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");
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue