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