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 "LookupCacheV4.h"
7 #include "HashStore.h"
8 #include "mozilla/Unused.h"
9 #include <string>
10
11 // MOZ_LOG=UrlClassifierDbService:5
12 extern mozilla::LazyLogModule gUrlClassifierDbServiceLog;
13 #define LOG(args) MOZ_LOG(gUrlClassifierDbServiceLog, mozilla::LogLevel::Debug, args)
14 #define LOG_ENABLED() MOZ_LOG_TEST(gUrlClassifierDbServiceLog, mozilla::LogLevel::Debug)
15
16 #define METADATA_SUFFIX NS_LITERAL_CSTRING(".metadata")
17
18 namespace mozilla {
19 namespace safebrowsing {
20
21 const int LookupCacheV4::VER = 4;
22
23 // Prefixes coming from updates and VLPrefixSet are both stored in the HashTable
24 // where the (key, value) pair is a prefix size and a lexicographic-sorted string.
25 // The difference is prefixes from updates use std:string(to avoid additional copies)
26 // and prefixes from VLPrefixSet use nsCString.
27 // This class provides a common interface for the partial update algorithm to make it
28 // easier to operate on two different kind prefix string map..
29 class VLPrefixSet
30 {
31 public:
32 explicit VLPrefixSet(const PrefixStringMap& aMap);
33 explicit VLPrefixSet(const TableUpdateV4::PrefixStdStringMap& aMap);
34
35 // This function will merge the prefix map in VLPrefixSet to aPrefixMap.
36 void Merge(PrefixStringMap& aPrefixMap);
37
38 // Find the smallest string from the map in VLPrefixSet.
39 bool GetSmallestPrefix(nsDependentCSubstring& aOutString);
40
41 // Return the number of prefixes in the map
Count() const42 uint32_t Count() const { return mCount; }
43
44 private:
45 // PrefixString structure contains a lexicographic-sorted string with
46 // a |pos| variable to indicate which substring we are pointing to right now.
47 // |pos| increases each time GetSmallestPrefix finds the smallest string.
48 struct PrefixString {
PrefixStringmozilla::safebrowsing::VLPrefixSet::PrefixString49 PrefixString(const nsACString& aStr, uint32_t aSize)
50 : pos(0)
51 , size(aSize)
52 {
53 data.Rebind(aStr.BeginReading(), aStr.Length());
54 }
55
getmozilla::safebrowsing::VLPrefixSet::PrefixString56 const char* get() {
57 return pos < data.Length() ? data.BeginReading() + pos : nullptr;
58 }
nextmozilla::safebrowsing::VLPrefixSet::PrefixString59 void next() { pos += size; }
remainingmozilla::safebrowsing::VLPrefixSet::PrefixString60 uint32_t remaining() { return data.Length() - pos; }
61
62 nsDependentCSubstring data;
63 uint32_t pos;
64 uint32_t size;
65 };
66
67 nsClassHashtable<nsUint32HashKey, PrefixString> mMap;
68 uint32_t mCount;
69 };
70
71 nsresult
Init()72 LookupCacheV4::Init()
73 {
74 mVLPrefixSet = new VariableLengthPrefixSet();
75 nsresult rv = mVLPrefixSet->Init(mTableName);
76 NS_ENSURE_SUCCESS(rv, rv);
77
78 return NS_OK;
79 }
80
81 nsresult
Has(const Completion & aCompletion,bool * aHas,bool * aComplete)82 LookupCacheV4::Has(const Completion& aCompletion,
83 bool* aHas, bool* aComplete)
84 {
85 *aHas = false;
86
87 uint32_t length = 0;
88 nsDependentCSubstring fullhash;
89 fullhash.Rebind((const char *)aCompletion.buf, COMPLETE_SIZE);
90
91 nsresult rv = mVLPrefixSet->Matches(fullhash, &length);
92 NS_ENSURE_SUCCESS(rv, rv);
93
94 *aHas = length >= PREFIX_SIZE;
95 *aComplete = length == COMPLETE_SIZE;
96
97 if (LOG_ENABLED()) {
98 uint32_t prefix = aCompletion.ToUint32();
99 LOG(("Probe in V4 %s: %X, found %d, complete %d", mTableName.get(),
100 prefix, *aHas, *aComplete));
101 }
102
103 return NS_OK;
104 }
105
106 nsresult
Build(PrefixStringMap & aPrefixMap)107 LookupCacheV4::Build(PrefixStringMap& aPrefixMap)
108 {
109 return mVLPrefixSet->SetPrefixes(aPrefixMap);
110 }
111
112 nsresult
GetPrefixes(PrefixStringMap & aPrefixMap)113 LookupCacheV4::GetPrefixes(PrefixStringMap& aPrefixMap)
114 {
115 return mVLPrefixSet->GetPrefixes(aPrefixMap);
116 }
117
118 nsresult
ClearPrefixes()119 LookupCacheV4::ClearPrefixes()
120 {
121 // Clear by seting a empty map
122 PrefixStringMap map;
123 return mVLPrefixSet->SetPrefixes(map);
124 }
125
126 nsresult
StoreToFile(nsIFile * aFile)127 LookupCacheV4::StoreToFile(nsIFile* aFile)
128 {
129 return mVLPrefixSet->StoreToFile(aFile);
130 }
131
132 nsresult
LoadFromFile(nsIFile * aFile)133 LookupCacheV4::LoadFromFile(nsIFile* aFile)
134 {
135 nsresult rv = mVLPrefixSet->LoadFromFile(aFile);
136 if (NS_FAILED(rv)) {
137 return rv;
138 }
139
140 nsCString state, checksum;
141 rv = LoadMetadata(state, checksum);
142 if (NS_FAILED(rv)) {
143 return rv;
144 }
145
146 rv = VerifyChecksum(checksum);
147 Telemetry::Accumulate(Telemetry::URLCLASSIFIER_VLPS_LOAD_CORRUPT,
148 rv == NS_ERROR_FILE_CORRUPTED);
149
150 return rv;
151 }
152
153 size_t
SizeOfPrefixSet()154 LookupCacheV4::SizeOfPrefixSet()
155 {
156 return mVLPrefixSet->SizeOfIncludingThis(moz_malloc_size_of);
157 }
158
159 static void
AppendPrefixToMap(PrefixStringMap & prefixes,nsDependentCSubstring & prefix)160 AppendPrefixToMap(PrefixStringMap& prefixes, nsDependentCSubstring& prefix)
161 {
162 if (!prefix.Length()) {
163 return;
164 }
165
166 nsCString* prefixString = prefixes.LookupOrAdd(prefix.Length());
167 prefixString->Append(prefix.BeginReading(), prefix.Length());
168 }
169
170 // Read prefix into a buffer and also update the hash which
171 // keeps track of the checksum
172 static void
UpdateChecksum(nsICryptoHash * aCrypto,const nsACString & aPrefix)173 UpdateChecksum(nsICryptoHash* aCrypto, const nsACString& aPrefix)
174 {
175 MOZ_ASSERT(aCrypto);
176 aCrypto->Update(reinterpret_cast<uint8_t*>(const_cast<char*>(
177 aPrefix.BeginReading())),
178 aPrefix.Length());
179 }
180
181 // Please see https://bug1287058.bmoattachments.org/attachment.cgi?id=8795366
182 // for detail about partial update algorithm.
183 nsresult
ApplyUpdate(TableUpdateV4 * aTableUpdate,PrefixStringMap & aInputMap,PrefixStringMap & aOutputMap)184 LookupCacheV4::ApplyUpdate(TableUpdateV4* aTableUpdate,
185 PrefixStringMap& aInputMap,
186 PrefixStringMap& aOutputMap)
187 {
188 MOZ_ASSERT(aOutputMap.IsEmpty());
189
190 nsCOMPtr<nsICryptoHash> crypto;
191 nsresult rv = InitCrypto(crypto);
192 if (NS_FAILED(rv)) {
193 return rv;
194 }
195
196 // oldPSet contains prefixes we already have or we just merged last round.
197 // addPSet contains prefixes stored in tableUpdate which should be merged with oldPSet.
198 VLPrefixSet oldPSet(aInputMap);
199 VLPrefixSet addPSet(aTableUpdate->Prefixes());
200
201 // RemovalIndiceArray is a sorted integer array indicating the index of prefix we should
202 // remove from the old prefix set(according to lexigraphic order).
203 // |removalIndex| is the current index of RemovalIndiceArray.
204 // |numOldPrefixPicked| is used to record how many prefixes we picked from the old map.
205 TableUpdateV4::RemovalIndiceArray& removalArray = aTableUpdate->RemovalIndices();
206 uint32_t removalIndex = 0;
207 int32_t numOldPrefixPicked = -1;
208
209 nsDependentCSubstring smallestOldPrefix;
210 nsDependentCSubstring smallestAddPrefix;
211
212 bool isOldMapEmpty = false, isAddMapEmpty = false;
213
214 // This is used to avoid infinite loop for partial update algorithm.
215 // The maximum loops will be the number of old prefixes plus the number of add prefixes.
216 int32_t index = oldPSet.Count() + addPSet.Count() + 1;
217 for(;index > 0; index--) {
218 // Get smallest prefix from the old prefix set if we don't have one
219 if (smallestOldPrefix.IsEmpty() && !isOldMapEmpty) {
220 isOldMapEmpty = !oldPSet.GetSmallestPrefix(smallestOldPrefix);
221 }
222
223 // Get smallest prefix from add prefix set if we don't have one
224 if (smallestAddPrefix.IsEmpty() && !isAddMapEmpty) {
225 isAddMapEmpty = !addPSet.GetSmallestPrefix(smallestAddPrefix);
226 }
227
228 bool pickOld;
229
230 // If both prefix sets are not empty, then compare to find the smaller one.
231 if (!isOldMapEmpty && !isAddMapEmpty) {
232 if (smallestOldPrefix == smallestAddPrefix) {
233 LOG(("Add prefix should not exist in the original prefix set."));
234 Telemetry::Accumulate(Telemetry::URLCLASSIFIER_UPDATE_ERROR_TYPE,
235 DUPLICATE_PREFIX);
236 return NS_ERROR_FAILURE;
237 }
238
239 // Compare the smallest string in old prefix set and add prefix set,
240 // merge the smaller one into new map to ensure merged string still
241 // follows lexigraphic order.
242 pickOld = smallestOldPrefix < smallestAddPrefix;
243 } else if (!isOldMapEmpty && isAddMapEmpty) {
244 pickOld = true;
245 } else if (isOldMapEmpty && !isAddMapEmpty) {
246 pickOld = false;
247 // If both maps are empty, then partial update is complete.
248 } else {
249 break;
250 }
251
252 if (pickOld) {
253 numOldPrefixPicked++;
254
255 // If the number of picks from old map matches the removalIndex, then this prefix
256 // will be removed by not merging it to new map.
257 if (removalIndex < removalArray.Length() &&
258 numOldPrefixPicked == removalArray[removalIndex]) {
259 removalIndex++;
260 } else {
261 AppendPrefixToMap(aOutputMap, smallestOldPrefix);
262 UpdateChecksum(crypto, smallestOldPrefix);
263 }
264 smallestOldPrefix.SetLength(0);
265 } else {
266 AppendPrefixToMap(aOutputMap, smallestAddPrefix);
267 UpdateChecksum(crypto, smallestAddPrefix);
268
269 smallestAddPrefix.SetLength(0);
270 }
271 }
272
273 // We expect index will be greater to 0 because max number of runs will be
274 // the number of original prefix plus add prefix.
275 if (index <= 0) {
276 LOG(("There are still prefixes remaining after reaching maximum runs."));
277 Telemetry::Accumulate(Telemetry::URLCLASSIFIER_UPDATE_ERROR_TYPE,
278 INFINITE_LOOP);
279 return NS_ERROR_FAILURE;
280 }
281
282 if (removalIndex < removalArray.Length()) {
283 LOG(("There are still prefixes to remove after exhausting the old PrefixSet."));
284 Telemetry::Accumulate(Telemetry::URLCLASSIFIER_UPDATE_ERROR_TYPE,
285 WRONG_REMOVAL_INDICES);
286 return NS_ERROR_FAILURE;
287 }
288
289 nsAutoCString checksum;
290 crypto->Finish(false, checksum);
291 if (aTableUpdate->Checksum().IsEmpty()) {
292 LOG(("Update checksum missing."));
293 Telemetry::Accumulate(Telemetry::URLCLASSIFIER_UPDATE_ERROR_TYPE,
294 MISSING_CHECKSUM);
295
296 // Generate our own checksum to tableUpdate to ensure there is always
297 // checksum in .metadata
298 std::string stdChecksum(checksum.BeginReading(), checksum.Length());
299 aTableUpdate->NewChecksum(stdChecksum);
300
301 } else if (aTableUpdate->Checksum() != checksum){
302 LOG(("Checksum mismatch after applying partial update"));
303 Telemetry::Accumulate(Telemetry::URLCLASSIFIER_UPDATE_ERROR_TYPE,
304 CHECKSUM_MISMATCH);
305 return NS_ERROR_FAILURE;
306 }
307
308 return NS_OK;
309 }
310
311 nsresult
InitCrypto(nsCOMPtr<nsICryptoHash> & aCrypto)312 LookupCacheV4::InitCrypto(nsCOMPtr<nsICryptoHash>& aCrypto)
313 {
314 nsresult rv;
315 aCrypto = do_CreateInstance(NS_CRYPTO_HASH_CONTRACTID, &rv);
316 if (NS_WARN_IF(NS_FAILED(rv))) {
317 return rv;
318 }
319
320 rv = aCrypto->Init(nsICryptoHash::SHA256);
321 NS_WARNING_ASSERTION(NS_SUCCEEDED(rv), "InitCrypto failed");
322
323 return rv;
324 }
325
326 nsresult
VerifyChecksum(const nsACString & aChecksum)327 LookupCacheV4::VerifyChecksum(const nsACString& aChecksum)
328 {
329 nsCOMPtr<nsICryptoHash> crypto;
330 nsresult rv = InitCrypto(crypto);
331 if (NS_FAILED(rv)) {
332 return rv;
333 }
334
335 PrefixStringMap map;
336 mVLPrefixSet->GetPrefixes(map);
337
338 VLPrefixSet loadPSet(map);
339 uint32_t index = loadPSet.Count() + 1;
340 for(;index > 0; index--) {
341 nsDependentCSubstring prefix;
342 if (!loadPSet.GetSmallestPrefix(prefix)) {
343 break;
344 }
345 UpdateChecksum(crypto, prefix);
346 }
347
348 nsAutoCString checksum;
349 crypto->Finish(false, checksum);
350
351 if (checksum != aChecksum) {
352 LOG(("Checksum mismatch when loading prefixes from file."));
353 return NS_ERROR_FILE_CORRUPTED;
354 }
355
356 return NS_OK;
357 }
358
359 //////////////////////////////////////////////////////////////////////////
360 // A set of lightweight functions for reading/writing value from/to file.
361
362 namespace {
363
364 template<typename T>
365 struct ValueTraits
366 {
Lengthmozilla::safebrowsing::__anondf4af98d0111::ValueTraits367 static uint32_t Length(const T& aValue) { return sizeof(T); }
WritePtrmozilla::safebrowsing::__anondf4af98d0111::ValueTraits368 static char* WritePtr(T& aValue, uint32_t aLength) { return (char*)&aValue; }
ReadPtrmozilla::safebrowsing::__anondf4af98d0111::ValueTraits369 static const char* ReadPtr(const T& aValue) { return (char*)&aValue; }
IsFixedLengthmozilla::safebrowsing::__anondf4af98d0111::ValueTraits370 static bool IsFixedLength() { return true; }
371 };
372
373 template<>
374 struct ValueTraits<nsACString>
375 {
IsFixedLengthmozilla::safebrowsing::__anondf4af98d0111::ValueTraits376 static bool IsFixedLength() { return false; }
377
Lengthmozilla::safebrowsing::__anondf4af98d0111::ValueTraits378 static uint32_t Length(const nsACString& aValue)
379 {
380 return aValue.Length();
381 }
382
WritePtrmozilla::safebrowsing::__anondf4af98d0111::ValueTraits383 static char* WritePtr(nsACString& aValue, uint32_t aLength)
384 {
385 aValue.SetLength(aLength);
386 return aValue.BeginWriting();
387 }
388
ReadPtrmozilla::safebrowsing::__anondf4af98d0111::ValueTraits389 static const char* ReadPtr(const nsACString& aValue)
390 {
391 return aValue.BeginReading();
392 }
393 };
394
395 template<typename T> static nsresult
WriteValue(nsIOutputStream * aOutputStream,const T & aValue)396 WriteValue(nsIOutputStream *aOutputStream, const T& aValue)
397 {
398 uint32_t writeLength = ValueTraits<T>::Length(aValue);
399 if (!ValueTraits<T>::IsFixedLength()) {
400 // We need to write out the variable value length.
401 nsresult rv = WriteValue(aOutputStream, writeLength);
402 NS_ENSURE_SUCCESS(rv, rv);
403 }
404
405 // Write out the value.
406 auto valueReadPtr = ValueTraits<T>::ReadPtr(aValue);
407 uint32_t written;
408 nsresult rv = aOutputStream->Write(valueReadPtr, writeLength, &written);
409 if (NS_FAILED(rv) || written != writeLength) {
410 LOG(("Failed to write the value."));
411 return NS_FAILED(rv) ? rv : NS_ERROR_FAILURE;
412 }
413
414 return rv;
415 }
416
417 template<typename T> static nsresult
ReadValue(nsIInputStream * aInputStream,T & aValue)418 ReadValue(nsIInputStream* aInputStream, T& aValue)
419 {
420 nsresult rv;
421
422 uint32_t readLength;
423 if (ValueTraits<T>::IsFixedLength()) {
424 readLength = ValueTraits<T>::Length(aValue);
425 } else {
426 // Read the variable value length from file.
427 nsresult rv = ReadValue(aInputStream, readLength);
428 NS_ENSURE_SUCCESS(rv, rv);
429 }
430
431 // Read the value.
432 uint32_t read;
433 auto valueWritePtr = ValueTraits<T>::WritePtr(aValue, readLength);
434 rv = aInputStream->Read(valueWritePtr, readLength, &read);
435 if (NS_FAILED(rv) || read != readLength) {
436 LOG(("Failed to read the value."));
437 return NS_FAILED(rv) ? rv : NS_ERROR_FAILURE;
438 }
439
440 return rv;
441 }
442
443 } // end of unnamed namespace.
444 ////////////////////////////////////////////////////////////////////////
445
446 nsresult
WriteMetadata(TableUpdateV4 * aTableUpdate)447 LookupCacheV4::WriteMetadata(TableUpdateV4* aTableUpdate)
448 {
449 NS_ENSURE_ARG_POINTER(aTableUpdate);
450 if (nsUrlClassifierDBService::ShutdownHasStarted()) {
451 return NS_ERROR_ABORT;
452 }
453
454 nsCOMPtr<nsIFile> metaFile;
455 nsresult rv = mStoreDirectory->Clone(getter_AddRefs(metaFile));
456 NS_ENSURE_SUCCESS(rv, rv);
457
458 rv = metaFile->AppendNative(mTableName + METADATA_SUFFIX);
459 NS_ENSURE_SUCCESS(rv, rv);
460
461 nsCOMPtr<nsIOutputStream> outputStream;
462 rv = NS_NewLocalFileOutputStream(getter_AddRefs(outputStream), metaFile,
463 PR_WRONLY | PR_TRUNCATE | PR_CREATE_FILE);
464 if (!NS_SUCCEEDED(rv)) {
465 LOG(("Unable to create file to store metadata."));
466 return rv;
467 }
468
469 // Write the state.
470 rv = WriteValue(outputStream, aTableUpdate->ClientState());
471 if (NS_FAILED(rv)) {
472 LOG(("Failed to write the list state."));
473 return rv;
474 }
475
476 // Write the checksum.
477 rv = WriteValue(outputStream, aTableUpdate->Checksum());
478 if (NS_FAILED(rv)) {
479 LOG(("Failed to write the list checksum."));
480 return rv;
481 }
482
483 return rv;
484 }
485
486 nsresult
LoadMetadata(nsACString & aState,nsACString & aChecksum)487 LookupCacheV4::LoadMetadata(nsACString& aState, nsACString& aChecksum)
488 {
489 nsCOMPtr<nsIFile> metaFile;
490 nsresult rv = mStoreDirectory->Clone(getter_AddRefs(metaFile));
491 NS_ENSURE_SUCCESS(rv, rv);
492
493 rv = metaFile->AppendNative(mTableName + METADATA_SUFFIX);
494 NS_ENSURE_SUCCESS(rv, rv);
495
496 nsCOMPtr<nsIInputStream> localInFile;
497 rv = NS_NewLocalFileInputStream(getter_AddRefs(localInFile), metaFile,
498 PR_RDONLY | nsIFile::OS_READAHEAD);
499 if (NS_FAILED(rv)) {
500 LOG(("Unable to open metadata file."));
501 return rv;
502 }
503
504 // Read the list state.
505 rv = ReadValue(localInFile, aState);
506 if (NS_FAILED(rv)) {
507 LOG(("Failed to read state."));
508 return rv;
509 }
510
511 // Read the checksum.
512 rv = ReadValue(localInFile, aChecksum);
513 if (NS_FAILED(rv)) {
514 LOG(("Failed to read checksum."));
515 return rv;
516 }
517
518 return rv;
519 }
520
VLPrefixSet(const PrefixStringMap & aMap)521 VLPrefixSet::VLPrefixSet(const PrefixStringMap& aMap)
522 : mCount(0)
523 {
524 for (auto iter = aMap.ConstIter(); !iter.Done(); iter.Next()) {
525 uint32_t size = iter.Key();
526 mMap.Put(size, new PrefixString(*iter.Data(), size));
527 mCount += iter.Data()->Length() / size;
528 }
529 }
530
VLPrefixSet(const TableUpdateV4::PrefixStdStringMap & aMap)531 VLPrefixSet::VLPrefixSet(const TableUpdateV4::PrefixStdStringMap& aMap)
532 : mCount(0)
533 {
534 for (auto iter = aMap.ConstIter(); !iter.Done(); iter.Next()) {
535 uint32_t size = iter.Key();
536 mMap.Put(size, new PrefixString(iter.Data()->GetPrefixString(), size));
537 mCount += iter.Data()->GetPrefixString().Length() / size;
538 }
539 }
540
541 void
Merge(PrefixStringMap & aPrefixMap)542 VLPrefixSet::Merge(PrefixStringMap& aPrefixMap) {
543 for (auto iter = mMap.ConstIter(); !iter.Done(); iter.Next()) {
544 nsCString* prefixString = aPrefixMap.LookupOrAdd(iter.Key());
545 PrefixString* str = iter.Data();
546
547 if (str->get()) {
548 prefixString->Append(str->get(), str->remaining());
549 }
550 }
551 }
552
553 bool
GetSmallestPrefix(nsDependentCSubstring & aOutString)554 VLPrefixSet::GetSmallestPrefix(nsDependentCSubstring& aOutString) {
555 PrefixString* pick = nullptr;
556 for (auto iter = mMap.ConstIter(); !iter.Done(); iter.Next()) {
557 PrefixString* str = iter.Data();
558
559 if (!str->get()) {
560 continue;
561 }
562
563 if (aOutString.IsEmpty()) {
564 aOutString.Rebind(str->get(), iter.Key());
565 pick = str;
566 continue;
567 }
568
569 nsDependentCSubstring cur(str->get(), iter.Key());
570 if (cur < aOutString) {
571 aOutString.Rebind(str->get(), iter.Key());
572 pick = str;
573 }
574 }
575
576 if (pick) {
577 pick->next();
578 }
579
580 return pick != nullptr;
581 }
582
583 } // namespace safebrowsing
584 } // namespace mozilla
585