1 files changed, 140 insertions, 0 deletions
diff --git a/toolkit/components/url-classifier/tests/gtest/TestProtocolParser.cpp b/toolkit/components/url-classifier/tests/gtest/TestProtocolParser.cpp
new file mode 100644
index 0000000000..2155140b4d
--- /dev/null
+++ b/toolkit/components/url-classifier/tests/gtest/TestProtocolParser.cpp
@@ -0,0 +1,140 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include "mozilla/EndianUtils.h"
+#include "ProtocolParser.h"
+
+#include "Common.h"
+
+typedef FetchThreatListUpdatesResponse_ListUpdateResponse ListUpdateResponse;
+
+static bool InitUpdateResponse(ListUpdateResponse* aUpdateResponse,
+                               ThreatType aThreatType, const nsACString& aState,
+                               const nsACString& aChecksum, bool isFullUpdate,
+                               const nsTArray<uint32_t>& aFixedLengthPrefixes,
+                               bool aDoPrefixEncoding) {
+  aUpdateResponse->set_threat_type(aThreatType);
+  aUpdateResponse->set_new_client_state(aState.BeginReading(), aState.Length());
+  aUpdateResponse->mutable_checksum()->set_sha256(aChecksum.BeginReading(),
+                                                  aChecksum.Length());
+  aUpdateResponse->set_response_type(isFullUpdate
+                                         ? ListUpdateResponse::FULL_UPDATE
+                                         : ListUpdateResponse::PARTIAL_UPDATE);
+
+  auto additions = aUpdateResponse->mutable_additions()->Add();
+
+  if (!aDoPrefixEncoding) {
+    additions->set_compression_type(RAW);
+    auto rawHashes = additions->mutable_raw_hashes();
+    rawHashes->set_prefix_size(4);
+    auto prefixes = rawHashes->mutable_raw_hashes();
+    for (auto p : aFixedLengthPrefixes) {
+      char buffer[4];
+      mozilla::NativeEndian::copyAndSwapToBigEndian(buffer, &p, 1);
+      prefixes->append(buffer, 4);
+    }
+    return true;
+  }
+
+  if (1 != aFixedLengthPrefixes.Length()) {
+    printf("This function only supports single value encoding.\n");
+    return false;
+  }
+
+  uint32_t firstValue = aFixedLengthPrefixes[0];
+  additions->set_compression_type(RICE);
+  auto riceHashes = additions->mutable_rice_hashes();
+  riceHashes->set_first_value(firstValue);
+  riceHashes->set_num_entries(0);
+
+  return true;
+}
+
+static void DumpBinary(const nsACString& aBinary) {
+  nsCString s;
+  for (size_t i = 0; i < aBinary.Length(); i++) {
+    s.AppendPrintf("\\x%.2X", (uint8_t)aBinary[i]);
+  }
+  printf("%s\n", s.get());
+}
+
+TEST(UrlClassifierProtocolParser, UpdateWait)
+{
+  // Top level response which contains a list of update response
+  // for different lists.
+  FetchThreatListUpdatesResponse response;
+
+  auto r = response.mutable_list_update_responses()->Add();
+  InitUpdateResponse(r, SOCIAL_ENGINEERING_PUBLIC, nsCString("sta\x00te", 6),
+                     nsCString("check\x0sum", 9), true, {0, 1, 2, 3},
+                     false /* aDoPrefixEncoding */);
+
+  // Set min wait duration.
+  auto minWaitDuration = response.mutable_minimum_wait_duration();
+  minWaitDuration->set_seconds(8);
+  minWaitDuration->set_nanos(1 * 1000000000);
+
+  std::string s;
+  response.SerializeToString(&s);
+
+  DumpBinary(nsCString(s.c_str(), s.length()));
+
+  ProtocolParser* p = new ProtocolParserProtobuf();
+  p->AppendStream(nsCString(s.c_str(), s.length()));
+  p->End();
+  ASSERT_EQ(p->UpdateWaitSec(), 9u);
+  delete p;
+}
+
+TEST(UrlClassifierProtocolParser, SingleValueEncoding)
+{
+  // Top level response which contains a list of update response
+  // for different lists.
+  FetchThreatListUpdatesResponse response;
+
+  auto r = response.mutable_list_update_responses()->Add();
+
+  const char* expectedPrefix = "\x00\x01\x02\x00";
+  if (!InitUpdateResponse(r, SOCIAL_ENGINEERING_PUBLIC,
+                          nsCString("sta\x00te", 6),
+                          nsCString("check\x0sum", 9), true,
+                          // As per spec, we should interpret the prefix as
+                          // uint32 in little endian before encoding.
+                          {mozilla::LittleEndian::readUint32(expectedPrefix)},
+                          true /* aDoPrefixEncoding */)) {
+    printf("Failed to initialize update response.");
+    ASSERT_TRUE(false);
+    return;
+  }
+
+  // Set min wait duration.
+  auto minWaitDuration = response.mutable_minimum_wait_duration();
+  minWaitDuration->set_seconds(8);
+  minWaitDuration->set_nanos(1 * 1000000000);
+
+  std::string s;
+  response.SerializeToString(&s);
+
+  // Feed data to the protocol parser.
+  ProtocolParser* p = new ProtocolParserProtobuf();
+  p->SetRequestedTables({nsCString("googpub-phish-proto")});
+  p->AppendStream(nsCString(s.c_str(), s.length()));
+  p->End();
+
+  const TableUpdateArray& tus = p->GetTableUpdates();
+  RefPtr<const TableUpdateV4> tuv4 = TableUpdate::Cast<TableUpdateV4>(tus[0]);
+  auto& prefixMap = tuv4->Prefixes();
+  for (const auto& entry : prefixMap) {
+    // This prefix map should contain only a single 4-byte prefixe.
+    ASSERT_EQ(entry.GetKey(), 4u);
+
+    // The fixed-length prefix string from ProtocolParser should
+    // exactly match the expected prefix string.
+    nsCString* prefix = entry.GetWeak();
+    ASSERT_TRUE(prefix->Equals(nsCString(expectedPrefix, 4)));
+  }
+
+  delete p;
+}