From eb79babf32327ef7011de7671daef622a062ddf6 Mon Sep 17 00:00:00 2001 From: Vasu Nori Date: Fri, 20 Aug 2021 13:53:45 -0700 Subject: [PATCH] Change checksum algorithm to CRC32C. And localize checksum handling to one file. --- .../client/ChecksumEnforcingInputStream.java | 18 +++++----- .../v1/client/EndToEndChecksumHandler.java | 33 ++++++++++--------- .../google/datastore/v1/client/RemoteRpc.java | 3 +- .../ChecksumEnforcingInputStreamTest.java | 11 ++----- .../datastore/v1/client/RemoteRpcTest.java | 2 +- 5 files changed, 31 insertions(+), 36 deletions(-) diff --git a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/ChecksumEnforcingInputStream.java b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/ChecksumEnforcingInputStream.java index 2b9674917..501d276b3 100644 --- a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/ChecksumEnforcingInputStream.java +++ b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/ChecksumEnforcingInputStream.java @@ -17,6 +17,7 @@ import com.google.api.client.http.HttpResponse; import com.google.common.annotations.VisibleForTesting; +import com.google.common.hash.Hasher; import java.io.IOException; import java.io.InputStream; import java.security.MessageDigest; @@ -24,20 +25,18 @@ /** This class provides End-to-End Checksum API for http protocol. */ class ChecksumEnforcingInputStream extends InputStream { private final InputStream delegate; - private final MessageDigest messageDigest; + private final EndToEndChecksumHandler endToEndChecksumHandler; private final String expectedChecksum; - ChecksumEnforcingInputStream( - InputStream originalInputStream, HttpResponse response, MessageDigest digest) { - this(originalInputStream, EndToEndChecksumHandler.getChecksumHeader(response), digest); + ChecksumEnforcingInputStream(InputStream originalInputStream, HttpResponse response) { + this(originalInputStream, EndToEndChecksumHandler.getChecksumHeader(response)); } @VisibleForTesting - ChecksumEnforcingInputStream( - InputStream originalInputStream, String checksum, MessageDigest digest) { + ChecksumEnforcingInputStream(InputStream originalInputStream, String checksum) { delegate = originalInputStream; expectedChecksum = checksum; - messageDigest = digest; + endToEndChecksumHandler = new EndToEndChecksumHandler(); } @Override @@ -76,11 +75,10 @@ public int read(byte[] b, int off, int len) throws IOException { if (len <= 0) return 0; int i = delegate.read(b, off, len); if (i > 0) { - messageDigest.update(b, off, i); + endToEndChecksumHandler.update(b, off, i); } else { // no more payload to read. compute checksum and verify - if (!expectedChecksum.equalsIgnoreCase( - com.google.common.io.BaseEncoding.base16().encode(messageDigest.digest()))) { + if (!expectedChecksum.equalsIgnoreCase(endToEndChecksumHandler.hash())) { throw new IOException("possible memory corruption on payload detected"); } } diff --git a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/EndToEndChecksumHandler.java b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/EndToEndChecksumHandler.java index 06d08d499..b20ba4b3b 100644 --- a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/EndToEndChecksumHandler.java +++ b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/EndToEndChecksumHandler.java @@ -16,8 +16,9 @@ package com.google.datastore.v1.client; import com.google.api.client.http.HttpResponse; -import java.security.MessageDigest; -import java.security.NoSuchAlgorithmException; +import com.google.common.hash.HashCode; +import com.google.common.hash.Hasher; +import com.google.common.hash.Hashing; /** This class provides End-to-End Checksum API for http protocol. */ class EndToEndChecksumHandler { @@ -25,22 +26,25 @@ class EndToEndChecksumHandler { static final String HTTP_REQUEST_CHECKSUM_HEADER = "x-request-checksum-348659783"; /** The checksum http header on http responses */ static final String HTTP_RESPONSE_CHECKSUM_HEADER = "x-response-checksum-348659783"; - /** Algorithm used for checksum */ - private static final String MD5 = "MD5"; + + final Hasher hasher = EndToEndChecksumHandler.getNewHasher(); /** * Create and return checksum as a string value for the input 'bytes'. * * @param bytes raw message for which the checksum is being computed * @return computed checksum as a hex string - * @throws RuntimeException if MD5 Algorithm is not found in the VM */ static String computeChecksum(byte[] bytes) { if (bytes == null || (bytes.length == 0)) { return null; } - return com.google.common.io.BaseEncoding.base16() - .encode(getMessageDigestInstance().digest(bytes)); + HashCode hc = getNewHasher().putBytes(bytes).hash(); + return hc.toString(); + } + + private static Hasher getNewHasher() { + return Hashing.crc32c().newHasher(); } /** @@ -58,14 +62,6 @@ static boolean validateChecksum(String checksum, byte[] bytes) { && checksum.equalsIgnoreCase(computeChecksum(bytes)); } - static MessageDigest getMessageDigestInstance() { - try { - return MessageDigest.getInstance(MD5); - } catch (NoSuchAlgorithmException e) { - throw new RuntimeException("MD5 algorithm is not found when computing checksum!"); - } - } - static boolean hasChecksumHeader(HttpResponse response) { String checksum = getChecksumHeader(response); return checksum != null && !checksum.isEmpty(); @@ -74,4 +70,11 @@ static boolean hasChecksumHeader(HttpResponse response) { static String getChecksumHeader(HttpResponse response) { return response.getHeaders().getFirstHeaderStringValue(HTTP_RESPONSE_CHECKSUM_HEADER); } + + void update(byte[] bytes, int off, int len) { + hasher.putBytes(bytes, off, len); + } + String hash() { + return hasher.hash().toString(); + } } diff --git a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/RemoteRpc.java b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/RemoteRpc.java index 2c9d4d34e..321eea72a 100644 --- a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/RemoteRpc.java +++ b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/RemoteRpc.java @@ -109,8 +109,7 @@ public InputStream call(String methodName, MessageLite request) throws Datastore } InputStream inputStream = httpResponse.getContent(); return enableE2EChecksum && EndToEndChecksumHandler.hasChecksumHeader(httpResponse) - ? new ChecksumEnforcingInputStream( - inputStream, httpResponse, EndToEndChecksumHandler.getMessageDigestInstance()) + ? new ChecksumEnforcingInputStream(inputStream, httpResponse) : inputStream; } catch (SocketTimeoutException e) { throw makeException(url, methodName, Code.DEADLINE_EXCEEDED, "Deadline exceeded", e); diff --git a/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/ChecksumEnforcingInputStreamTest.java b/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/ChecksumEnforcingInputStreamTest.java index 47592f987..f90c99cae 100644 --- a/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/ChecksumEnforcingInputStreamTest.java +++ b/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/ChecksumEnforcingInputStreamTest.java @@ -21,7 +21,6 @@ import java.io.ByteArrayInputStream; import java.io.IOException; -import java.security.MessageDigest; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -29,8 +28,6 @@ /** Test for {@link ChecksumEnforcingInputStream}. */ @RunWith(JUnit4.class) public class ChecksumEnforcingInputStreamTest { - private final MessageDigest digest = EndToEndChecksumHandler.getMessageDigestInstance(); - public void test(int payloadSize) throws Exception { // read 1000 bytes at a time // Since checksum should be correct, do not expect IOException @@ -40,7 +37,7 @@ public void test(int payloadSize) throws Exception { // do nothing with the bytes read } } catch (IOException e) { - fail("checksum verification failed!"); + fail("checksum verification failed! " + e.getMessage()); } } @@ -67,8 +64,7 @@ public void read_withInvalidChecksum() { try (ChecksumEnforcingInputStream instance = new ChecksumEnforcingInputStream( new ByteArrayInputStream("hello there".getBytes(UTF_8)), - "this checksum is invalid", - digest)) { + "this checksum is invalid")) { byte[] buf = new byte[1000]; while (instance.read(buf, 0, 1000) != -1) { // do nothing with the bytes read @@ -103,7 +99,6 @@ private ChecksumEnforcingInputStream setUpData(int payloadSize) throws Exception } byte[] bytes = payload.getBytes(UTF_8); String expectedChecksum = EndToEndChecksumHandler.computeChecksum(bytes); - return new ChecksumEnforcingInputStream( - new ByteArrayInputStream(bytes), expectedChecksum, digest); + return new ChecksumEnforcingInputStream(new ByteArrayInputStream(bytes), expectedChecksum); } } diff --git a/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/RemoteRpcTest.java b/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/RemoteRpcTest.java index b5f82369b..ebcb12396 100644 --- a/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/RemoteRpcTest.java +++ b/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/RemoteRpcTest.java @@ -174,7 +174,7 @@ public void testHttpHeaders_expectE2eChecksumHeader() throws IOException { httpRequest .getHeaders() .getFirstHeaderStringValue(EndToEndChecksumHandler.HTTP_REQUEST_CHECKSUM_HEADER); - assertEquals(32, header.length()); + assertEquals(8, header.length()); } @Test