fune/toolkit/components/url-classifier/tests/gtest/TestProtocolParser.cpp

/* -*- 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;
}