From 284b683b048fac6e00b4b1c4895c684c98b29c8e Mon Sep 17 00:00:00 2001 From: Gus Brodman Date: Thu, 3 Sep 2026 14:45:38 -0400 Subject: [PATCH] Reduce SafeBrowsing API calls by prefix hashing SafeBrowsing provides an API endpoint that returns, for each threat type provided, a base64-encoded string that represents the concatenation of 4-byte prefixes of sha256 hashes of harmful domains. That means we can do the following - decode the base64 string into a byte[] - interpret the byte[] as an int[], because each int is four bytes - for each domain: - hash the domain and take the hash's first four bytes as an int - if that int is in the prefix hash array, query SafeBrowsing directly with that domain name - if the int is not in the prefix hash array, we know the domain is not harmful (according to SafeBrowsing at least) The prefix list provided by SafeBrowsing contains about 2.5 million entries which corresponds to only about 10 MB of heap memory and only about 0.1% of the space of all integers. Thus, we can assume there won't be too many hash collisions. This will reduce our number of API calls from ~5100 calls (each containing 490 domains) to 6 API calls + (# actually-harmful domains / 490). --- .../beam/spec11/SafeBrowsingTransforms.java | 249 +++++++++++++++--- .../registry/beam/spec11/Spec11Pipeline.java | 60 ++++- .../spec11/SafeBrowsingTransformsTest.java | 178 ++++++++++++- .../beam/spec11/Spec11PipelineTest.java | 20 +- .../registry/beam/spec11/test_output.txt | 4 +- 5 files changed, 440 insertions(+), 71 deletions(-) diff --git a/core/src/main/java/google/registry/beam/spec11/SafeBrowsingTransforms.java b/core/src/main/java/google/registry/beam/spec11/SafeBrowsingTransforms.java index ae3f3a49b65..f536960f3a9 100644 --- a/core/src/main/java/google/registry/beam/spec11/SafeBrowsingTransforms.java +++ b/core/src/main/java/google/registry/beam/spec11/SafeBrowsingTransforms.java @@ -14,28 +14,35 @@ package google.registry.beam.spec11; +import static com.google.common.base.Preconditions.checkArgument; import static java.nio.charset.StandardCharsets.UTF_8; -import static org.apache.http.HttpHeaders.RETRY_AFTER; import static org.apache.http.HttpStatus.SC_OK; -import static org.apache.http.HttpStatus.SC_TOO_MANY_REQUESTS; import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Ascii; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.google.common.flogger.FluentLogger; +import com.google.common.hash.Hashing; import com.google.common.io.CharStreams; import google.registry.util.Clock; import google.registry.util.Retrier; +import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.InputStreamReader; import java.io.Serializable; import java.net.URISyntaxException; +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.Base64; import java.util.LinkedHashMap; import java.util.Map; import java.util.function.Supplier; +import javax.annotation.Nullable; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.windowing.GlobalWindow; import org.apache.beam.sdk.values.KV; -import org.apache.http.Header; +import org.apache.beam.sdk.values.PCollectionView; import org.apache.http.client.methods.CloseableHttpResponse; import org.apache.http.client.methods.HttpPost; import org.apache.http.client.utils.URIBuilder; @@ -43,7 +50,6 @@ import org.apache.http.entity.ContentType; import org.apache.http.impl.client.CloseableHttpClient; import org.apache.http.impl.client.HttpClients; -import org.apache.http.protocol.HTTP; import org.json.JSONArray; import org.json.JSONException; import org.json.JSONObject; @@ -53,14 +59,180 @@ public class SafeBrowsingTransforms { private static final FluentLogger logger = FluentLogger.forEnclosingClass(); - /** The URL to send SafeBrowsing API calls (POSTS) to. */ + /** The URL to send SafeBrowsing Lookup API calls (threatMatches:find) to. */ private static final String SAFE_BROWSING_URL = "https://safebrowsing.googleapis.com/v4/threatMatches:find"; + /** The URL to fetch SafeBrowsing threat list updates (threatListUpdates:fetch) from. */ + private static final String THREAT_LIST_UPDATES_URL = + "https://safebrowsing.googleapis.com/v4/threatListUpdates:fetch"; + + /** The threat types evaluated for Spec11 reporting. */ + static final ImmutableList THREAT_TYPES = + ImmutableList.of("MALWARE", "SOCIAL_ENGINEERING", "UNWANTED_SOFTWARE"); + + /** Computes the SHA-256 hash for a given URL or host/path expression. */ + public static byte[] computeSha256(String expression) { + return Hashing.sha256().hashString(expression, UTF_8).asBytes(); + } + + /** Converts the first 4 bytes of a SHA-256 hash into an integer. */ + public static int getPrefixInt(byte[] sha256Bytes) { + return ByteBuffer.wrap(sha256Bytes).getInt(); + } + + /** + * {@link DoFn} fetching threat list hash prefixes from the SafeBrowsing Update API. + * + *

SafeBrowsing provides an endpoint returning 4-byte hash prefixes of harmful URLs. This + * serves as an in-memory prefix filter; domains whose hashes do not match any prefix are + * guaranteed clean and require no API calls (similar to a bloom filter). For the small fraction + * of domains that match a prefix (actual threats or hash collisions), we confirm them via the + * Lookup API. + */ + public static class FetchThreatListPrefixesFn extends DoFn { + + private final Retrier retrier; + private final Supplier closeableHttpClientSupplier; + + FetchThreatListPrefixesFn(Retrier retrier) { + this(retrier, (Supplier & Serializable) HttpClients::createDefault); + } + + @VisibleForTesting + FetchThreatListPrefixesFn(Retrier retrier, Supplier clientSupplier) { + this.retrier = retrier; + this.closeableHttpClientSupplier = clientSupplier; + } + + @ProcessElement + public void processElement(@Element String apiKey, OutputReceiver output) { + try { + URIBuilder uriBuilder = new URIBuilder(THREAT_LIST_UPDATES_URL); + uriBuilder.addParameter("key", apiKey); + + HttpPost httpPost = new HttpPost(uriBuilder.build()); + JSONObject requestBody = createFetchRequestBody(); + httpPost.setEntity( + new ByteArrayEntity( + requestBody.toString().getBytes(UTF_8), ContentType.APPLICATION_JSON)); + + int[] prefixes = + retrier.callWithRetry( + () -> { + try (CloseableHttpClient client = closeableHttpClientSupplier.get(); + CloseableHttpResponse response = client.execute(httpPost)) { + return processSafeBrowsingFetchResponse(response); + } + }, + IOException.class); + output.output(prefixes); + } catch (URISyntaxException | JSONException e) { + throw new RuntimeException("Caught exception fetching threat list prefixes.", e); + } + } + } + + private static JSONObject createFetchRequestBody() throws JSONException { + JSONArray listUpdateRequests = new JSONArray(); + for (String threatType : THREAT_TYPES) { + listUpdateRequests.put( + new JSONObject() + .put("threatType", threatType) + .put("platformType", "ANY_PLATFORM") + .put("threatEntryType", "URL") + .put( + "constraints", + new JSONObject().put("supportedCompressions", new JSONArray().put("RAW")))); + } + return new JSONObject() + .put( + "client", + new JSONObject().put("clientId", "domainregistry").put("clientVersion", "0.0.1")) + .put("listUpdateRequests", listUpdateRequests); + } + + /** + * Fetches and unpacks threat-match hash prefixes from a SafeBrowsing response. + * + *

SafeBrowsing provides an API endpoint (threatListUpdates) that includes a base64 string + * representing hash prefixes. When we decode that base64 string, we get a byte array representing + * the concatenation of 4-byte prefixes of sha256 hashes of harmful domains. + * + *

We represent the result as an int array, with each 4-byte int representing the first four + * bytes of a sha256 hash of a harmful domain. + */ + private static int[] processSafeBrowsingFetchResponse(CloseableHttpResponse response) + throws IOException, JSONException { + int statusCode = response.getStatusLine().getStatusCode(); + if (statusCode != SC_OK) { + throw new IOException( + String.format( + "Got unexpected status code %s from threatListUpdates response.", statusCode)); + } + try (InputStreamReader reader = + new InputStreamReader(response.getEntity().getContent(), UTF_8)) { + JSONObject responseBody = new JSONObject(CharStreams.toString(reader)); + JSONArray listUpdateResponses = responseBody.optJSONArray("listUpdateResponses"); + if (listUpdateResponses == null || listUpdateResponses.isEmpty()) { + throw new IOException("No value for listUpdateResponses in threatListUpdates response"); + } + // We expect 10-20 MB of data, so initialize the stream with a relatively large size + ByteArrayOutputStream bytes = new ByteArrayOutputStream(10 * 1024 * 1024); + // We ask SafeBrowsing for multiple threat types, and they return a different object for each + for (int i = 0; i < listUpdateResponses.length(); i++) { + JSONArray additions = listUpdateResponses.getJSONObject(i).optJSONArray("additions"); + if (additions != null) { + for (int j = 0; j < additions.length(); j++) { + JSONObject rawHashes = additions.getJSONObject(j).optJSONObject("rawHashes"); + if (rawHashes != null) { + int prefixSize = rawHashes.optInt("prefixSize"); + checkArgument(prefixSize == 4, "Expected prefixSize of 4, got %s", prefixSize); + String base64 = rawHashes.optString("rawHashes", ""); + if (!base64.isEmpty()) { + bytes.writeBytes(Base64.getDecoder().decode(base64)); + } + } + } + } + } + byte[] byteArray = bytes.toByteArray(); + if (byteArray.length == 0 || byteArray.length % 4 != 0) { + throw new IOException( + String.format( + "Expected prefix byte array length to be a positive multiple of 4, got %d", + byteArray.length)); + } + int[] prefixes = new int[byteArray.length / 4]; + ByteBuffer.wrap(byteArray).asIntBuffer().get(prefixes); + // Sort the array to normalize signed integer order and enable binary search + Arrays.sort(prefixes); + // Remove any duplicates to reduce broadcast size + return removeDuplicates(prefixes); + } + } + + /** Removes duplicates from an already sorted prefix array in-place. */ + private static int[] removeDuplicates(int[] prefixes) { + int uniquePrefix = 0; + for (int i = 0; i < prefixes.length; i++) { + if (i == 0 || prefixes[i] != prefixes[i - 1]) { + prefixes[uniquePrefix] = prefixes[i]; + uniquePrefix++; + } + } + if (uniquePrefix < prefixes.length) { + prefixes = Arrays.copyOf(prefixes, uniquePrefix); + } + return prefixes; + } + /** * {@link DoFn} mapping a {@link DomainNameInfo} to its evaluation report from SafeBrowsing. * - *

Refer to the Lookup API documentation for the request/response format and other details. + *

Incoming domains are first checked locally against the 4-byte hash prefixes side input. Any + * domains matching a prefix (potential threats or hash collisions) are then confirmed in batches + * using the SafeBrowsing Lookup API. * * @see Lookup API */ @@ -80,8 +252,8 @@ static class EvaluateSafeBrowsingFn private final Clock clock; /** - * Maps a domain name's {@code domainName} to its corresponding {@link DomainNameInfo} to - * facilitate batching SafeBrowsing API requests. + * Maps a domain's {@code domainName} to its corresponding {@link DomainNameInfo} to facilitate + * batching SafeBrowsing API requests. */ private final Map domainNameInfoBuffer = new LinkedHashMap<>(BATCH_SIZE); @@ -97,6 +269,12 @@ static class EvaluateSafeBrowsingFn /** Retries on receiving transient failures such as {@link IOException}. */ private final Retrier retrier; + /** Side input providing the sorted 4-byte hash prefixes of harmful domains. */ + private final PCollectionView prefixesView; + + /** Cached reference to the hash prefix side input to avoid repeated lookups per element. */ + private transient @Nullable int[] prefixes; + /** * Constructs a {@link EvaluateSafeBrowsingFn} with a given API key. * @@ -105,12 +283,14 @@ static class EvaluateSafeBrowsingFn * because class methods are generally serializable, especially a static function such as {@link * HttpClients#createDefault()}. */ - EvaluateSafeBrowsingFn(String apiKey, Retrier retrier, Clock clock) { + EvaluateSafeBrowsingFn( + String apiKey, Retrier retrier, Clock clock, PCollectionView prefixesView) { this( apiKey, retrier, clock, - (Supplier & Serializable) HttpClients::createDefault); + (Supplier & Serializable) HttpClients::createDefault, + prefixesView); } /** @@ -121,11 +301,16 @@ static class EvaluateSafeBrowsingFn */ @VisibleForTesting EvaluateSafeBrowsingFn( - String apiKey, Retrier retrier, Clock clock, Supplier clientSupplier) { + String apiKey, + Retrier retrier, + Clock clock, + Supplier clientSupplier, + PCollectionView prefixesView) { this.apiKey = apiKey; this.retrier = retrier; this.clock = clock; this.closeableHttpClientSupplier = clientSupplier; + this.prefixesView = prefixesView; } /** Evaluates any buffered {@link DomainNameInfo} objects upon completing the bundle. */ @@ -145,12 +330,24 @@ public void finishBundle(FinishBundleContext context) { } /** - * Buffers {@link DomainNameInfo} objects until we reach the batch size, then bulk-evaluate the - * URLs with the SafeBrowsing API. + * Checks each domain against the hash prefix side input. Matching domains (potential threats or + * collisions) are buffered until reaching {@link #BATCH_SIZE} and then evaluated in bulk via + * the SafeBrowsing Lookup API. */ @ProcessElement public void processElement(ProcessContext context) { + if (prefixes == null) { + prefixes = context.sideInput(prefixesView); + } DomainNameInfo domainNameInfo = context.element(); + // The canonical domain form is e.g. "mydomain.tld/". See + // https://developers.google.com/safe-browsing/v4/urls-hashing for more details + byte[] hash = computeSha256(Ascii.toLowerCase(domainNameInfo.domainName()) + "/"); + int prefixInt = getPrefixInt(hash); + if (Arrays.binarySearch(prefixes, prefixInt) < 0) { + // The domain's prefix does not match any threat prefix; it cannot be marked as harmful + return; + } domainNameInfoBuffer.put(domainNameInfo.domainName(), domainNameInfo); if (domainNameInfoBuffer.size() >= BATCH_SIZE) { ImmutableSet> results = evaluateAndFlush(); @@ -173,10 +370,11 @@ private ImmutableSet> evaluateAndFlush() { uriBuilder.addParameter("key", apiKey); HttpPost httpPost = new HttpPost(uriBuilder.build()); - httpPost.addHeader(HTTP.CONTENT_TYPE, ContentType.APPLICATION_JSON.toString()); JSONObject requestBody = createRequestBody(); - httpPost.setEntity(new ByteArrayEntity(requestBody.toString().getBytes(UTF_8))); + httpPost.setEntity( + new ByteArrayEntity( + requestBody.toString().getBytes(UTF_8), ContentType.APPLICATION_JSON)); // Retry transient exceptions such as IOException retrier.callWithRetry( () -> { @@ -187,7 +385,7 @@ private ImmutableSet> evaluateAndFlush() { }, IOException.class); } catch (URISyntaxException | JSONException e) { - // Fail the pipeline on a parsing exception- this indicates the API likely changed. + // Fail the pipeline on a parsing exception. This indicates the API likely changed. throw new RuntimeException("Caught parsing exception, failing pipeline.", e); } finally { // Flush the buffer @@ -211,12 +409,7 @@ private JSONObject createRequestBody() throws JSONException { .put( "threatInfo", new JSONObject() - .put( - "threatTypes", - new JSONArray() - .put("MALWARE") - .put("SOCIAL_ENGINEERING") - .put("UNWANTED_SOFTWARE")) + .put("threatTypes", new JSONArray(THREAT_TYPES)) .put("platformTypes", new JSONArray().put("ANY_PLATFORM")) .put("threatEntryTypes", new JSONArray().put("URL")) .put("threatEntries", threatArray)); @@ -232,14 +425,6 @@ private void processResponse( throws IOException { int statusCode = response.getStatusLine().getStatusCode(); if (statusCode != SC_OK) { - if (statusCode == SC_TOO_MANY_REQUESTS) { - Header retryAfterHeader = response.getFirstHeader(RETRY_AFTER); - if (retryAfterHeader != null) { - logger.atWarning().log( - "SafeBrowsing API returned 429 with Retry-After header: %s", - retryAfterHeader.getValue()); - } - } throw new IOException( String.format("Got unexpected status code %s from response.", statusCode)); } @@ -256,6 +441,10 @@ private void processResponse( JSONObject match = threatMatches.getJSONObject(i); String url = match.getJSONObject("threat").getString("url"); DomainNameInfo domainNameInfo = domainNameInfoBuffer.get(url); + if (domainNameInfo == null) { + throw new IOException( + String.format("Received threat match for domain not present in buffer: %s", url)); + } resultBuilder.add( KV.of( domainNameInfo, diff --git a/core/src/main/java/google/registry/beam/spec11/Spec11Pipeline.java b/core/src/main/java/google/registry/beam/spec11/Spec11Pipeline.java index 081f84710a9..9416d48f902 100644 --- a/core/src/main/java/google/registry/beam/spec11/Spec11Pipeline.java +++ b/core/src/main/java/google/registry/beam/spec11/Spec11Pipeline.java @@ -16,6 +16,7 @@ import static com.google.common.base.Preconditions.checkArgument; +import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableSet; import dagger.Component; import dagger.Module; @@ -23,6 +24,7 @@ import google.registry.beam.common.RegistryJpaIO; import google.registry.beam.common.RegistryJpaIO.Read; import google.registry.beam.spec11.SafeBrowsingTransforms.EvaluateSafeBrowsingFn; +import google.registry.beam.spec11.SafeBrowsingTransforms.FetchThreatListPrefixesFn; import google.registry.config.RegistryConfig.ConfigModule; import google.registry.model.reporting.Spec11ThreatMatch; import google.registry.model.reporting.Spec11ThreatMatch.ThreatType; @@ -35,20 +37,26 @@ import java.io.Serializable; import java.time.LocalDate; import java.time.YearMonth; +import java.util.function.Supplier; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.coders.SerializableCoder; import org.apache.beam.sdk.io.TextIO; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.GroupByKey; import org.apache.beam.sdk.transforms.MapElements; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.Reshuffle; +import org.apache.beam.sdk.transforms.View; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.sdk.values.TypeDescriptor; import org.apache.beam.sdk.values.TypeDescriptors; +import org.apache.http.impl.client.CloseableHttpClient; +import org.apache.http.impl.client.HttpClients; import org.json.JSONArray; import org.json.JSONException; import org.json.JSONObject; @@ -84,11 +92,28 @@ public static String getSpec11ReportFilePath(LocalDate localDate) { public static final String THREAT_MATCHES_FIELD = "threatMatches"; private final Spec11PipelineOptions options; - private final EvaluateSafeBrowsingFn safeBrowsingFn; + private final Clock clock; + private final Retrier retrier; + private final Supplier clientSupplier; - Spec11Pipeline(Spec11PipelineOptions options, EvaluateSafeBrowsingFn safeBrowsingFn) { + Spec11Pipeline(Spec11PipelineOptions options, Clock clock, Retrier retrier) { + this( + options, + clock, + retrier, + (Supplier & Serializable) HttpClients::createDefault); + } + + @VisibleForTesting + Spec11Pipeline( + Spec11PipelineOptions options, + Clock clock, + Retrier retrier, + Supplier clientSupplier) { this.options = options; - this.safeBrowsingFn = safeBrowsingFn; + this.clock = clock; + this.retrier = retrier; + this.clientSupplier = clientSupplier; } PipelineResult run() { @@ -99,10 +124,25 @@ PipelineResult run() { void setupPipeline(Pipeline pipeline) { options.setIsolationOverride(TransactionIsolationLevel.TRANSACTION_READ_COMMITTED); + + PCollectionView prefixView = + pipeline + .apply("Get SafeBrowsing API key", Create.of(options.getSafeBrowsingApiKey())) + .apply( + "Fetch threat prefixes from SafeBrowsing", + ParDo.of(new FetchThreatListPrefixesFn(retrier, clientSupplier))) + .setCoder(SerializableCoder.of(int[].class)) + .apply("Create prefix view", View.asSingleton()); + PCollection domains = readFromCloudSql(pipeline); + EvaluateSafeBrowsingFn safeBrowsingFn = + new EvaluateSafeBrowsingFn( + options.getSafeBrowsingApiKey(), retrier, clock, clientSupplier, prefixView); + PCollection> threatMatches = - domains.apply("Run through SafeBrowsing API", ParDo.of(safeBrowsingFn)); + domains.apply( + "Run through SafeBrowsing API", ParDo.of(safeBrowsingFn).withSideInputs(prefixView)); saveToSql(threatMatches, options); saveToGcs(threatMatches, options); @@ -238,18 +278,10 @@ Spec11PipelineOptions provideOptions() { } @Provides - EvaluateSafeBrowsingFn provideSafeBrowsingFn( - Spec11PipelineOptions options, Clock clock, Sleeper sleeper) { + Spec11Pipeline providePipeline(Spec11PipelineOptions options, Clock clock, Sleeper sleeper) { // Have a noticeably longer backoff for SafeBrowsing retries to mitigate any 429s Retrier safeBrowsingRetrier = new Retrier(sleeper, 9, 1000L); - return new EvaluateSafeBrowsingFn( - options.getSafeBrowsingApiKey(), safeBrowsingRetrier, clock); - } - - @Provides - Spec11Pipeline providePipeline( - Spec11PipelineOptions options, EvaluateSafeBrowsingFn safeBrowsingFn) { - return new Spec11Pipeline(options, safeBrowsingFn); + return new Spec11Pipeline(options, clock, safeBrowsingRetrier); } } diff --git a/core/src/test/java/google/registry/beam/spec11/SafeBrowsingTransformsTest.java b/core/src/test/java/google/registry/beam/spec11/SafeBrowsingTransformsTest.java index 54bbf6eb490..e49e9636b63 100644 --- a/core/src/test/java/google/registry/beam/spec11/SafeBrowsingTransformsTest.java +++ b/core/src/test/java/google/registry/beam/spec11/SafeBrowsingTransformsTest.java @@ -14,7 +14,10 @@ package google.registry.beam.spec11; +import static com.google.common.truth.Truth.assertThat; +import static google.registry.beam.spec11.SafeBrowsingTransforms.THREAT_TYPES; import static java.nio.charset.StandardCharsets.UTF_8; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -23,24 +26,33 @@ import com.google.common.base.Suppliers; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Iterables; import com.google.common.io.CharStreams; import google.registry.beam.TestPipelineExtension; import google.registry.beam.spec11.SafeBrowsingTransforms.EvaluateSafeBrowsingFn; +import google.registry.beam.spec11.SafeBrowsingTransforms.FetchThreatListPrefixesFn; import google.registry.testing.FakeClock; import google.registry.testing.FakeSleeper; import google.registry.util.Retrier; import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.InputStreamReader; import java.io.ObjectInputStream; import java.io.ObjectOutputStream; import java.io.Serializable; +import java.util.Arrays; +import java.util.Base64; +import java.util.Map; +import org.apache.beam.sdk.Pipeline.PipelineExecutionException; import org.apache.beam.sdk.coders.SerializableCoder; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.View; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionView; import org.apache.http.ProtocolVersion; import org.apache.http.client.methods.CloseableHttpResponse; import org.apache.http.client.methods.HttpPost; @@ -70,9 +82,9 @@ class SafeBrowsingTransformsTest { "party-night.net", "SOCIAL_ENGINEERING", "bitcoin.bank", - "POTENTIALLY_HARMFUL_APPLICATION", + "MALWARE", "no-email.com", - "THREAT_TYPE_UNSPECIFIED", + "SOCIAL_ENGINEERING", "anti-anti-anti-virus.dev", "UNWANTED_SOFTWARE"); @@ -82,22 +94,37 @@ class SafeBrowsingTransformsTest { private static ImmutableMap THREAT_MATCH_MAP; + private static boolean threatListEmpty = false; + private final CloseableHttpClient mockHttpClient = mock(CloseableHttpClient.class, withSettings().serializable()); private final FakeClock clock = new FakeClock(); - private final EvaluateSafeBrowsingFn safeBrowsingFn = - new EvaluateSafeBrowsingFn( - "API_KEY", - new Retrier(new FakeSleeper(clock), 1), - clock, - Suppliers.ofInstance(mockHttpClient)); - @RegisterExtension final TestPipelineExtension pipeline = TestPipelineExtension.create().enableAbandonedNodeEnforcement(true); + private PCollectionView createPrefixView() { + FetchThreatListPrefixesFn fetchFn = + new FetchThreatListPrefixesFn( + new Retrier(new FakeSleeper(clock), 1), Suppliers.ofInstance(mockHttpClient)); + return pipeline + .apply("Retrieve API key", Create.of("API_KEY")) + .apply("Fetch threat prefixes", ParDo.of(fetchFn)) + .setCoder(SerializableCoder.of(int[].class)) + .apply("Create prefix side input", View.asSingleton()); + } + + private EvaluateSafeBrowsingFn createSafeBrowsingFn(PCollectionView prefixView) { + return new EvaluateSafeBrowsingFn( + "API_KEY", + new Retrier(new FakeSleeper(clock), 1), + clock, + Suppliers.ofInstance(mockHttpClient), + prefixView); + } + private static DomainNameInfo createDomainNameInfo(String url) { return DomainNameInfo.create(url, REPO_ID, REGISTRAR_ID, REGISTRAR_EMAIL); } @@ -122,11 +149,13 @@ static void beforeAll() { @BeforeEach void beforeEach() throws Exception { + threatListEmpty = false; when(mockHttpClient.execute(any(HttpPost.class))).thenAnswer(new HttpResponder()); } @Test void testSuccess_someBadDomains() throws Exception { + PCollectionView prefixView = createPrefixView(); ImmutableList domainNameInfos = ImmutableList.of( createDomainNameInfo("111.com"), @@ -137,7 +166,7 @@ void testSuccess_someBadDomains() throws Exception { PCollection> threats = pipeline .apply(Create.of(domainNameInfos).withCoder(SerializableCoder.of(DomainNameInfo.class))) - .apply(ParDo.of(safeBrowsingFn)); + .apply(ParDo.of(createSafeBrowsingFn(prefixView)).withSideInputs(prefixView)); PAssert.that(threats) .containsInAnyOrder( @@ -150,6 +179,7 @@ void testSuccess_someBadDomains() throws Exception { @Test void testSuccess_noBadDomains() throws Exception { + PCollectionView prefixView = createPrefixView(); ImmutableList domainNameInfos = ImmutableList.of( createDomainNameInfo("hello_kitty.dev"), @@ -158,12 +188,80 @@ void testSuccess_noBadDomains() throws Exception { PCollection> threats = pipeline .apply(Create.of(domainNameInfos).withCoder(SerializableCoder.of(DomainNameInfo.class))) - .apply(ParDo.of(safeBrowsingFn)); + .apply(ParDo.of(createSafeBrowsingFn(prefixView)).withSideInputs(prefixView)); + + PAssert.that(threats).empty(); + pipeline.run().waitUntilFinish(); + } + + @Test + void testSuccess_hashCollisionFalsePositive() { + byte[] hash = SafeBrowsingTransforms.computeSha256("clean-domain.com/"); + int prefixInt = SafeBrowsingTransforms.getPrefixInt(hash); + PCollectionView prefixView = + pipeline + .apply("Create colliding prefix", Create.of(new int[] {prefixInt})) + .apply("View colliding prefix", View.asSingleton()); + + ImmutableList domainNameInfos = + ImmutableList.of(createDomainNameInfo("clean-domain.com")); + PCollection> threats = + pipeline + .apply(Create.of(domainNameInfos).withCoder(SerializableCoder.of(DomainNameInfo.class))) + .apply(ParDo.of(createSafeBrowsingFn(prefixView)).withSideInputs(prefixView)); PAssert.that(threats).empty(); pipeline.run().waitUntilFinish(); } + @Test + void testSuccess_fetchThreatListPrefixes() { + FetchThreatListPrefixesFn fetchFn = + new FetchThreatListPrefixesFn( + new Retrier(new FakeSleeper(clock), 1), Suppliers.ofInstance(mockHttpClient)); + PCollection prefixDb = + pipeline + .apply(Create.of("API_KEY")) + .apply(ParDo.of(fetchFn)) + .setCoder(SerializableCoder.of(int[].class)); + + PAssert.that(prefixDb) + .satisfies( + iterable -> { + int[] prefixes = Iterables.getOnlyElement(iterable); + assertThat(prefixes.length).isGreaterThan(0); + for (String badDomain : THREAT_MAP.keySet()) { + byte[] hash = SafeBrowsingTransforms.computeSha256(badDomain + "/"); + int prefixInt = SafeBrowsingTransforms.getPrefixInt(hash); + assertThat(Arrays.binarySearch(prefixes, prefixInt)).isAtLeast(0); + } + byte[] cleanHash = + SafeBrowsingTransforms.computeSha256("clean-domain-never-bad.com/"); + int cleanPrefixInt = SafeBrowsingTransforms.getPrefixInt(cleanHash); + assertThat(Arrays.binarySearch(prefixes, cleanPrefixInt)).isLessThan(0); + return null; + }); + pipeline.run().waitUntilFinish(); + } + + @Test + void testFailure_emptyPrefixes() { + threatListEmpty = true; + FetchThreatListPrefixesFn fetchFn = + new FetchThreatListPrefixesFn( + new Retrier(new FakeSleeper(clock), 1), Suppliers.ofInstance(mockHttpClient)); + pipeline.apply(Create.of("API_KEY")).apply(ParDo.of(fetchFn)); + + PipelineExecutionException thrown = + assertThrows(PipelineExecutionException.class, () -> pipeline.run().waitUntilFinish()); + assertThat(thrown).hasCauseThat().hasCauseThat().isInstanceOf(IOException.class); + assertThat(thrown) + .hasCauseThat() + .hasCauseThat() + .hasMessageThat() + .isEqualTo("Expected prefix byte array length to be a positive multiple of 4, got 0"); + } + /** * A serializable {@link Answer} that returns a mock HTTP response based on the HTTP request's * content. @@ -171,11 +269,63 @@ void testSuccess_noBadDomains() throws Exception { static class HttpResponder implements Answer, Serializable { @Override public CloseableHttpResponse answer(InvocationOnMock invocation) throws Throwable { + HttpPost post = (HttpPost) invocation.getArguments()[0]; + String uri = post.getURI().toString(); + if (uri.contains("threatListUpdates:fetch")) { + return getThreatListUpdatesMockResponse(threatListEmpty); + } return getMockResponse( - CharStreams.toString( - new InputStreamReader( - ((HttpPost) invocation.getArguments()[0]).getEntity().getContent(), UTF_8))); + CharStreams.toString(new InputStreamReader(post.getEntity().getContent(), UTF_8))); + } + } + + private static CloseableHttpResponse getThreatListUpdatesMockResponse() throws JSONException { + return getThreatListUpdatesMockResponse(false); + } + + private static CloseableHttpResponse getThreatListUpdatesMockResponse(boolean empty) + throws JSONException { + JSONObject response = new JSONObject(); + JSONArray listUpdateResponses = new JSONArray(); + + for (String threatType : THREAT_TYPES) { + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + if (!empty) { + for (Map.Entry entry : THREAT_MAP.entrySet()) { + if (entry.getValue().equals(threatType)) { + byte[] hash = SafeBrowsingTransforms.computeSha256(entry.getKey() + "/"); + bytes.write(hash, 0, 4); + } + } + } + JSONObject listResponse = + new JSONObject() + .put("threatType", threatType) + .put("platformType", "ANY_PLATFORM") + .put("threatEntryType", "URL") + .put("responseType", "FULL_UPDATE"); + listResponse.put( + "additions", + new JSONArray() + .put( + new JSONObject() + .put( + "rawHashes", + new JSONObject() + .put("prefixSize", 4) + .put( + "rawHashes", + Base64.getEncoder().encodeToString(bytes.toByteArray()))))); + listUpdateResponses.put(listResponse); } + response.put("listUpdateResponses", listUpdateResponses); + + CloseableHttpResponse httpResponse = + mock(CloseableHttpResponse.class, withSettings().serializable()); + when(httpResponse.getStatusLine()) + .thenReturn(new BasicStatusLine(new ProtocolVersion("HTTP", 1, 1), 200, "Done")); + when(httpResponse.getEntity()).thenReturn(new FakeHttpEntity(response.toString())); + return httpResponse; } /** diff --git a/core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java b/core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java index b618bc483cd..14709fb5931 100644 --- a/core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java +++ b/core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java @@ -35,7 +35,6 @@ import com.google.common.truth.Correspondence; import com.google.common.truth.Correspondence.BinaryPredicate; import google.registry.beam.TestPipelineExtension; -import google.registry.beam.spec11.SafeBrowsingTransforms.EvaluateSafeBrowsingFn; import google.registry.beam.spec11.SafeBrowsingTransformsTest.HttpResponder; import google.registry.model.domain.Domain; import google.registry.model.domain.DomainAuthInfo; @@ -105,8 +104,8 @@ class Spec11PipelineTest { ImmutableList.of( ThreatMatch.create("MALWARE", "111.com"), ThreatMatch.create("SOCIAL_ENGINEERING", "party-night.net"), - ThreatMatch.create("POTENTIALLY_HARMFUL_APPLICATION", "bitcoin.bank"), - ThreatMatch.create("THREAT_TYPE_UNSPECIFIED", "no-eamil.com"), + ThreatMatch.create("MALWARE", "bitcoin.bank"), + ThreatMatch.create("SOCIAL_ENGINEERING", "no-eamil.com"), ThreatMatch.create("UNWANTED_SOFTWARE", "anti-anti-anti-virus.dev")); @TempDir Path tmpDir; @@ -164,14 +163,14 @@ void beforeEach() throws Exception { .setDomainRepoId("1C3D5E7F9-BANK") .setRegistrarId("hello-registrar") .setCheckDate(LocalDate.of(2020, 1, 27)) - .setThreatTypes(ImmutableSet.of(ThreatType.POTENTIALLY_HARMFUL_APPLICATION)) + .setThreatTypes(ImmutableSet.of(ThreatType.MALWARE)) .build(), new Spec11ThreatMatch.Builder() .setDomainName("no-email.com") .setDomainRepoId("2A4BA9BBC-COM") .setRegistrarId("kitty-registrar") .setCheckDate(LocalDate.of(2020, 1, 27)) - .setThreatTypes(ImmutableSet.of(ThreatType.THREAT_TYPE_UNSPECIFIED)) + .setThreatTypes(ImmutableSet.of(ThreatType.SOCIAL_ENGINEERING)) .build(), new Spec11ThreatMatch.Builder() .setDomainName("anti-anti-anti-virus.dev") @@ -185,14 +184,13 @@ void beforeEach() throws Exception { @Test void testSuccess_fullSqlPipeline() throws Exception { setupCloudSql(); - EvaluateSafeBrowsingFn safeBrowsingFn = - new EvaluateSafeBrowsingFn( - SAFE_BROWSING_API_KEY, - new Retrier(new FakeSleeper(fakeClock), 1), + when(mockHttpClient.execute(any(HttpPost.class))).thenAnswer(new HttpResponder()); + Spec11Pipeline spec11Pipeline = + new Spec11Pipeline( + options, fakeClock, + new Retrier(new FakeSleeper(fakeClock), 1), Suppliers.ofInstance(mockHttpClient)); - when(mockHttpClient.execute(any(HttpPost.class))).thenAnswer(new HttpResponder()); - Spec11Pipeline spec11Pipeline = new Spec11Pipeline(options, safeBrowsingFn); spec11Pipeline.setupPipeline(pipeline); pipeline.run(options).waitUntilFinish(); verifySaveToGcs(); diff --git a/core/src/test/resources/google/registry/beam/spec11/test_output.txt b/core/src/test/resources/google/registry/beam/spec11/test_output.txt index b89b6a31e0d..d7ebb17a21c 100644 --- a/core/src/test/resources/google/registry/beam/spec11/test_output.txt +++ b/core/src/test/resources/google/registry/beam/spec11/test_output.txt @@ -1,4 +1,4 @@ Map from registrar email / name to detected domain name threats: {"threatMatches":[{"threatType":"UNWANTED_SOFTWARE","domainName":"anti-anti-anti-virus.dev"}],"registrarClientId":"cool-registrar","registrarEmailAddress":"cool@aid.net"} -{"threatMatches":[{"threatType":"MALWARE","domainName":"111.com"},{"threatType":"POTENTIALLY_HARMFUL_APPLICATION","domainName":"bitcoin.bank"}],"registrarClientId":"hello-registrar","registrarEmailAddress":"email@hello.net"} -{"threatMatches":[{"threatType":"THREAT_TYPE_UNSPECIFIED","domainName":"no-eamil.com"},{"threatType":"SOCIAL_ENGINEERING","domainName":"party-night.net"}],"registrarClientId":"kitty-registrar","registrarEmailAddress":"contact@kit.ty"} \ No newline at end of file +{"threatMatches":[{"threatType":"MALWARE","domainName":"111.com"},{"threatType":"MALWARE","domainName":"bitcoin.bank"}],"registrarClientId":"hello-registrar","registrarEmailAddress":"email@hello.net"} +{"threatMatches":[{"threatType":"SOCIAL_ENGINEERING","domainName":"no-eamil.com"},{"threatType":"SOCIAL_ENGINEERING","domainName":"party-night.net"}],"registrarClientId":"kitty-registrar","registrarEmailAddress":"contact@kit.ty"} \ No newline at end of file