mirrors/nomulus
1
0
Fork 0
mirror of https://github.com/google/nomulus synced 2026-10-21 22:02:44 +00:00
Code Issues Releases Wiki Activity

Reduce SafeBrowsing API calls by prefix hashing (#3227)

Browse Source 
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).
This commit is contained in:
gbrodman
2026-10-15 20:51:38 +00:00
committed by GitHub
parent 19b20b7313
commit 052c083bcb
5 changed files with 444 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/src/main/java/google/registry/beam/spec11/SafeBrowsingTransforms.java
+222 -30
View File
@@ -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.
*
* <p>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<String, int[]> {
private final Retrier retrier;
private final Supplier<CloseableHttpClient> closeableHttpClientSupplier;
FetchThreatListPrefixesFn(Retrier retrier) {
this(retrier, (Supplier<CloseableHttpClient> & Serializable) HttpClients::createDefault);
}
@VisibleForTesting
FetchThreatListPrefixesFn(Retrier retrier, Supplier<CloseableHttpClient> clientSupplier) {
this.retrier = retrier;
this.closeableHttpClientSupplier = clientSupplier;
}
@ProcessElement
public void processElement(@Element String apiKey, OutputReceiver<int[]> 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.
*
* <p>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.
*
* <p>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<Integer>}) 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.
*
* <p>Refer to the Lookup API documentation for the request/response format and other details.
* <p>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 <a href=https://developers.google.com/safe-browsing/v4/lookup-api>Lookup API</a>
*/
@@ -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<String, DomainNameInfo> 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<int[]> 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<int[]> prefixesView) {
this(
apiKey,
retrier,
clock,
(Supplier<CloseableHttpClient> & Serializable) HttpClients::createDefault);
(Supplier<CloseableHttpClient> & Serializable) HttpClients::createDefault,
prefixesView);
}
/**
@@ -121,11 +304,16 @@ public class SafeBrowsingTransforms {
*/
@VisibleForTesting
EvaluateSafeBrowsingFn(
String apiKey, Retrier retrier, Clock clock, Supplier<CloseableHttpClient> clientSupplier) {
String apiKey,
Retrier retrier,
Clock clock,
Supplier<CloseableHttpClient> clientSupplier,
PCollectionView<int[]> 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<KV<DomainNameInfo, ThreatMatch>> 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,
core/src/main/java/google/registry/beam/spec11/Spec11Pipeline.java
+46 -14
View File
@@ -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<CloseableHttpClient> clientSupplier;
Spec11Pipeline(Spec11PipelineOptions options, EvaluateSafeBrowsingFn safeBrowsingFn) {
Spec11Pipeline(Spec11PipelineOptions options, Clock clock, Retrier retrier) {
this(
options,
clock,
retrier,
(Supplier<CloseableHttpClient> & Serializable) HttpClients::createDefault);
}
@VisibleForTesting
Spec11Pipeline(
Spec11PipelineOptions options,
Clock clock,
Retrier retrier,
Supplier<CloseableHttpClient> 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<int[]> 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<DomainNameInfo> domains = readFromCloudSql(pipeline);
EvaluateSafeBrowsingFn safeBrowsingFn =
new EvaluateSafeBrowsingFn(
options.getSafeBrowsingApiKey(), retrier, clock, clientSupplier, prefixView);
PCollection<KV<DomainNameInfo, ThreatMatch>> 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);
}
}
core/src/test/java/google/registry/beam/spec11/SafeBrowsingTransformsTest.java
+164 -14
View File
@@ -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<DomainNameInfo, ThreatMatch> 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<int[]> 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<int[]> 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<int[]> prefixView = createPrefixView();
ImmutableList<DomainNameInfo> domainNameInfos =
ImmutableList.of(
createDomainNameInfo("111.com"),
@@ -137,7 +166,7 @@ class SafeBrowsingTransformsTest {
PCollection<KV<DomainNameInfo, ThreatMatch>> 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<int[]> prefixView = createPrefixView();
ImmutableList<DomainNameInfo> domainNameInfos =
ImmutableList.of(
createDomainNameInfo("hello_kitty.dev"),
@@ -158,12 +188,80 @@ class SafeBrowsingTransformsTest {
PCollection<KV<DomainNameInfo, ThreatMatch>> 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<int[]> prefixView =
pipeline
.apply("Create colliding prefix", Create.of(new int[] {prefixInt}))
.apply("View colliding prefix", View.asSingleton());
ImmutableList<DomainNameInfo> domainNameInfos =
ImmutableList.of(createDomainNameInfo("clean-domain.com"));
PCollection<KV<DomainNameInfo, ThreatMatch>> 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<int[]> 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<CloseableHttpResponse>, 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<String, String> 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.
core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java
+10 -12
View File
@@ -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();
core/src/test/resources/google/registry/beam/spec11/test_output.txt
+2 -2
View File
@@ -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"}
{"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"}