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 ae3f3a49b..19261200e 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,34 @@ package google.registry.beam.spec11; +import static com.google.common.base.Preconditions.checkState; 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.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 +49,6 @@ import org.apache.http.entity.ByteArrayEntity; 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 +58,184 @@ 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 JSONArray THREAT_TYPES = + new JSONArray(new String[] {"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 (Object 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. NB: the number of elements (approximately 2.5 + * million as of October 2026) means that boxing the ints into Integers and using a hash table + * ({@code HashSet}) would use a prohibitively high amount of memory relative to a + * primitive array. Each binary search will only do about 21 primitive comparisons, which is + * easily fast enough. + */ + 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"); + checkState(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 +255,8 @@ public class SafeBrowsingTransforms { 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 +272,12 @@ public class SafeBrowsingTransforms { /** 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 +286,14 @@ public class SafeBrowsingTransforms { * 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 +304,16 @@ public class SafeBrowsingTransforms { */ @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 +333,24 @@ public class SafeBrowsingTransforms { } /** - * 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 +373,11 @@ public class SafeBrowsingTransforms { 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 +388,7 @@ public class SafeBrowsingTransforms { }, 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 +412,7 @@ public class SafeBrowsingTransforms { .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 +428,6 @@ public class SafeBrowsingTransforms { 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 +444,10 @@ public class SafeBrowsingTransforms { 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 081f84710..9416d48f9 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 @@ package google.registry.beam.spec11; 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 dagger.Provides; 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 jakarta.inject.Singleton; 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 class Spec11Pipeline implements Serializable { 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 @@ public class Spec11Pipeline implements Serializable { 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 @@ public class Spec11Pipeline implements Serializable { } @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 54bbf6eb4..a012f4330 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 static org.mockito.Mockito.withSettings; 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 @@ class SafeBrowsingTransformsTest { @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 @@ class SafeBrowsingTransformsTest { 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 @@ class SafeBrowsingTransformsTest { @Test void testSuccess_noBadDomains() throws Exception { + PCollectionView prefixView = createPrefixView(); ImmutableList domainNameInfos = ImmutableList.of( createDomainNameInfo("hello_kitty.dev"), @@ -158,12 +188,80 @@ class SafeBrowsingTransformsTest { 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,13 +269,65 @@ class SafeBrowsingTransformsTest { 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 (Object 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; + } + /** * Returns a {@link CloseableHttpResponse} containing either positive (threat found) or negative * (no threat) API examples based on the request data. 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 b618bc483..14709fb59 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.collect.Streams; 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 @@ class Spec11PipelineTest { .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 @@ class Spec11PipelineTest { @Test void testSuccess_fullSqlPipeline() throws Exception { setupCloudSql(); - EvaluateSafeBrowsingFn safeBrowsingFn = - new EvaluateSafeBrowsingFn( - SAFE_BROWSING_API_KEY, - new Retrier(new FakeSleeper(fakeClock), 1), - fakeClock, - Suppliers.ofInstance(mockHttpClient)); when(mockHttpClient.execute(any(HttpPost.class))).thenAnswer(new HttpResponder()); - Spec11Pipeline spec11Pipeline = new Spec11Pipeline(options, safeBrowsingFn); + Spec11Pipeline spec11Pipeline = + new Spec11Pipeline( + options, + fakeClock, + new Retrier(new FakeSleeper(fakeClock), 1), + Suppliers.ofInstance(mockHttpClient)); 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 b89b6a31e..d7ebb17a2 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