Add support for UUID buckets in remote config

This commit is contained in:
Moxie Marlinspike
2020-01-21 14:38:47 -08:00
parent 08a70664f4
commit e4e20c2d25
8 changed files with 135 additions and 31 deletions

View File

@@ -21,9 +21,12 @@ import javax.ws.rs.Produces;
import javax.ws.rs.WebApplicationException;
import javax.ws.rs.core.MediaType;
import javax.ws.rs.core.Response;
import java.nio.ByteBuffer;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.List;
import java.util.Set;
import java.util.UUID;
import java.util.stream.Collectors;
import io.dropwizard.auth.Auth;
@@ -46,12 +49,12 @@ public class RemoteConfigController {
public UserRemoteConfigList getAll(@Auth Account account) {
try {
MessageDigest digest = MessageDigest.getInstance("SHA1");
byte[] number = account.getNumber().getBytes();
return new UserRemoteConfigList(remoteConfigsManager.getAll().stream().map(config -> new UserRemoteConfig(config.getName(),
isInBucket(digest, number,
isInBucket(digest, account.getUuid(),
config.getName().getBytes(),
config.getPercentage())))
config.getPercentage(),
config.getUuids())))
.collect(Collectors.toList()));
} catch (NoSuchAlgorithmException e) {
throw new AssertionError(e);
@@ -82,8 +85,14 @@ public class RemoteConfigController {
}
@VisibleForTesting
public static boolean isInBucket(MessageDigest digest, byte[] user, byte[] configName, int configPercentage) {
digest.update(user);
public static boolean isInBucket(MessageDigest digest, UUID uid, byte[] configName, int configPercentage, Set<UUID> uuidsInBucket) {
if (uuidsInBucket.contains(uid)) return true;
ByteBuffer bb = ByteBuffer.wrap(new byte[16]);
bb.putLong(uid.getMostSignificantBits());
bb.putLong(uid.getLeastSignificantBits());
digest.update(bb.array());
byte[] hash = digest.digest(configName);
int bucket = (int)(Math.abs(Conversions.byteArrayToLong(hash)) % 100);
@@ -93,7 +102,7 @@ public class RemoteConfigController {
@SuppressWarnings("BooleanMethodIsAlwaysInverted")
private boolean isAuthorized(String configToken) {
return configAuthTokens.stream().anyMatch(authorized -> MessageDigest.isEqual(authorized.getBytes(), configToken.getBytes()));
return configToken != null && configAuthTokens.stream().anyMatch(authorized -> MessageDigest.isEqual(authorized.getBytes(), configToken.getBytes()));
}
}

View File

@@ -6,6 +6,11 @@ import javax.validation.constraints.Max;
import javax.validation.constraints.Min;
import javax.validation.constraints.NotNull;
import javax.validation.constraints.Pattern;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;
import java.util.UUID;
public class RemoteConfig {
@@ -19,11 +24,16 @@ public class RemoteConfig {
@Max(100)
private int percentage;
@JsonProperty
@NotNull
private Set<UUID> uuids = new HashSet<>();
public RemoteConfig() {}
public RemoteConfig(String name, int percentage) {
public RemoteConfig(String name, int percentage, Set<UUID> uuids) {
this.name = name;
this.percentage = percentage;
this.uuids = uuids;
}
public int getPercentage() {
@@ -33,4 +43,8 @@ public class RemoteConfig {
public String getName() {
return name;
}
public Set<UUID> getUuids() {
return uuids;
}
}

View File

@@ -11,6 +11,7 @@ import org.whispersystems.textsecuregcm.storage.mappers.RemoteConfigRowMapper;
import org.whispersystems.textsecuregcm.util.Constants;
import org.whispersystems.textsecuregcm.util.SystemMapper;
import java.util.LinkedList;
import java.util.List;
import java.util.Optional;
import java.util.UUID;
@@ -22,8 +23,7 @@ public class RemoteConfigs {
public static final String ID = "id";
public static final String NAME = "name";
public static final String PERCENTAGE = "percentage";
private static final ObjectMapper mapper = SystemMapper.getMapper();
public static final String UUIDS = "uuids";
private final MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
private final Timer setTimer = metricRegistry.timer(name(Accounts.class, "set" ));
@@ -35,14 +35,16 @@ public class RemoteConfigs {
public RemoteConfigs(FaultTolerantDatabase database) {
this.database = database;
this.database.getDatabase().registerRowMapper(new RemoteConfigRowMapper());
this.database.getDatabase().registerArrayType(UUID.class, "uuid");
}
public void set(RemoteConfig remoteConfig) {
database.use(jdbi -> jdbi.useHandle(handle -> {
try (Timer.Context ignored = setTimer.time()) {
handle.createUpdate("INSERT INTO remote_config (" + NAME + ", " + PERCENTAGE + ") VALUES (:name, :percentage) ON CONFLICT(" + NAME + ") DO UPDATE SET " + PERCENTAGE + " = EXCLUDED." + PERCENTAGE)
handle.createUpdate("INSERT INTO remote_config (" + NAME + ", " + PERCENTAGE + ", " + UUIDS + ") VALUES (:name, :percentage, :uuids) ON CONFLICT(" + NAME + ") DO UPDATE SET " + PERCENTAGE + " = EXCLUDED." + PERCENTAGE + ", " + UUIDS + " = EXCLUDED." + UUIDS)
.bind("name", remoteConfig.getName())
.bind("percentage", remoteConfig.getPercentage())
.bind("uuids", remoteConfig.getUuids().toArray(new UUID[0]))
.execute();
}
}));

View File

@@ -7,12 +7,15 @@ import org.whispersystems.textsecuregcm.storage.RemoteConfigs;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Arrays;
import java.util.HashSet;
import java.util.UUID;
public class RemoteConfigRowMapper implements RowMapper<RemoteConfig> {
@Override
public RemoteConfig map(ResultSet rs, StatementContext ctx) throws SQLException {
return new RemoteConfig(rs.getString(RemoteConfigs.NAME), rs.getInt(RemoteConfigs.PERCENTAGE));
return new RemoteConfig(rs.getString(RemoteConfigs.NAME), rs.getInt(RemoteConfigs.PERCENTAGE), new HashSet<>(Arrays.asList((UUID[])rs.getArray(RemoteConfigs.UUIDS).getArray())));
}
}