1 /* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
2 /* This Source Code Form is subject to the terms of the Mozilla Public
3 * License, v. 2.0. If a copy of the MPL was not distributed with this
4 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
5
6 #include "mozilla/EndianUtils.h"
7 #include "ProtocolParser.h"
8
9 #include "Common.h"
10
11 typedef FetchThreatListUpdatesResponse_ListUpdateResponse ListUpdateResponse;
12
InitUpdateResponse(ListUpdateResponse * aUpdateResponse,ThreatType aThreatType,const nsACString & aState,const nsACString & aChecksum,bool isFullUpdate,const nsTArray<uint32_t> & aFixedLengthPrefixes,bool aDoPrefixEncoding)13 static bool InitUpdateResponse(ListUpdateResponse* aUpdateResponse,
14 ThreatType aThreatType, const nsACString& aState,
15 const nsACString& aChecksum, bool isFullUpdate,
16 const nsTArray<uint32_t>& aFixedLengthPrefixes,
17 bool aDoPrefixEncoding) {
18 aUpdateResponse->set_threat_type(aThreatType);
19 aUpdateResponse->set_new_client_state(aState.BeginReading(), aState.Length());
20 aUpdateResponse->mutable_checksum()->set_sha256(aChecksum.BeginReading(),
21 aChecksum.Length());
22 aUpdateResponse->set_response_type(isFullUpdate
23 ? ListUpdateResponse::FULL_UPDATE
24 : ListUpdateResponse::PARTIAL_UPDATE);
25
26 auto additions = aUpdateResponse->mutable_additions()->Add();
27
28 if (!aDoPrefixEncoding) {
29 additions->set_compression_type(RAW);
30 auto rawHashes = additions->mutable_raw_hashes();
31 rawHashes->set_prefix_size(4);
32 auto prefixes = rawHashes->mutable_raw_hashes();
33 for (auto p : aFixedLengthPrefixes) {
34 char buffer[4];
35 NativeEndian::copyAndSwapToBigEndian(buffer, &p, 1);
36 prefixes->append(buffer, 4);
37 }
38 return true;
39 }
40
41 if (1 != aFixedLengthPrefixes.Length()) {
42 printf("This function only supports single value encoding.\n");
43 return false;
44 }
45
46 uint32_t firstValue = aFixedLengthPrefixes[0];
47 additions->set_compression_type(RICE);
48 auto riceHashes = additions->mutable_rice_hashes();
49 riceHashes->set_first_value(firstValue);
50 riceHashes->set_num_entries(0);
51
52 return true;
53 }
54
DumpBinary(const nsACString & aBinary)55 static void DumpBinary(const nsACString& aBinary) {
56 nsCString s;
57 for (size_t i = 0; i < aBinary.Length(); i++) {
58 s.AppendPrintf("\\x%.2X", (uint8_t)aBinary[i]);
59 }
60 printf("%s\n", s.get());
61 }
62
TEST(UrlClassifierProtocolParser,UpdateWait)63 TEST(UrlClassifierProtocolParser, UpdateWait)
64 {
65 // Top level response which contains a list of update response
66 // for different lists.
67 FetchThreatListUpdatesResponse response;
68
69 auto r = response.mutable_list_update_responses()->Add();
70 InitUpdateResponse(r, SOCIAL_ENGINEERING_PUBLIC, nsCString("sta\x00te", 6),
71 nsCString("check\x0sum", 9), true, {0, 1, 2, 3},
72 false /* aDoPrefixEncoding */);
73
74 // Set min wait duration.
75 auto minWaitDuration = response.mutable_minimum_wait_duration();
76 minWaitDuration->set_seconds(8);
77 minWaitDuration->set_nanos(1 * 1000000000);
78
79 std::string s;
80 response.SerializeToString(&s);
81
82 DumpBinary(nsCString(s.c_str(), s.length()));
83
84 ProtocolParser* p = new ProtocolParserProtobuf();
85 p->AppendStream(nsCString(s.c_str(), s.length()));
86 p->End();
87 ASSERT_EQ(p->UpdateWaitSec(), 9u);
88 delete p;
89 }
90
TEST(UrlClassifierProtocolParser,SingleValueEncoding)91 TEST(UrlClassifierProtocolParser, SingleValueEncoding)
92 {
93 // Top level response which contains a list of update response
94 // for different lists.
95 FetchThreatListUpdatesResponse response;
96
97 auto r = response.mutable_list_update_responses()->Add();
98
99 const char* expectedPrefix = "\x00\x01\x02\x00";
100 if (!InitUpdateResponse(r, SOCIAL_ENGINEERING_PUBLIC,
101 nsCString("sta\x00te", 6),
102 nsCString("check\x0sum", 9), true,
103 // As per spec, we should interpret the prefix as
104 // uint32 in little endian before encoding.
105 {LittleEndian::readUint32(expectedPrefix)},
106 true /* aDoPrefixEncoding */)) {
107 printf("Failed to initialize update response.");
108 ASSERT_TRUE(false);
109 return;
110 }
111
112 // Set min wait duration.
113 auto minWaitDuration = response.mutable_minimum_wait_duration();
114 minWaitDuration->set_seconds(8);
115 minWaitDuration->set_nanos(1 * 1000000000);
116
117 std::string s;
118 response.SerializeToString(&s);
119
120 // Feed data to the protocol parser.
121 ProtocolParser* p = new ProtocolParserProtobuf();
122 p->SetRequestedTables({nsCString("googpub-phish-proto")});
123 p->AppendStream(nsCString(s.c_str(), s.length()));
124 p->End();
125
126 const TableUpdateArray& tus = p->GetTableUpdates();
127 RefPtr<const TableUpdateV4> tuv4 = TableUpdate::Cast<TableUpdateV4>(tus[0]);
128 auto& prefixMap = tuv4->Prefixes();
129 for (auto iter = prefixMap.ConstIter(); !iter.Done(); iter.Next()) {
130 // This prefix map should contain only a single 4-byte prefixe.
131 ASSERT_EQ(iter.Key(), 4u);
132
133 // The fixed-length prefix string from ProtocolParser should
134 // exactly match the expected prefix string.
135 nsCString* prefix = iter.UserData();
136 ASSERT_TRUE(prefix->Equals(nsCString(expectedPrefix, 4)));
137 }
138
139 delete p;
140 }
141