Skip to content

Reduce SafeBrowsing API calls by prefix hashing #3227

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gbrodman wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Reduce SafeBrowsing API calls by prefix hashing #3227

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
249 changes: 219 additions & 30 deletions core/src/main/java/google/registry/beam/spec11/SafeBrowsingTransforms.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,36 +14,42 @@

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;
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;
Expand All @@ -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<String> 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.
*
* <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 (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.
*
* <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.
*/
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.
*
* <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>
*/
Expand All @@ -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<String, DomainNameInfo> domainNameInfoBuffer =
new LinkedHashMap<>(BATCH_SIZE);
Expand All @@ -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<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.
*
Expand All @@ -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<int[]> prefixesView) {
this(
apiKey,
retrier,
clock,
(Supplier<CloseableHttpClient> & Serializable) HttpClients::createDefault);
(Supplier<CloseableHttpClient> & Serializable) HttpClients::createDefault,
prefixesView);
}

/**
Expand All @@ -121,11 +301,16 @@ static class EvaluateSafeBrowsingFn
*/
@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. */
Expand All @@ -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<KV<DomainNameInfo, ThreatMatch>> results = evaluateAndFlush();
Expand All @@ -173,10 +370,11 @@ private ImmutableSet<KV<DomainNameInfo, ThreatMatch>> 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(
() -> {
Expand All @@ -187,7 +385,7 @@ private ImmutableSet<KV<DomainNameInfo, ThreatMatch>> 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
Expand All @@ -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));
Expand All @@ -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));
}
Expand All @@ -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,
Expand Down
Loading
Loading