1 : //* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
2 : // Originally based on Chrome sources:
3 : // Copyright (c) 2010 The Chromium Authors. All rights reserved.
4 : //
5 : // Redistribution and use in source and binary forms, with or without
6 : // modification, are permitted provided that the following conditions are
7 : // met:
8 : //
9 : // * Redistributions of source code must retain the above copyright
10 : // notice, this list of conditions and the following disclaimer.
11 : // * Redistributions in binary form must reproduce the above
12 : // copyright notice, this list of conditions and the following disclaimer
13 : // in the documentation and/or other materials provided with the
14 : // distribution.
15 : // * Neither the name of Google Inc. nor the names of its
16 : // contributors may be used to endorse or promote products derived from
17 : // this software without specific prior written permission.
18 : //
19 : // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
20 : // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
21 : // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
22 : // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
23 : // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
24 : // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
25 : // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
26 : // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
27 : // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
28 : // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
29 : // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
30 :
31 :
32 : #include "HashStore.h"
33 : #include "nsAutoPtr.h"
34 : #include "nsICryptoHash.h"
35 : #include "nsISeekableStream.h"
36 : #include "nsIStreamConverterService.h"
37 : #include "nsNetUtil.h"
38 : #include "nsCheckSummedOutputStream.h"
39 : #include "prlog.h"
40 : #include "zlib.h"
41 :
42 : // Main store for SafeBrowsing protocol data. We store
43 : // known add/sub chunks, prefixes and completions in memory
44 : // during an update, and serialize to disk.
45 : // We do not store the add prefixes, those are retrieved by
46 : // decompressing the PrefixSet cache whenever we need to apply
47 : // an update.
48 : //
49 : // byte slicing: Many of the 4-byte values stored here are strongly
50 : // correlated in the upper bytes, and uncorrelated in the lower
51 : // bytes. Because zlib/DEFLATE requires match lengths of at least
52 : // 3 to achieve good compression, and we don't get those if only
53 : // the upper 16-bits are correlated, it is worthwhile to slice 32-bit
54 : // values into 4 1-byte slices and compress the slices individually.
55 : // The slices corresponding to MSBs will compress very well, and the
56 : // slice corresponding to LSB almost nothing. Because of this, we
57 : // only apply DEFLATE to the 3 most significant bytes, and store the
58 : // LSB uncompressed.
59 : //
60 : // byte sliced (numValues) data format:
61 : // uint32 compressed-size
62 : // compressed-size bytes zlib DEFLATE data
63 : // 0...numValues byte MSB of 4-byte numValues data
64 : // uint32 compressed-size
65 : // compressed-size bytes zlib DEFLATE data
66 : // 0...numValues byte 2nd byte of 4-byte numValues data
67 : // uint32 compressed-size
68 : // compressed-size bytes zlib DEFLATE data
69 : // 0...numValues byte 3rd byte of 4-byte numValues data
70 : // 0...numValues byte LSB of 4-byte numValues data
71 : //
72 : // Store data format:
73 : // uint32 magic
74 : // uint32 version
75 : // uint32 numAddChunks
76 : // uint32 numSubChunks
77 : // uint32 numAddPrefixes
78 : // uint32 numSubPrefixes
79 : // uint32 numAddCompletes
80 : // uint32 numSubCompletes
81 : // 0...numAddChunks uint32 addChunk
82 : // 0...numSubChunks uint32 subChunk
83 : // byte sliced (numAddPrefixes) uint32 add chunk of AddPrefixes
84 : // byte sliced (numSubPrefixes) uint32 sub chunk of SubPrefixes
85 : // byte sliced (numSubPrefixes) uint32 add chunk of SubPrefixes
86 : // byte sliced (numSubPrefixes) uint32 SubPrefixes
87 : // 0...numAddCompletes 32-byte Completions
88 : // 0...numSubCompletes 32-byte Completions
89 : // 16-byte MD5 of all preceding data
90 :
91 : // NSPR_LOG_MODULES=UrlClassifierDbService:5
92 : extern PRLogModuleInfo *gUrlClassifierDbServiceLog;
93 : #if defined(PR_LOGGING)
94 : #define LOG(args) PR_LOG(gUrlClassifierDbServiceLog, PR_LOG_DEBUG, args)
95 : #define LOG_ENABLED() PR_LOG_TEST(gUrlClassifierDbServiceLog, 4)
96 : #else
97 : #define LOG(args)
98 : #define LOG_ENABLED() (PR_FALSE)
99 : #endif
100 :
101 : namespace mozilla {
102 : namespace safebrowsing {
103 :
104 : const uint32 STORE_MAGIC = 0x1231af3b;
105 : const uint32 CURRENT_VERSION = 2;
106 :
107 : void
108 45 : TableUpdate::NewAddPrefix(PRUint32 aAddChunk, const Prefix& aHash)
109 : {
110 45 : AddPrefix *add = mAddPrefixes.AppendElement();
111 45 : add->addChunk = aAddChunk;
112 45 : add->prefix = aHash;
113 45 : }
114 :
115 : void
116 2 : TableUpdate::NewSubPrefix(PRUint32 aAddChunk, const Prefix& aHash, PRUint32 aSubChunk)
117 : {
118 2 : SubPrefix *sub = mSubPrefixes.AppendElement();
119 2 : sub->addChunk = aAddChunk;
120 2 : sub->prefix = aHash;
121 2 : sub->subChunk = aSubChunk;
122 2 : }
123 :
124 : void
125 141 : TableUpdate::NewAddComplete(PRUint32 aAddChunk, const Completion& aHash)
126 : {
127 141 : AddComplete *add = mAddCompletes.AppendElement();
128 141 : add->addChunk = aAddChunk;
129 141 : add->hash.complete = aHash;
130 141 : }
131 :
132 : void
133 24 : TableUpdate::NewSubComplete(PRUint32 aAddChunk, const Completion& aHash, PRUint32 aSubChunk)
134 : {
135 24 : SubComplete *sub = mSubCompletes.AppendElement();
136 24 : sub->addChunk = aAddChunk;
137 24 : sub->hash.complete = aHash;
138 24 : sub->subChunk = aSubChunk;
139 24 : }
140 :
141 :
142 244 : HashStore::HashStore(const nsACString& aTableName, nsIFile* aStoreDir)
143 : : mTableName(aTableName)
144 : , mStoreDirectory(aStoreDir)
145 244 : , mInUpdate(false)
146 : {
147 244 : }
148 :
149 244 : HashStore::~HashStore()
150 : {
151 244 : }
152 :
153 : nsresult
154 0 : HashStore::Reset()
155 : {
156 0 : LOG(("HashStore resetting"));
157 :
158 0 : nsCOMPtr<nsIFile> storeFile;
159 0 : nsresult rv = mStoreDirectory->Clone(getter_AddRefs(storeFile));
160 0 : NS_ENSURE_SUCCESS(rv, rv);
161 :
162 0 : rv = storeFile->AppendNative(mTableName + NS_LITERAL_CSTRING(".sbstore"));
163 0 : NS_ENSURE_SUCCESS(rv, rv);
164 :
165 0 : rv = storeFile->Remove(PR_FALSE);
166 0 : NS_ENSURE_SUCCESS(rv, rv);
167 :
168 0 : Clear();
169 :
170 0 : return NS_OK;
171 : }
172 :
173 : nsresult
174 199 : HashStore::CheckChecksum(nsIFile* aStoreFile)
175 : {
176 : // Check for file corruption by
177 : // comparing the stored checksum to actual checksum of data
178 398 : nsCAutoString hash;
179 398 : nsCAutoString compareHash;
180 : char *data;
181 : PRUint32 read;
182 :
183 : PRInt64 fileSize;
184 199 : nsresult rv = aStoreFile->GetFileSize(&fileSize);
185 199 : NS_ENSURE_SUCCESS(rv, rv);
186 :
187 199 : if (fileSize < 0) {
188 0 : return NS_ERROR_FAILURE;
189 : }
190 :
191 199 : rv = CalculateChecksum(hash, true);
192 199 : NS_ENSURE_SUCCESS(rv, rv);
193 :
194 199 : compareHash.GetMutableData(&data, hash.Length());
195 :
196 398 : nsCOMPtr<nsISeekableStream> seekIn = do_QueryInterface(mInputStream);
197 199 : rv = seekIn->Seek(nsISeekableStream::NS_SEEK_SET, fileSize-hash.Length());
198 199 : NS_ENSURE_SUCCESS(rv, rv);
199 :
200 199 : rv = mInputStream->Read(data, hash.Length(), &read);
201 199 : NS_ENSURE_SUCCESS(rv, rv);
202 199 : NS_ASSERTION(read == hash.Length(), "Could not read hash bytes");
203 :
204 199 : if (!hash.Equals(compareHash)) {
205 0 : NS_WARNING("Safebrowing file failed checksum.");
206 0 : return NS_ERROR_FAILURE;
207 : }
208 :
209 199 : return NS_OK;
210 : }
211 :
212 : nsresult
213 244 : HashStore::Open()
214 : {
215 488 : nsCOMPtr<nsIFile> storeFile;
216 244 : nsresult rv = mStoreDirectory->Clone(getter_AddRefs(storeFile));
217 244 : NS_ENSURE_SUCCESS(rv, rv);
218 :
219 244 : rv = storeFile->AppendNative(mTableName + NS_LITERAL_CSTRING(".sbstore"));
220 244 : NS_ENSURE_SUCCESS(rv, rv);
221 :
222 488 : nsCOMPtr<nsIInputStream> origStream;
223 244 : rv = NS_NewLocalFileInputStream(getter_AddRefs(origStream), storeFile,
224 244 : PR_RDONLY);
225 :
226 244 : if (NS_FAILED(rv) && rv != NS_ERROR_FILE_NOT_FOUND) {
227 0 : Reset();
228 0 : return rv;
229 : }
230 :
231 244 : if (rv == NS_ERROR_FILE_NOT_FOUND) {
232 45 : Clear();
233 45 : UpdateHeader();
234 45 : return NS_OK;
235 : }
236 :
237 : PRInt64 fileSize;
238 199 : rv = storeFile->GetFileSize(&fileSize);
239 199 : NS_ENSURE_SUCCESS(rv, rv);
240 :
241 199 : rv = NS_NewBufferedInputStream(getter_AddRefs(mInputStream), origStream,
242 199 : fileSize);
243 199 : NS_ENSURE_SUCCESS(rv, rv);
244 :
245 199 : rv = CheckChecksum(storeFile);
246 199 : if (NS_FAILED(rv)) {
247 0 : Reset();
248 0 : return rv;
249 : }
250 :
251 199 : rv = ReadHeader();
252 199 : if (NS_FAILED(rv)) {
253 0 : Reset();
254 0 : return rv;
255 : }
256 :
257 199 : rv = SanityCheck(storeFile);
258 199 : if (NS_FAILED(rv)) {
259 0 : NS_WARNING("Safebrowsing file failed sanity check. probably out of date.");
260 0 : Reset();
261 0 : return rv;
262 : }
263 :
264 199 : rv = ReadChunkNumbers();
265 199 : if (NS_FAILED(rv)) {
266 0 : Reset();
267 0 : return rv;
268 : }
269 :
270 199 : return NS_OK;
271 : }
272 :
273 : void
274 232 : HashStore::Clear()
275 : {
276 232 : mAddChunks.Clear();
277 232 : mSubChunks.Clear();
278 232 : mAddExpirations.Clear();
279 232 : mSubExpirations.Clear();
280 232 : mAddPrefixes.Clear();
281 232 : mSubPrefixes.Clear();
282 232 : mAddCompletes.Clear();
283 232 : mSubCompletes.Clear();
284 232 : }
285 :
286 : nsresult
287 97 : HashStore::ReadEntireStore()
288 : {
289 97 : Clear();
290 :
291 97 : nsresult rv = ReadHeader();
292 97 : NS_ENSURE_SUCCESS(rv, rv);
293 :
294 97 : rv = ReadChunkNumbers();
295 97 : NS_ENSURE_SUCCESS(rv, rv);
296 :
297 97 : rv = ReadHashes();
298 97 : if (NS_FAILED(rv)) {
299 : // we are the only one reading this so it's up to us to detect corruption
300 0 : Reset();
301 : }
302 :
303 97 : return rv;
304 : }
305 :
306 : nsresult
307 296 : HashStore::ReadHeader()
308 : {
309 296 : if (!mInputStream) {
310 45 : Clear();
311 45 : UpdateHeader();
312 45 : return NS_OK;
313 : }
314 :
315 502 : nsCOMPtr<nsISeekableStream> seekable = do_QueryInterface(mInputStream);
316 251 : nsresult rv = seekable->Seek(nsISeekableStream::NS_SEEK_SET, 0);
317 251 : NS_ENSURE_SUCCESS(rv, rv);
318 :
319 251 : void *buffer = &mHeader;
320 : rv = NS_ReadInputStreamToBuffer(mInputStream,
321 : &buffer,
322 251 : sizeof(Header));
323 251 : NS_ENSURE_SUCCESS(rv, rv);
324 :
325 251 : return NS_OK;
326 : }
327 :
328 : nsresult
329 199 : HashStore::SanityCheck(nsIFile *storeFile)
330 : {
331 199 : if (mHeader.magic != STORE_MAGIC || mHeader.version != CURRENT_VERSION) {
332 0 : NS_WARNING("Unexpected header data in the store.");
333 0 : return NS_ERROR_FAILURE;
334 : }
335 :
336 199 : return NS_OK;
337 : }
338 :
339 : nsresult
340 199 : HashStore::CalculateChecksum(nsCAutoString& aChecksum, bool aChecksumPresent)
341 : {
342 199 : aChecksum.Truncate();
343 :
344 398 : nsCOMPtr<nsIFile> storeFile;
345 199 : nsresult rv = mStoreDirectory->Clone(getter_AddRefs(storeFile));
346 199 : NS_ENSURE_SUCCESS(rv, rv);
347 :
348 199 : rv = storeFile->AppendNative(mTableName + NS_LITERAL_CSTRING(".sbstore"));
349 199 : NS_ENSURE_SUCCESS(rv, rv);
350 :
351 398 : nsCOMPtr<nsIInputStream> hashStream;
352 :
353 199 : rv = NS_NewLocalFileInputStream(getter_AddRefs(hashStream), storeFile,
354 199 : PR_RDONLY);
355 :
356 199 : if (NS_FAILED(rv) && rv != NS_ERROR_FILE_NOT_FOUND) {
357 0 : Reset();
358 0 : return rv;
359 : }
360 :
361 : PRInt64 fileSize;
362 199 : rv = storeFile->GetFileSize(&fileSize);
363 199 : NS_ENSURE_SUCCESS(rv, rv);
364 :
365 199 : if (fileSize < 0) {
366 0 : return NS_ERROR_FAILURE;
367 : }
368 :
369 398 : nsCOMPtr<nsICryptoHash> hash = do_CreateInstance(NS_CRYPTO_HASH_CONTRACTID, &rv);
370 199 : NS_ENSURE_SUCCESS(rv, rv);
371 :
372 : // Size of MD5 hash in bytes
373 199 : const uint32 CHECKSUM_SIZE = 16;
374 :
375 199 : rv = hash->Init(nsICryptoHash::MD5);
376 199 : NS_ENSURE_SUCCESS(rv, rv);
377 :
378 199 : if (!aChecksumPresent) {
379 : // Hash entire file
380 0 : rv = hash->UpdateFromStream(hashStream, PR_UINT32_MAX);
381 : } else {
382 : // Hash everything but last checksum bytes
383 199 : rv = hash->UpdateFromStream(hashStream, fileSize-CHECKSUM_SIZE);
384 : }
385 199 : NS_ENSURE_SUCCESS(rv, rv);
386 :
387 199 : rv = hash->Finish(PR_FALSE, aChecksum);
388 199 : NS_ENSURE_SUCCESS(rv, rv);
389 :
390 199 : return NS_OK;
391 : }
392 :
393 : void
394 187 : HashStore::UpdateHeader()
395 : {
396 187 : mHeader.magic = STORE_MAGIC;
397 187 : mHeader.version = CURRENT_VERSION;
398 :
399 187 : mHeader.numAddChunks = mAddChunks.Length();
400 187 : mHeader.numSubChunks = mSubChunks.Length();
401 187 : mHeader.numAddPrefixes = mAddPrefixes.Length();
402 187 : mHeader.numSubPrefixes = mSubPrefixes.Length();
403 187 : mHeader.numAddCompletes = mAddCompletes.Length();
404 187 : mHeader.numSubCompletes = mSubCompletes.Length();
405 187 : }
406 :
407 : nsresult
408 296 : HashStore::ReadChunkNumbers()
409 : {
410 296 : if (!mInputStream) {
411 45 : LOG(("Clearing."));
412 45 : Clear();
413 45 : return NS_OK;
414 : }
415 :
416 502 : nsCOMPtr<nsISeekableStream> seekable = do_QueryInterface(mInputStream);
417 251 : nsresult rv = seekable->Seek(nsISeekableStream::NS_SEEK_SET,
418 251 : sizeof(Header));
419 :
420 251 : rv = mAddChunks.Read(mInputStream, mHeader.numAddChunks);
421 251 : NS_ENSURE_SUCCESS(rv, rv);
422 251 : NS_ASSERTION(mAddChunks.Length() == mHeader.numAddChunks, "Read the right amount of add chunks.");
423 :
424 251 : rv = mSubChunks.Read(mInputStream, mHeader.numSubChunks);
425 251 : NS_ENSURE_SUCCESS(rv, rv);
426 251 : NS_ASSERTION(mSubChunks.Length() == mHeader.numSubChunks, "Read the right amount of sub chunks.");
427 :
428 251 : return NS_OK;
429 : }
430 :
431 : nsresult
432 97 : HashStore::ReadHashes()
433 : {
434 97 : if (!mInputStream) {
435 45 : return NS_OK;
436 : }
437 :
438 104 : nsCOMPtr<nsISeekableStream> seekable = do_QueryInterface(mInputStream);
439 :
440 52 : uint32 offset = sizeof(Header);
441 52 : offset += (mHeader.numAddChunks + mHeader.numSubChunks) * sizeof(uint32);
442 52 : nsresult rv = seekable->Seek(nsISeekableStream::NS_SEEK_SET, offset);
443 :
444 52 : rv = ReadAddPrefixes();
445 52 : NS_ENSURE_SUCCESS(rv, rv);
446 :
447 52 : rv = ReadSubPrefixes();
448 52 : NS_ENSURE_SUCCESS(rv, rv);
449 :
450 52 : rv = ReadTArray(mInputStream, &mAddCompletes, mHeader.numAddCompletes);
451 52 : NS_ENSURE_SUCCESS(rv, rv);
452 :
453 52 : rv = ReadTArray(mInputStream, &mSubCompletes, mHeader.numSubCompletes);
454 52 : NS_ENSURE_SUCCESS(rv, rv);
455 :
456 52 : return NS_OK;
457 : }
458 :
459 : nsresult
460 97 : HashStore::BeginUpdate()
461 : {
462 97 : mInUpdate = true;
463 :
464 97 : nsresult rv = ReadEntireStore();
465 97 : NS_ENSURE_SUCCESS(rv, rv);
466 :
467 97 : return NS_OK;
468 : }
469 :
470 : template<class T>
471 : static nsresult
472 488 : Merge(ChunkSet* aStoreChunks,
473 : nsTArray<T>* aStorePrefixes,
474 : ChunkSet& aUpdateChunks,
475 : nsTArray<T>& aUpdatePrefixes)
476 : {
477 488 : EntrySort(aUpdatePrefixes);
478 :
479 488 : T* updateIter = aUpdatePrefixes.Elements();
480 488 : T* updateEnd = aUpdatePrefixes.Elements() + aUpdatePrefixes.Length();
481 :
482 488 : T* storeIter = aStorePrefixes->Elements();
483 488 : T* storeEnd = aStorePrefixes->Elements() + aStorePrefixes->Length();
484 :
485 : // use a separate array so we can keep the iterators valid
486 : // if the nsTArray grows
487 976 : nsTArray<T> adds;
488 :
489 695 : for (; updateIter != updateEnd; updateIter++) {
490 : // XXX: binary search for insertion point might be faster in common
491 : // case?
492 485 : while (storeIter < storeEnd && (storeIter->Compare(*updateIter) < 0)) {
493 : // skip forward to matching element (or not...)
494 71 : storeIter++;
495 : }
496 : // no match, add
497 207 : if (storeIter == storeEnd
498 : || storeIter->Compare(*updateIter) != 0) {
499 198 : if (!adds.AppendElement(*updateIter))
500 0 : return NS_ERROR_OUT_OF_MEMORY;
501 : }
502 : }
503 :
504 : // chunks can be empty, but we should still report we have them
505 : // to make the chunkranges continuous
506 488 : aStoreChunks->Merge(aUpdateChunks);
507 :
508 488 : aStorePrefixes->AppendElements(adds);
509 488 : EntrySort(*aStorePrefixes);
510 :
511 488 : return NS_OK;
512 : }
513 :
514 : nsresult
515 122 : HashStore::ApplyUpdate(TableUpdate &update)
516 : {
517 122 : nsresult rv = mAddExpirations.Merge(update.AddExpirations());
518 122 : NS_ENSURE_SUCCESS(rv, rv);
519 :
520 122 : rv = mSubExpirations.Merge(update.SubExpirations());
521 122 : NS_ENSURE_SUCCESS(rv, rv);
522 :
523 122 : rv = Expire();
524 122 : NS_ENSURE_SUCCESS(rv, rv);
525 :
526 : rv = Merge(&mAddChunks, &mAddPrefixes,
527 122 : update.AddChunks(), update.AddPrefixes());
528 122 : NS_ENSURE_SUCCESS(rv, rv);
529 :
530 : rv = Merge(&mAddChunks, &mAddCompletes,
531 122 : update.AddChunks(), update.AddCompletes());
532 122 : NS_ENSURE_SUCCESS(rv, rv);
533 :
534 : rv = Merge(&mSubChunks, &mSubPrefixes,
535 122 : update.SubChunks(), update.SubPrefixes());
536 122 : NS_ENSURE_SUCCESS(rv, rv);
537 :
538 : rv = Merge(&mSubChunks, &mSubCompletes,
539 122 : update.SubChunks(), update.SubCompletes());
540 122 : NS_ENSURE_SUCCESS(rv, rv);
541 :
542 122 : return NS_OK;
543 : }
544 :
545 : nsresult
546 97 : HashStore::Rebuild()
547 : {
548 97 : NS_ASSERTION(mInUpdate, "Must be in update to rebuild.");
549 :
550 97 : nsresult rv = ProcessSubs();
551 97 : NS_ENSURE_SUCCESS(rv, rv);
552 :
553 97 : UpdateHeader();
554 :
555 97 : return NS_OK;
556 : }
557 :
558 : template<class T>
559 : static void
560 682 : ExpireEntries(nsTArray<T>* aEntries, ChunkSet& aExpirations)
561 : {
562 682 : T* addIter = aEntries->Elements();
563 682 : T* end = aEntries->Elements() + aEntries->Length();
564 :
565 1112 : for (T *iter = addIter; iter != end; iter++) {
566 430 : if (!aExpirations.Has(iter->Chunk())) {
567 406 : *addIter = *iter;
568 406 : addIter++;
569 : }
570 : }
571 :
572 682 : aEntries->SetLength(addIter - aEntries->Elements());
573 682 : }
574 :
575 : nsresult
576 122 : HashStore::Expire()
577 : {
578 122 : ExpireEntries(&mAddPrefixes, mAddExpirations);
579 122 : ExpireEntries(&mAddCompletes, mAddExpirations);
580 122 : ExpireEntries(&mSubPrefixes, mSubExpirations);
581 122 : ExpireEntries(&mSubCompletes, mSubExpirations);
582 :
583 122 : mAddChunks.Remove(mAddExpirations);
584 122 : mSubChunks.Remove(mSubExpirations);
585 :
586 122 : mAddExpirations.Clear();
587 122 : mSubExpirations.Clear();
588 :
589 122 : return NS_OK;
590 : }
591 :
592 : template<class T>
593 1164 : nsresult DeflateWriteTArray(nsIOutputStream* aStream, nsTArray<T>& aIn)
594 : {
595 1164 : uLongf insize = aIn.Length() * sizeof(T);
596 1164 : uLongf outsize = compressBound(insize);
597 2328 : nsTArray<char> outBuff;
598 1164 : outBuff.SetLength(outsize);
599 :
600 : int zerr = compress(reinterpret_cast<Bytef*>(outBuff.Elements()),
601 : &outsize,
602 : reinterpret_cast<const Bytef*>(aIn.Elements()),
603 1164 : insize);
604 1164 : if (zerr != Z_OK) {
605 0 : return NS_ERROR_FAILURE;
606 : }
607 1164 : LOG(("DeflateWriteTArray: %d in %d out", insize, outsize));
608 :
609 1164 : outBuff.TruncateLength(outsize);
610 :
611 : // Length of compressed data stream
612 1164 : PRUint32 dataLen = outBuff.Length();
613 : PRUint32 written;
614 1164 : nsresult rv = aStream->Write(reinterpret_cast<char*>(&dataLen), sizeof(dataLen), &written);
615 1164 : NS_ENSURE_SUCCESS(rv, rv);
616 :
617 1164 : NS_ASSERTION(written == sizeof(dataLen), "Error writing deflate length");
618 :
619 : // Store to stream
620 1164 : rv = WriteTArray(aStream, outBuff);
621 1164 : NS_ENSURE_SUCCESS(rv, rv);
622 :
623 1164 : return NS_OK;
624 : }
625 :
626 : template<class T>
627 624 : nsresult InflateReadTArray(nsIInputStream* aStream, nsTArray<T>* aOut,
628 : PRUint32 aExpectedSize)
629 : {
630 :
631 : PRUint32 inLen;
632 : PRUint32 read;
633 624 : nsresult rv = aStream->Read(reinterpret_cast<char*>(&inLen), sizeof(inLen), &read);
634 624 : NS_ENSURE_SUCCESS(rv, rv);
635 :
636 624 : NS_ASSERTION(read == sizeof(inLen), "Error reading inflate length");
637 :
638 1248 : nsTArray<char> inBuff;
639 624 : inBuff.SetLength(inLen);
640 :
641 624 : rv = ReadTArray(aStream, &inBuff, inLen);
642 624 : NS_ENSURE_SUCCESS(rv, rv);
643 :
644 624 : uLongf insize = inLen;
645 624 : uLongf outsize = aExpectedSize * sizeof(T);
646 624 : aOut->SetLength(aExpectedSize);
647 :
648 : int zerr = uncompress(reinterpret_cast<Bytef*>(aOut->Elements()),
649 : &outsize,
650 : reinterpret_cast<const Bytef*>(inBuff.Elements()),
651 624 : insize);
652 624 : if (zerr != Z_OK) {
653 0 : return NS_ERROR_FAILURE;
654 : }
655 624 : LOG(("InflateReadTArray: %d in %d out", insize, outsize));
656 :
657 624 : NS_ASSERTION(outsize == aExpectedSize * sizeof(T), "Decompression size mismatch");
658 :
659 624 : return NS_OK;
660 : }
661 :
662 : static nsresult
663 388 : ByteSliceWrite(nsIOutputStream* aOut, nsTArray<PRUint32>& aData)
664 : {
665 776 : nsTArray<PRUint8> slice1;
666 776 : nsTArray<PRUint8> slice2;
667 776 : nsTArray<PRUint8> slice3;
668 776 : nsTArray<PRUint8> slice4;
669 388 : PRUint32 count = aData.Length();
670 :
671 388 : slice1.SetCapacity(count);
672 388 : slice2.SetCapacity(count);
673 388 : slice3.SetCapacity(count);
674 388 : slice4.SetCapacity(count);
675 :
676 495 : for (PRUint32 i = 0; i < count; i++) {
677 107 : slice1.AppendElement( aData[i] >> 24);
678 107 : slice2.AppendElement((aData[i] >> 16) & 0xFF);
679 107 : slice3.AppendElement((aData[i] >> 8) & 0xFF);
680 107 : slice4.AppendElement( aData[i] & 0xFF);
681 : }
682 :
683 388 : nsresult rv = DeflateWriteTArray(aOut, slice1);
684 388 : NS_ENSURE_SUCCESS(rv, rv);
685 388 : rv = DeflateWriteTArray(aOut, slice2);
686 388 : NS_ENSURE_SUCCESS(rv, rv);
687 388 : rv = DeflateWriteTArray(aOut, slice3);
688 388 : NS_ENSURE_SUCCESS(rv, rv);
689 : // The LSB slice is generally uncompressible, don't bother
690 : // compressing it.
691 388 : rv = WriteTArray(aOut, slice4);
692 388 : NS_ENSURE_SUCCESS(rv, rv);
693 :
694 388 : return NS_OK;
695 : }
696 :
697 : static nsresult
698 208 : ByteSliceRead(nsIInputStream* aInStream, nsTArray<PRUint32>* aData, PRUint32 count)
699 : {
700 416 : nsTArray<PRUint8> slice1;
701 416 : nsTArray<PRUint8> slice2;
702 416 : nsTArray<PRUint8> slice3;
703 416 : nsTArray<PRUint8> slice4;
704 :
705 208 : nsresult rv = InflateReadTArray(aInStream, &slice1, count);
706 208 : NS_ENSURE_SUCCESS(rv, rv);
707 :
708 208 : rv = InflateReadTArray(aInStream, &slice2, count);
709 208 : NS_ENSURE_SUCCESS(rv, rv);
710 :
711 208 : rv = InflateReadTArray(aInStream, &slice3, count);
712 208 : NS_ENSURE_SUCCESS(rv, rv);
713 :
714 208 : rv = ReadTArray(aInStream, &slice4, count);
715 208 : NS_ENSURE_SUCCESS(rv, rv);
716 :
717 208 : aData->SetCapacity(count);
718 :
719 272 : for (uint32 i = 0; i < count; i++) {
720 128 : aData->AppendElement((slice1[i] << 24) | (slice2[i] << 16)
721 128 : | (slice3[i] << 8) | (slice4[i]));
722 : }
723 :
724 208 : return NS_OK;
725 : }
726 :
727 : nsresult
728 52 : HashStore::ReadAddPrefixes()
729 : {
730 104 : nsTArray<PRUint32> chunks;
731 52 : PRUint32 count = mHeader.numAddPrefixes;
732 :
733 52 : nsresult rv = ByteSliceRead(mInputStream, &chunks, count);
734 52 : NS_ENSURE_SUCCESS(rv, rv);
735 :
736 52 : mAddPrefixes.SetCapacity(count);
737 116 : for (PRUint32 i = 0; i < count; i++) {
738 64 : AddPrefix *add = mAddPrefixes.AppendElement();
739 64 : add->prefix.FromUint32(0);
740 64 : add->addChunk = chunks[i];
741 : }
742 :
743 52 : return NS_OK;
744 : }
745 :
746 : nsresult
747 52 : HashStore::ReadSubPrefixes()
748 : {
749 104 : nsTArray<PRUint32> addchunks;
750 104 : nsTArray<PRUint32> subchunks;
751 104 : nsTArray<PRUint32> prefixes;
752 52 : PRUint32 count = mHeader.numSubPrefixes;
753 :
754 52 : nsresult rv = ByteSliceRead(mInputStream, &addchunks, count);
755 52 : NS_ENSURE_SUCCESS(rv, rv);
756 :
757 52 : rv = ByteSliceRead(mInputStream, &subchunks, count);
758 52 : NS_ENSURE_SUCCESS(rv, rv);
759 :
760 52 : rv = ByteSliceRead(mInputStream, &prefixes, count);
761 52 : NS_ENSURE_SUCCESS(rv, rv);
762 :
763 52 : mSubPrefixes.SetCapacity(count);
764 52 : for (uint32 i = 0; i < count; i++) {
765 0 : SubPrefix *sub = mSubPrefixes.AppendElement();
766 0 : sub->addChunk = addchunks[i];
767 0 : sub->prefix.FromUint32(prefixes[i]);
768 0 : sub->subChunk = subchunks[i];
769 : }
770 :
771 52 : return NS_OK;
772 : }
773 :
774 : // Split up PrefixArray back into the constituents
775 : nsresult
776 97 : HashStore::WriteAddPrefixes(nsIOutputStream* aOut)
777 : {
778 194 : nsTArray<PRUint32> chunks;
779 97 : PRUint32 count = mAddPrefixes.Length();
780 97 : chunks.SetCapacity(count);
781 :
782 204 : for (uint32 i = 0; i < count; i++) {
783 107 : chunks.AppendElement(mAddPrefixes[i].Chunk());
784 : }
785 :
786 97 : nsresult rv = ByteSliceWrite(aOut, chunks);
787 97 : NS_ENSURE_SUCCESS(rv, rv);
788 :
789 97 : return NS_OK;
790 : }
791 :
792 : nsresult
793 97 : HashStore::WriteSubPrefixes(nsIOutputStream* aOut)
794 : {
795 194 : nsTArray<PRUint32> addchunks;
796 194 : nsTArray<PRUint32> subchunks;
797 194 : nsTArray<PRUint32> prefixes;
798 97 : PRUint32 count = mSubPrefixes.Length();
799 97 : addchunks.SetCapacity(count);
800 97 : subchunks.SetCapacity(count);
801 97 : prefixes.SetCapacity(count);
802 :
803 97 : for (uint32 i = 0; i < count; i++) {
804 0 : addchunks.AppendElement(mSubPrefixes[i].AddChunk());
805 0 : prefixes.AppendElement(mSubPrefixes[i].PrefixHash().ToUint32());
806 0 : subchunks.AppendElement(mSubPrefixes[i].Chunk());
807 : }
808 :
809 97 : nsresult rv = ByteSliceWrite(aOut, addchunks);
810 97 : NS_ENSURE_SUCCESS(rv, rv);
811 :
812 97 : rv = ByteSliceWrite(aOut, subchunks);
813 97 : NS_ENSURE_SUCCESS(rv, rv);
814 :
815 97 : rv = ByteSliceWrite(aOut, prefixes);
816 97 : NS_ENSURE_SUCCESS(rv, rv);
817 :
818 97 : return NS_OK;
819 : }
820 :
821 : nsresult
822 97 : HashStore::WriteFile()
823 : {
824 194 : nsCOMPtr<nsIFile> storeFile;
825 97 : nsresult rv = mStoreDirectory->Clone(getter_AddRefs(storeFile));
826 97 : NS_ENSURE_SUCCESS(rv, rv);
827 97 : rv = storeFile->AppendNative(mTableName + NS_LITERAL_CSTRING(".sbstore"));
828 97 : NS_ENSURE_SUCCESS(rv, rv);
829 :
830 : // Need to close the inputstream here *before* rewriting its file.
831 : // Windows will fail with an access violation if we don't.
832 97 : if (mInputStream) {
833 52 : rv = mInputStream->Close();
834 52 : NS_ENSURE_SUCCESS(rv, rv);
835 : }
836 :
837 194 : nsCOMPtr<nsIOutputStream> out;
838 97 : rv = NS_NewCheckSummedOutputStream(getter_AddRefs(out), storeFile,
839 97 : PR_WRONLY | PR_TRUNCATE | PR_CREATE_FILE);
840 97 : NS_ENSURE_SUCCESS(rv, rv);
841 :
842 : PRUint32 written;
843 97 : rv = out->Write(reinterpret_cast<char*>(&mHeader), sizeof(mHeader), &written);
844 97 : NS_ENSURE_SUCCESS(rv, rv);
845 :
846 : // Write chunk numbers...
847 97 : rv = mAddChunks.Write(out);
848 97 : NS_ENSURE_SUCCESS(rv, rv);
849 :
850 97 : rv = mSubChunks.Write(out);
851 97 : NS_ENSURE_SUCCESS(rv, rv);
852 :
853 : // Write hashes..
854 97 : rv = WriteAddPrefixes(out);
855 97 : NS_ENSURE_SUCCESS(rv, rv);
856 :
857 97 : rv = WriteSubPrefixes(out);
858 97 : NS_ENSURE_SUCCESS(rv, rv);
859 :
860 97 : rv = WriteTArray(out, mAddCompletes);
861 97 : NS_ENSURE_SUCCESS(rv, rv);
862 :
863 97 : rv = WriteTArray(out, mSubCompletes);
864 97 : NS_ENSURE_SUCCESS(rv, rv);
865 :
866 194 : nsCOMPtr<nsISafeOutputStream> safeOut = do_QueryInterface(out, &rv);
867 97 : NS_ENSURE_SUCCESS(rv, rv);
868 :
869 97 : rv = safeOut->Finish();
870 97 : NS_ENSURE_SUCCESS(rv, rv);
871 :
872 : PRInt64 fileSize;
873 97 : rv = storeFile->GetFileSize(&fileSize);
874 97 : NS_ENSURE_SUCCESS(rv, rv);
875 :
876 : // Reopen the file now that we've rewritten it.
877 194 : nsCOMPtr<nsIInputStream> origStream;
878 97 : rv = NS_NewLocalFileInputStream(getter_AddRefs(origStream), storeFile,
879 97 : PR_RDONLY);
880 97 : NS_ENSURE_SUCCESS(rv, rv);
881 :
882 97 : rv = NS_NewBufferedInputStream(getter_AddRefs(mInputStream), origStream,
883 97 : fileSize);
884 97 : NS_ENSURE_SUCCESS(rv, rv);
885 :
886 97 : return NS_OK;
887 : }
888 :
889 : nsresult
890 97 : HashStore::FinishUpdate()
891 : {
892 : // Drop add/sub data, it's only used during updates.
893 97 : mAddPrefixes.Clear();
894 97 : mSubPrefixes.Clear();
895 97 : mAddCompletes.Clear();
896 97 : mSubCompletes.Clear();
897 :
898 97 : return NS_OK;
899 : }
900 :
901 : template <class T>
902 : static void
903 582 : Erase(nsTArray<T>* array, T* iterStart, T* iterEnd)
904 : {
905 582 : uint32 start = iterStart - array->Elements();
906 582 : uint32 count = iterEnd - iterStart;
907 :
908 582 : if (count > 0) {
909 27 : array->RemoveElementsAt(start, count);
910 : }
911 582 : }
912 :
913 : // Find items matching between |subs| and |adds|, and remove them,
914 : // recording the item from |adds| in |adds_removed|. To minimize
915 : // copies, the inputs are processing in parallel, so |subs| and |adds|
916 : // should be compatibly ordered (either by SBAddPrefixLess or
917 : // SBAddPrefixHashLess).
918 : //
919 : // |predAS| provides add < sub, |predSA| provides sub < add, for the
920 : // tightest compare appropriate (see calls in SBProcessSubs).
921 : template<class TSub, class TAdd>
922 : static void
923 194 : KnockoutSubs(nsTArray<TSub>* aSubs, nsTArray<TAdd>* aAdds)
924 : {
925 : // Keep a pair of output iterators for writing kept items. Due to
926 : // deletions, these may lag the main iterators. Using erase() on
927 : // individual items would result in O(N^2) copies. Using a list
928 : // would work around that, at double or triple the memory cost.
929 194 : TAdd* addOut = aAdds->Elements();
930 194 : TAdd* addIter = aAdds->Elements();
931 :
932 194 : TSub* subOut = aSubs->Elements();
933 194 : TSub* subIter = aSubs->Elements();
934 :
935 194 : TAdd* addEnd = addIter + aAdds->Length();
936 194 : TSub* subEnd = subIter + aSubs->Length();
937 :
938 425 : while (addIter != addEnd && subIter != subEnd) {
939 : // additer compare, so it compares on add chunk
940 37 : int32 cmp = addIter->Compare(*subIter);
941 37 : if (cmp > 0) {
942 : // If |*sub_iter| < |*add_iter|, retain the sub.
943 9 : *subOut = *subIter;
944 9 : ++subOut;
945 9 : ++subIter;
946 28 : } else if (cmp < 0) {
947 : // If |*add_iter| < |*sub_iter|, retain the add.
948 13 : *addOut = *addIter;
949 13 : ++addOut;
950 13 : ++addIter;
951 : } else {
952 : // Drop equal items
953 15 : ++addIter;
954 15 : ++subIter;
955 : }
956 : }
957 :
958 194 : Erase(aAdds, addOut, addIter);
959 194 : Erase(aSubs, subOut, subIter);
960 194 : }
961 :
962 : // Remove items in |removes| from |fullHashes|. |fullHashes| and
963 : // |removes| should be ordered by SBAddPrefix component.
964 : template <class T>
965 : static void
966 194 : RemoveMatchingPrefixes(const SubPrefixArray& aSubs, nsTArray<T>* aFullHashes)
967 : {
968 : // Where to store kept items.
969 194 : T* out = aFullHashes->Elements();
970 194 : T* hashIter = out;
971 194 : T* hashEnd = aFullHashes->Elements() + aFullHashes->Length();
972 :
973 194 : SubPrefix const * removeIter = aSubs.Elements();
974 194 : SubPrefix const * removeEnd = aSubs.Elements() + aSubs.Length();
975 :
976 389 : while (hashIter != hashEnd && removeIter != removeEnd) {
977 1 : int32 cmp = removeIter->CompareAlt(*hashIter);
978 1 : if (cmp > 0) {
979 : // Keep items less than |*removeIter|.
980 0 : *out = *hashIter;
981 0 : ++out;
982 0 : ++hashIter;
983 1 : } else if (cmp < 0) {
984 : // No hit for |*removeIter|, bump it forward.
985 0 : ++removeIter;
986 : } else {
987 : // Drop equal items, there may be multiple hits.
988 1 : do {
989 1 : ++hashIter;
990 : } while (hashIter != hashEnd &&
991 : !(removeIter->CompareAlt(*hashIter) < 0));
992 1 : ++removeIter;
993 : }
994 : }
995 194 : Erase(aFullHashes, out, hashIter);
996 194 : }
997 :
998 : nsresult
999 97 : HashStore::ProcessSubs()
1000 : {
1001 97 : EntrySort(mAddPrefixes);
1002 97 : EntrySort(mSubPrefixes);
1003 97 : EntrySort(mAddCompletes);
1004 97 : EntrySort(mSubCompletes);
1005 :
1006 97 : KnockoutSubs(&mSubPrefixes, &mAddPrefixes);
1007 :
1008 97 : RemoveMatchingPrefixes(mSubPrefixes, &mAddCompletes);
1009 97 : RemoveMatchingPrefixes(mSubPrefixes, &mSubCompletes);
1010 :
1011 97 : KnockoutSubs(&mSubCompletes, &mAddCompletes);
1012 :
1013 : // Clean up temporary subs used for knocking out completes
1014 194 : ChunkSet dummyChunks;
1015 97 : dummyChunks.Set(0);
1016 97 : ExpireEntries(&mSubPrefixes, dummyChunks);
1017 97 : ExpireEntries(&mSubCompletes, dummyChunks);
1018 97 : mSubChunks.Remove(dummyChunks);
1019 :
1020 97 : return NS_OK;
1021 : }
1022 :
1023 : nsresult
1024 97 : HashStore::AugmentAdds(const nsTArray<PRUint32>& aPrefixes)
1025 : {
1026 97 : uint32 cnt = aPrefixes.Length();
1027 97 : if (cnt != mAddPrefixes.Length()) {
1028 0 : LOG(("Amount of prefixes in cache not consistent with store (%d vs %d)",
1029 : aPrefixes.Length(), mAddPrefixes.Length()));
1030 0 : return NS_ERROR_FAILURE;
1031 : }
1032 161 : for (uint32 i = 0; i < cnt; i++) {
1033 64 : mAddPrefixes[i].prefix.FromUint32(aPrefixes[i]);
1034 : }
1035 97 : return NS_OK;
1036 : }
1037 :
1038 : }
1039 : }
|