michael@0: /* This Source Code Form is subject to the terms of the Mozilla Public
michael@0:  * License, v. 2.0. If a copy of the MPL was not distributed with this
michael@0:  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
michael@0: 
michael@0: #ifndef HashStore_h__
michael@0: #define HashStore_h__
michael@0: 
michael@0: #include "Entries.h"
michael@0: #include "ChunkSet.h"
michael@0: 
michael@0: #include "nsString.h"
michael@0: #include "nsTArray.h"
michael@0: #include "nsIFile.h"
michael@0: #include "nsIFileStreams.h"
michael@0: #include "nsCOMPtr.h"
michael@0: 
michael@0: namespace mozilla {
michael@0: namespace safebrowsing {
michael@0: 
michael@0: // A table update is built from a single update chunk from the server. As the
michael@0: // protocol parser processes each chunk, it constructs a table update with the
michael@0: // new hashes.
michael@0: class TableUpdate {
michael@0: public:
michael@0:   TableUpdate(const nsACString& aTable)
michael@0:       : mTable(aTable), mLocalUpdate(false) {}
michael@0:   const nsCString& TableName() const { return mTable; }
michael@0: 
michael@0:   bool Empty() const {
michael@0:     return mAddChunks.Length() == 0 &&
michael@0:       mSubChunks.Length() == 0 &&
michael@0:       mAddExpirations.Length() == 0 &&
michael@0:       mSubExpirations.Length() == 0 &&
michael@0:       mAddPrefixes.Length() == 0 &&
michael@0:       mSubPrefixes.Length() == 0 &&
michael@0:       mAddCompletes.Length() == 0 &&
michael@0:       mSubCompletes.Length() == 0;
michael@0:   }
michael@0: 
michael@0:   // Throughout, uint32_t aChunk refers only to the chunk number. Chunk data is
michael@0:   // stored in the Prefix structures.
michael@0:   void NewAddChunk(uint32_t aChunk) { mAddChunks.Set(aChunk); }
michael@0:   void NewSubChunk(uint32_t aChunk) { mSubChunks.Set(aChunk); }
michael@0: 
michael@0:   void NewAddExpiration(uint32_t aChunk) { mAddExpirations.Set(aChunk); }
michael@0:   void NewSubExpiration(uint32_t aChunk) { mSubExpirations.Set(aChunk); }
michael@0: 
michael@0:   void NewAddPrefix(uint32_t aAddChunk, const Prefix& aPrefix);
michael@0:   void NewSubPrefix(uint32_t aAddChunk, const Prefix& aPrefix, uint32_t aSubChunk);
michael@0: 
michael@0:   void NewAddComplete(uint32_t aChunk, const Completion& aCompletion);
michael@0:   void NewSubComplete(uint32_t aAddChunk, const Completion& aCompletion,
michael@0:                       uint32_t aSubChunk);
michael@0:   void SetLocalUpdate(void) { mLocalUpdate = true; }
michael@0:   bool IsLocalUpdate(void) { return mLocalUpdate; }
michael@0: 
michael@0:   ChunkSet& AddChunks() { return mAddChunks; }
michael@0:   ChunkSet& SubChunks() { return mSubChunks; }
michael@0: 
michael@0:   // Expirations for chunks.
michael@0:   ChunkSet& AddExpirations() { return mAddExpirations; }
michael@0:   ChunkSet& SubExpirations() { return mSubExpirations; }
michael@0: 
michael@0:   // Hashes associated with this chunk.
michael@0:   AddPrefixArray& AddPrefixes() { return mAddPrefixes; }
michael@0:   SubPrefixArray& SubPrefixes() { return mSubPrefixes; }
michael@0:   AddCompleteArray& AddCompletes() { return mAddCompletes; }
michael@0:   SubCompleteArray& SubCompletes() { return mSubCompletes; }
michael@0: 
michael@0: private:
michael@0:   nsCString mTable;
michael@0:   // Update not from the remote server (no freshness)
michael@0:   bool mLocalUpdate;
michael@0: 
michael@0:   // The list of chunk numbers that we have for each of the type of chunks.
michael@0:   ChunkSet mAddChunks;
michael@0:   ChunkSet mSubChunks;
michael@0:   ChunkSet mAddExpirations;
michael@0:   ChunkSet mSubExpirations;
michael@0: 
michael@0:   // 4-byte sha256 prefixes.
michael@0:   AddPrefixArray mAddPrefixes;
michael@0:   SubPrefixArray mSubPrefixes;
michael@0: 
michael@0:   // 32-byte hashes.
michael@0:   AddCompleteArray mAddCompletes;
michael@0:   SubCompleteArray mSubCompletes;
michael@0: };
michael@0: 
michael@0: // There is one hash store per table.
michael@0: class HashStore {
michael@0: public:
michael@0:   HashStore(const nsACString& aTableName, nsIFile* aStoreFile);
michael@0:   ~HashStore();
michael@0: 
michael@0:   const nsCString& TableName() const { return mTableName; }
michael@0: 
michael@0:   nsresult Open();
michael@0:   // Add Prefixes are stored partly in the PrefixSet (contains the
michael@0:   // Prefix data organized for fast lookup/low RAM usage) and partly in the
michael@0:   // HashStore (Add Chunk numbers - only used for updates, slow retrieval).
michael@0:   // AugmentAdds function joins the separate datasets into one complete
michael@0:   // prefixes+chunknumbers dataset.
michael@0:   nsresult AugmentAdds(const nsTArray<uint32_t>& aPrefixes);
michael@0: 
michael@0:   ChunkSet& AddChunks() { return mAddChunks; }
michael@0:   ChunkSet& SubChunks() { return mSubChunks; }
michael@0:   AddPrefixArray& AddPrefixes() { return mAddPrefixes; }
michael@0:   AddCompleteArray& AddCompletes() { return mAddCompletes; }
michael@0:   SubPrefixArray& SubPrefixes() { return mSubPrefixes; }
michael@0:   SubCompleteArray& SubCompletes() { return mSubCompletes; }
michael@0: 
michael@0:   // =======
michael@0:   // Updates
michael@0:   // =======
michael@0:   // Begin the update process.  Reads the store into memory.
michael@0:   nsresult BeginUpdate();
michael@0: 
michael@0:   // Imports the data from a TableUpdate.
michael@0:   nsresult ApplyUpdate(TableUpdate &aUpdate);
michael@0: 
michael@0:   // Process expired chunks
michael@0:   nsresult Expire();
michael@0: 
michael@0:   // Rebuild the store, Incorporating all the applied updates.
michael@0:   nsresult Rebuild();
michael@0: 
michael@0:   // Write the current state of the store to disk.
michael@0:   // If you call between ApplyUpdate() and Rebuild(), you'll
michael@0:   // have a mess on your hands.
michael@0:   nsresult WriteFile();
michael@0: 
michael@0:   // Wipe out all Completes.
michael@0:   void ClearCompletes();
michael@0: 
michael@0: private:
michael@0:   nsresult Reset();
michael@0: 
michael@0:   nsresult ReadHeader();
michael@0:   nsresult SanityCheck();
michael@0:   nsresult CalculateChecksum(nsAutoCString& aChecksum, uint32_t aFileSize,
michael@0:                              bool aChecksumPresent);
michael@0:   nsresult CheckChecksum(nsIFile* aStoreFile, uint32_t aFileSize);
michael@0:   void UpdateHeader();
michael@0: 
michael@0:   nsresult ReadChunkNumbers();
michael@0:   nsresult ReadHashes();
michael@0: 
michael@0:   nsresult ReadAddPrefixes();
michael@0:   nsresult ReadSubPrefixes();
michael@0: 
michael@0:   nsresult WriteAddPrefixes(nsIOutputStream* aOut);
michael@0:   nsresult WriteSubPrefixes(nsIOutputStream* aOut);
michael@0: 
michael@0:   nsresult ProcessSubs();
michael@0: 
michael@0:  // This is used for checking that the database is correct and for figuring out
michael@0:  // the number of chunks, etc. to read from disk on restart.
michael@0:   struct Header {
michael@0:     uint32_t magic;
michael@0:     uint32_t version;
michael@0:     uint32_t numAddChunks;
michael@0:     uint32_t numSubChunks;
michael@0:     uint32_t numAddPrefixes;
michael@0:     uint32_t numSubPrefixes;
michael@0:     uint32_t numAddCompletes;
michael@0:     uint32_t numSubCompletes;
michael@0:   };
michael@0: 
michael@0:   Header mHeader;
michael@0: 
michael@0:   // The name of the table (must end in -shavar or -digest256, or evidently
michael@0:   // -simple for unittesting.
michael@0:   nsCString mTableName;
michael@0:   nsCOMPtr<nsIFile> mStoreDirectory;
michael@0: 
michael@0:   bool mInUpdate;
michael@0: 
michael@0:   nsCOMPtr<nsIInputStream> mInputStream;
michael@0: 
michael@0:   // Chunk numbers, stored as uint32_t arrays.
michael@0:   ChunkSet mAddChunks;
michael@0:   ChunkSet mSubChunks;
michael@0: 
michael@0:   ChunkSet mAddExpirations;
michael@0:   ChunkSet mSubExpirations;
michael@0: 
michael@0:   // Chunk data for shavar tables. See Entries.h for format.
michael@0:   AddPrefixArray mAddPrefixes;
michael@0:   SubPrefixArray mSubPrefixes;
michael@0: 
michael@0:   // See bug 806422 for background. We must be able to distinguish between
michael@0:   // updates from the completion server and updates from the regular server.
michael@0:   AddCompleteArray mAddCompletes;
michael@0:   SubCompleteArray mSubCompletes;
michael@0: };
michael@0: 
michael@0: }
michael@0: }
michael@0: #endif