Add ability to do unused reads from CDSv2 to test server load.

This commit is contained in:
Greyson Parrelli 2022-08-17 18:23:40 -04:00 committed by Cody Henthorne
parent 84717b95f7
commit 15e52a8b88
7 changed files with 144 additions and 46 deletions

View file

@ -201,7 +201,7 @@ android {
buildConfigField "String[]", "SIGNAL_CONTENT_PROXY_IPS", content_proxy_ips
buildConfigField "String", "SIGNAL_AGENT", "\"OWA\""
buildConfigField "String", "CDS_MRENCLAVE", "\"74778bb0f93ae1f78c26e67152bab0bbeb693cd56d1bb9b4e9244157acc58081\""
buildConfigField "String", "CDSI_MRENCLAVE", "\"7b75dd6e862decef9b37132d54be082441917a7790e82fe44f9cf653de03a75f\""
buildConfigField "String", "CDSI_MRENCLAVE", "\"ef4787a56a154ac6d009138cac17155acd23cfe4329281252365dd7c252e7fbf\""
buildConfigField "org.thoughtcrime.securesms.KbsEnclave", "KBS_ENCLAVE", "new org.thoughtcrime.securesms.KbsEnclave(\"0cedba03535b41b67729ce9924185f831d7767928a1d1689acb689bc079c375f\", " +
"\"187d2739d22be65e74b65f0055e74d31310e4267e5fac2b1246cc8beba81af39\", " +
"\"ee19f1965b1eefa3dc4204eb70c04f397755f771b8c1909d080c04dad2a6a9ba\")"

View file

@ -13,6 +13,7 @@ import org.signal.contacts.SystemContactsRepository.ContactIterator
import org.signal.contacts.SystemContactsRepository.ContactPhoneDetails
import org.signal.core.util.Stopwatch
import org.signal.core.util.StringUtil
import org.signal.core.util.concurrent.SignalExecutors
import org.signal.core.util.logging.Log
import org.thoughtcrime.securesms.BuildConfig
import org.thoughtcrime.securesms.R
@ -36,7 +37,11 @@ import org.thoughtcrime.securesms.util.Util
import org.whispersystems.signalservice.api.push.SignalServiceAddress
import org.whispersystems.signalservice.api.util.UuidUtil
import java.io.IOException
import java.lang.Exception
import java.util.Calendar
import java.util.concurrent.Callable
import java.util.concurrent.ExecutionException
import java.util.concurrent.Future
/**
* Methods for discovering which users are registered and marking them as such in the database.
@ -76,6 +81,8 @@ object ContactDiscovery {
refresh = {
if (FeatureFlags.phoneNumberPrivacy()) {
ContactDiscoveryRefreshV2.refreshAll(context)
} else if (FeatureFlags.cdsV2LoadTesting()) {
loadTestRefreshAll(context)
} else {
ContactDiscoveryRefreshV1.refreshAll(context)
}
@ -97,6 +104,8 @@ object ContactDiscovery {
refresh = {
if (FeatureFlags.phoneNumberPrivacy()) {
ContactDiscoveryRefreshV2.refresh(context, recipients)
} else if (FeatureFlags.cdsV2LoadTesting()) {
loadTestRefresh(context, recipients)
} else {
ContactDiscoveryRefreshV1.refresh(context, recipients)
}
@ -116,6 +125,8 @@ object ContactDiscovery {
refresh = {
if (FeatureFlags.phoneNumberPrivacy()) {
ContactDiscoveryRefreshV2.refresh(context, listOf(recipient))
} else if (FeatureFlags.cdsV2LoadTesting()) {
loadTestRefresh(context, listOf(recipient))
} else {
ContactDiscoveryRefreshV1.refresh(context, listOf(recipient))
}
@ -367,6 +378,38 @@ object ContactDiscovery {
ApplicationDependencies.getProtocolStore().pni().containsSession(protocolAddress)
}
private fun loadTestRefreshAll(context: Context): RefreshResult {
return loadTestOperation(
{ ContactDiscoveryRefreshV1.refreshAll(context) },
{ ContactDiscoveryRefreshV2.refreshAll(context, ignoreResults = true) }
)
}
private fun loadTestRefresh(context: Context, recipients: List<Recipient>): RefreshResult {
return loadTestOperation(
{ ContactDiscoveryRefreshV1.refresh(context, recipients) },
{ ContactDiscoveryRefreshV2.refresh(context, recipients, ignoreResults = true) }
)
}
private fun loadTestOperation(operationV1: Callable<RefreshResult>, operationV2: Callable<RefreshResult>): RefreshResult {
val v1Future: Future<RefreshResult> = SignalExecutors.UNBOUNDED.submit(operationV1)
val v2Future: Future<RefreshResult> = SignalExecutors.UNBOUNDED.submit(operationV2)
try {
v2Future.get()
} catch (e: Exception) {
Log.w(TAG, "Failed to complete the V2 fetch!", e)
}
try {
return v1Future.get()
} catch (e: ExecutionException) {
Log.w(TAG, "Hit exception during V1 fetch!", e)
throw e.cause!!
}
}
class RefreshResult(
val registeredIds: Set<RecipientId>,
val rewrites: Map<String, String>

View file

@ -62,7 +62,8 @@ import io.reactivex.rxjava3.schedulers.Schedulers;
*/
class ContactDiscoveryRefreshV1 {
private static final String TAG = Log.tag(ContactDiscoveryRefreshV1.class);
// Using Log.tag will cut off the version number
private static final String TAG = "CdsRefreshV1";
private static final int MAX_NUMBERS = 20_500;

View file

@ -24,7 +24,8 @@ import java.util.Optional
*/
object ContactDiscoveryRefreshV2 {
private val TAG = Log.tag(ContactDiscoveryRefreshV2::class.java)
// Using Log.tag will cut off the version number
private const val TAG = "CdsRefreshV2"
/**
* The maximum number items we will allow in a 'one-off' request.
@ -38,7 +39,7 @@ object ContactDiscoveryRefreshV2 {
@WorkerThread
@Synchronized
@JvmStatic
fun refreshAll(context: Context): ContactDiscovery.RefreshResult {
fun refreshAll(context: Context, ignoreResults: Boolean = false): ContactDiscovery.RefreshResult {
val stopwatch = Stopwatch("refresh-all")
val previousE164s: Set<String> = if (SignalStore.misc().cdsToken != null) {
@ -59,28 +60,45 @@ object ContactDiscoveryRefreshV2 {
val newE164s: Set<String> = newRecipientE164s + newSystemE164s
val tokenToUse: ByteArray? = if (previousE164s.isNotEmpty()) {
SignalStore.misc().cdsToken
} else {
if (SignalStore.misc().cdsToken != null) {
Log.w(TAG, "We have a token, but our previousE164 list is empty! We cannot provide a token.")
}
null
}
val response: CdsiV2Service.Response = makeRequest(
previousE164s = previousE164s,
newE164s = newE164s,
serviceIds = SignalDatabase.recipients.getAllServiceIdProfileKeyPairs(),
token = SignalStore.misc().cdsToken,
saveToken = true
token = tokenToUse,
saveToken = true,
tag = "refresh-all"
)
stopwatch.split("network")
SignalDatabase.cds.updateAfterCdsQuery(newE164s, recipientE164s + systemE164s)
stopwatch.split("cds-db")
val registeredIds: Set<RecipientId> = SignalDatabase.recipients.bulkProcessCdsV2Result(
response.results
.mapValues { entry -> RecipientDatabase.CdsV2Result(entry.value.pni, entry.value.aci.orElse(null)) }
)
stopwatch.split("recipient-db")
var registeredIds: Set<RecipientId> = emptySet()
SignalDatabase.recipients.bulkUpdatedRegisteredStatus(registeredIds.associateWith { null }, emptyList())
stopwatch.split("update-registered")
if (ignoreResults) {
Log.w(TAG, "[refresh-all] Ignoring CDSv2 results.")
} else {
registeredIds = SignalDatabase.recipients.bulkProcessCdsV2Result(
response.results
.mapValues { entry -> RecipientDatabase.CdsV2Result(entry.value.pni, entry.value.aci.orElse(null)) }
)
stopwatch.split("recipient-db")
SignalDatabase.recipients.bulkUpdatedRegisteredStatus(registeredIds.associateWith { null }, emptyList())
stopwatch.split("update-registered")
}
stopwatch.stop(TAG)
Log.d(TAG, "[refresh-all] Used ${response.quotaUsedDebugOnly} units of our quota.")
return ContactDiscovery.RefreshResult(registeredIds, emptyMap())
}
@ -89,7 +107,7 @@ object ContactDiscoveryRefreshV2 {
@WorkerThread
@Synchronized
@JvmStatic
fun refresh(context: Context, inputRecipients: List<Recipient>): ContactDiscovery.RefreshResult {
fun refresh(context: Context, inputRecipients: List<Recipient>, ignoreResults: Boolean = false): ContactDiscovery.RefreshResult {
val stopwatch = Stopwatch("refresh-some")
val recipients = inputRecipients.map { it.resolve() }
@ -100,7 +118,7 @@ object ContactDiscoveryRefreshV2 {
if (inputE164s.size > MAXIMUM_ONE_OFF_REQUEST_SIZE) {
Log.i(TAG, "List of specific recipients to refresh is too large! (Size: ${recipients.size}). Doing a full refresh instead.")
val fullResult: ContactDiscovery.RefreshResult = refreshAll(context)
val fullResult: ContactDiscovery.RefreshResult = refreshAll(context, ignoreResults)
return ContactDiscovery.RefreshResult(
registeredIds = fullResult.registeredIds.intersect(inputIds),
@ -120,35 +138,44 @@ object ContactDiscoveryRefreshV2 {
newE164s = inputE164s,
serviceIds = SignalDatabase.recipients.getAllServiceIdProfileKeyPairs(),
token = null,
saveToken = false
saveToken = false,
tag = "refresh-some"
)
stopwatch.split("network")
val registeredIds: Set<RecipientId> = SignalDatabase.recipients.bulkProcessCdsV2Result(
response.results
.mapValues { entry -> RecipientDatabase.CdsV2Result(entry.value.pni, entry.value.aci.orElse(null)) }
)
stopwatch.split("recipient-db")
var registeredIds: Set<RecipientId> = emptySet()
SignalDatabase.recipients.bulkUpdatedRegisteredStatus(registeredIds.associateWith { null }, emptyList())
stopwatch.split("update-registered")
if (ignoreResults) {
Log.w(TAG, "[refresh-some] Ignoring CDSv2 results.")
} else {
registeredIds = SignalDatabase.recipients.bulkProcessCdsV2Result(
response.results
.mapValues { entry -> RecipientDatabase.CdsV2Result(entry.value.pni, entry.value.aci.orElse(null)) }
)
stopwatch.split("recipient-db")
SignalDatabase.recipients.bulkUpdatedRegisteredStatus(registeredIds.associateWith { null }, emptyList())
stopwatch.split("update-registered")
}
Log.d(TAG, "[refresh-some] Used ${response.quotaUsedDebugOnly} units of our quota.")
stopwatch.stop(TAG)
return ContactDiscovery.RefreshResult(registeredIds, emptyMap())
}
@Throws(IOException::class)
private fun makeRequest(previousE164s: Set<String>, newE164s: Set<String>, serviceIds: Map<ServiceId, ProfileKey>, token: ByteArray?, saveToken: Boolean): CdsiV2Service.Response {
private fun makeRequest(previousE164s: Set<String>, newE164s: Set<String>, serviceIds: Map<ServiceId, ProfileKey>, token: ByteArray?, saveToken: Boolean, tag: String): CdsiV2Service.Response {
return ApplicationDependencies.getSignalServiceAccountManager().getRegisteredUsersWithCdsi(
previousE164s,
newE164s,
serviceIds,
Optional.ofNullable(token),
BuildConfig.CDSI_MRENCLAVE
) { token ->
) { tokenToSave ->
if (saveToken) {
SignalStore.misc().cdsToken = token
SignalStore.misc().cdsToken = tokenToSave
Log.d(TAG, "[$tag] Token saved!")
}
}
}

View file

@ -101,6 +101,7 @@ public final class FeatureFlags {
private static final String TELECOM_MODEL_BLOCKLIST = "android.calling.telecomModelBlockList";
private static final String CAMERAX_MODEL_BLOCKLIST = "android.cameraXModelBlockList";
private static final String RECIPIENT_MERGE_V2 = "android.recipientMergeV2";
private static final String CDS_V2_LOAD_TEST = "android.csdV2LoadTest";
/**
* We will only store remote values for flags in this set. If you want a flag to be controllable
@ -154,7 +155,8 @@ public final class FeatureFlags {
TELECOM_MANUFACTURER_ALLOWLIST,
TELECOM_MODEL_BLOCKLIST,
CAMERAX_MODEL_BLOCKLIST,
RECIPIENT_MERGE_V2
RECIPIENT_MERGE_V2,
CDS_V2_LOAD_TEST
);
@VisibleForTesting
@ -217,7 +219,8 @@ public final class FeatureFlags {
TELECOM_MANUFACTURER_ALLOWLIST,
TELECOM_MODEL_BLOCKLIST,
CAMERAX_MODEL_BLOCKLIST,
RECIPIENT_MERGE_V2
RECIPIENT_MERGE_V2,
CDS_V2_LOAD_TEST
);
/**
@ -545,6 +548,13 @@ public final class FeatureFlags {
return getBoolean(RECIPIENT_MERGE_V2, false);
}
/**
* Whether or not we should also query CDSv2 as a form of load test.
*/
public static boolean cdsV2LoadTesting() {
return getBoolean(CDS_V2_LOAD_TEST, false);
}
/** Only for rendering debug info. */
public static synchronized @NonNull Map<String, Object> getMemoryValues() {
return new TreeMap<>(REMOTE_VALUES);

View file

@ -52,13 +52,15 @@ public final class CdsiV2Service {
.map(CdsiV2Service::parseEntries)
.collect(Collectors.toList())
.flatMap(pages -> {
Map<String, ResponseItem> all = new HashMap<>();
Map<String, ResponseItem> all = new HashMap<>();
int quotaUsed = 0;
for (Response page : pages) {
all.putAll(page.getResults());
quotaUsed += page.getQuotaUsedDebugOnly();
}
return Single.just(new Response(all));
return Single.just(new Response(all, quotaUsed));
})
.map(result -> ServiceResponse.forResult(result, 200, null))
.onErrorReturn(error -> {
@ -87,7 +89,7 @@ public final class CdsiV2Service {
}
}
return new Response(results);
return new Response(results, clientResponse.getDebugPermitsUsed());
}
private static ClientRequest buildClientRequest(Request request) {
@ -146,13 +148,13 @@ public final class CdsiV2Service {
}
public static final class Request {
private final Set<String> previousE164s;
private final Set<String> newE164s;
private final Set<String> removedE164s;
final Set<String> previousE164s;
final Set<String> newE164s;
final Set<String> removedE164s;
private final Map<ServiceId, ProfileKey> serviceIds;
final Map<ServiceId, ProfileKey> serviceIds;
private final byte[] token;
final byte[] token;
public Request(Set<String> previousE164s, Set<String> newE164s, Map<ServiceId, ProfileKey> serviceIds, Optional<byte[]> token) {
if (previousE164s.size() > 0 && !token.isPresent()) {
@ -166,10 +168,6 @@ public final class CdsiV2Service {
this.token = token.orElse(null);
}
public int totalE164s() {
return previousE164s.size() + newE164s.size() - removedE164s.size();
}
public int serviceIdSize() {
return previousE164s.size() + newE164s.size() + removedE164s.size() + serviceIds.size();
}
@ -177,14 +175,24 @@ public final class CdsiV2Service {
public static final class Response {
private final Map<String, ResponseItem> results;
private final int quotaUsed;
public Response(Map<String, ResponseItem> results) {
this.results = results;
public Response(Map<String, ResponseItem> results, int quoteUsed) {
this.results = results;
this.quotaUsed = quoteUsed;
}
public Map<String, ResponseItem> getResults() {
return results;
}
/**
* Tells you how much quota you used in the request. This should only be used for debugging/logging purposed, and should never be relied upon for making
* actual decisions.
*/
public int getQuotaUsedDebugOnly() {
return quotaUsed;
}
}
public static final class ResponseItem {

View file

@ -15,7 +15,10 @@ message ClientRequest {
bytes new_e164s = 3;
bytes discard_e164s = 4;
reserved /*bool has_more*/ 5;
// If true, the client has more pairs or e164s to send. If false or unset,
// this is the client's last request, and processing should commence.
// NOT NECESSARY FOR CDSI
// bool has_more = 5;
// If set, a token which allows rate limiting to discount the e164s in
// the request's prev_e164s, only counting new_e164s. If not set, then
@ -53,6 +56,11 @@ message ClientResponse {
// e164s sent up in this request, only counting those in the next
// request's new_e164s.
bytes token = 3;
// On a successful response to a token_ack request, the number of permits
// that were deducted from the user's rate-limit in order to process the
// request
int32 debug_permits_used = 4;
}
message EnclaveLoad {
@ -71,12 +79,13 @@ message EnclaveLoad {
}
message ClientHandshakeStart {
// Public key associated with this server's enclave
bytes pubkey = 1;
// Public key associated with this server's enclave. For use in test-only
// contexts where attestation is not available
bytes test_only_pubkey = 1;
// Remote-attestation evidence associated with the public key
bytes evidence = 2;
// Endorsements of remote-attestation evidence.
bytes endorsement = 3;
}
}