Add second (migration) database to AbusiveHostRules

This commit is contained in:
Chris Eager
2021-12-08 12:46:05 -08:00
committed by GitHub
parent 9a5ffea0ad
commit a70b057e1c
9 changed files with 450 additions and 154 deletions

View File

@@ -6,39 +6,63 @@
package org.whispersystems.textsecuregcm.tests.storage;
import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import io.zonky.test.db.postgres.embedded.LiquibasePreparer;
import io.zonky.test.db.postgres.junit.EmbeddedPostgresRules;
import io.zonky.test.db.postgres.junit.PreparedDbRule;
import io.zonky.test.db.postgres.junit5.EmbeddedPostgresExtension;
import io.zonky.test.db.postgres.junit5.PreparedDbExtension;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Arrays;
import java.util.List;
import org.jdbi.v3.core.Jdbi;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicAbusiveHostRulesMigrationConfiguration;
import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
import org.whispersystems.textsecuregcm.storage.AbusiveHostRule;
import org.whispersystems.textsecuregcm.storage.AbusiveHostRules;
import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase;
public class AbusiveHostRulesTest {
class AbusiveHostRulesTest {
@Rule
public PreparedDbRule db = EmbeddedPostgresRules.preparedDatabase(LiquibasePreparer.forClasspathLocation("abusedb.xml"));
@RegisterExtension
PreparedDbExtension db = EmbeddedPostgresExtension.preparedDatabase(
LiquibasePreparer.forClasspathLocation("abusedb.xml"));
@RegisterExtension
PreparedDbExtension newDb = EmbeddedPostgresExtension.preparedDatabase(
LiquibasePreparer.forClasspathLocation("abusedb.xml"));
private AbusiveHostRules abusiveHostRules;
@Before
public void setup() {
this.abusiveHostRules = new AbusiveHostRules(new FaultTolerantDatabase("abusive_hosts-test", Jdbi.create(db.getTestDatabase()), new CircuitBreakerConfiguration()));
@BeforeEach
void setup() {
//noinspection unchecked
final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager = mock(
DynamicConfigurationManager.class);
final DynamicConfiguration dynamicConfiguration = mock(DynamicConfiguration.class);
when(dynamicConfigurationManager.getConfiguration()).thenReturn(dynamicConfiguration);
when(dynamicConfiguration.getAbusiveHostRulesMigrationConfiguration()).thenReturn(
new DynamicAbusiveHostRulesMigrationConfiguration());
this.abusiveHostRules = new AbusiveHostRules(
new FaultTolerantDatabase("abusive_hosts-test", Jdbi.create(db.getTestDatabase()),
new CircuitBreakerConfiguration()),
new FaultTolerantDatabase("abusive_hosts-test", Jdbi.create(newDb.getTestDatabase()),
new CircuitBreakerConfiguration()),
dynamicConfigurationManager);
}
@Test
public void testBlockedHost() throws SQLException {
PreparedStatement statement = db.getTestDatabase().getConnection().prepareStatement("INSERT INTO abusive_host_rules (host, blocked) VALUES (?::INET, ?)");
void testBlockedHost() throws SQLException {
PreparedStatement statement = db.getTestDatabase().getConnection()
.prepareStatement("INSERT INTO abusive_host_rules (host, blocked) VALUES (?::INET, ?)");
statement.setString(1, "192.168.1.1");
statement.setInt(2, 1);
statement.execute();
@@ -51,8 +75,9 @@ public class AbusiveHostRulesTest {
}
@Test
public void testBlockedCidr() throws SQLException {
PreparedStatement statement = db.getTestDatabase().getConnection().prepareStatement("INSERT INTO abusive_host_rules (host, blocked) VALUES (?::INET, ?)");
void testBlockedCidr() throws SQLException {
PreparedStatement statement = db.getTestDatabase().getConnection()
.prepareStatement("INSERT INTO abusive_host_rules (host, blocked) VALUES (?::INET, ?)");
statement.setString(1, "192.168.1.0/24");
statement.setInt(2, 1);
statement.execute();
@@ -65,8 +90,9 @@ public class AbusiveHostRulesTest {
}
@Test
public void testUnblocked() throws SQLException {
PreparedStatement statement = db.getTestDatabase().getConnection().prepareStatement("INSERT INTO abusive_host_rules (host, blocked) VALUES (?::INET, ?)");
void testUnblocked() throws SQLException {
PreparedStatement statement = db.getTestDatabase().getConnection()
.prepareStatement("INSERT INTO abusive_host_rules (host, blocked) VALUES (?::INET, ?)");
statement.setString(1, "192.168.1.0/24");
statement.setInt(2, 1);
statement.execute();
@@ -76,8 +102,9 @@ public class AbusiveHostRulesTest {
}
@Test
public void testRestricted() throws SQLException {
PreparedStatement statement = db.getTestDatabase().getConnection().prepareStatement("INSERT INTO abusive_host_rules (host, blocked, regions) VALUES (?::INET, ?, ?)");
void testRestricted() throws SQLException {
PreparedStatement statement = db.getTestDatabase().getConnection()
.prepareStatement("INSERT INTO abusive_host_rules (host, blocked, regions) VALUES (?::INET, ?, ?)");
statement.setString(1, "192.168.1.0/24");
statement.setInt(2, 0);
statement.setString(3, "+1,+49");
@@ -90,10 +117,11 @@ public class AbusiveHostRulesTest {
}
@Test
public void testInsertBlocked() throws Exception {
void testInsertBlocked() throws Exception {
abusiveHostRules.setBlockedHost("172.17.0.1", "Testing one two");
PreparedStatement statement = db.getTestDatabase().getConnection().prepareStatement("SELECT * from abusive_host_rules WHERE host = ?::inet");
PreparedStatement statement = db.getTestDatabase().getConnection()
.prepareStatement("SELECT * from abusive_host_rules WHERE host = ?::inet");
statement.setString(1, "172.17.0.1");
ResultSet resultSet = statement.executeQuery();
@@ -106,8 +134,8 @@ public class AbusiveHostRulesTest {
abusiveHostRules.setBlockedHost("172.17.0.1", "Different notes");
statement = db.getTestDatabase().getConnection().prepareStatement("SELECT * from abusive_host_rules WHERE host = ?::inet");
statement = db.getTestDatabase().getConnection()
.prepareStatement("SELECT * from abusive_host_rules WHERE host = ?::inet");
statement.setString(1, "172.17.0.1");
resultSet = statement.executeQuery();
@@ -119,4 +147,30 @@ public class AbusiveHostRulesTest {
assertThat(resultSet.getString("notes")).isEqualTo("Testing one two");
}
@Test
void testMigrate() throws Exception {
final int rules = 20;
for (int i = 1; i <= rules; i++) {
abusiveHostRules.setBlockedHost("172.17.0." + i, "Testing one two " + i);
}
PreparedStatement statement = newDb.getTestDatabase().getConnection()
.prepareStatement("SELECT * from abusive_host_rules");
assertThat(queryResultSize(statement.executeQuery())).isEqualTo(0);
abusiveHostRules.forEachInOldDatabase((rule, host) -> abusiveHostRules.migrateAbusiveHostRule(rule, host), 5);
assertThat(queryResultSize(statement.executeQuery())).isEqualTo(rules);
}
private int queryResultSize(ResultSet resultSet) throws SQLException {
int migrated = 0;
while (resultSet.next()) {
migrated++;
}
return migrated;
}
}