diff options
Diffstat (limited to '')
| -rw-r--r-- | toolkit/components/url-classifier/HashStore.h | 321 |
1 files changed, 321 insertions, 0 deletions
diff --git a/toolkit/components/url-classifier/HashStore.h b/toolkit/components/url-classifier/HashStore.h new file mode 100644 index 0000000000..ccb665efa7 --- /dev/null +++ b/toolkit/components/url-classifier/HashStore.h @@ -0,0 +1,321 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ + +#ifndef HashStore_h__ +#define HashStore_h__ + +#include "Entries.h" +#include "ChunkSet.h" + +#include "nsString.h" +#include "nsTArray.h" +#include "nsIFile.h" +#include "nsISupports.h" +#include "nsCOMPtr.h" +#include <string> + +namespace mozilla { +namespace safebrowsing { + +// The abstract class of TableUpdateV2 and TableUpdateV4. This +// is convenient for passing the TableUpdate* around associated +// with v2 and v4 instance. +class TableUpdate { + public: + TableUpdate(const nsACString& aTable) : mTable(aTable) {} + + NS_INLINE_DECL_THREADSAFE_REFCOUNTING(TableUpdate); + + // To be overriden. + virtual bool Empty() const = 0; + + // Common interfaces. + const nsCString& TableName() const { return mTable; } + + template <typename T> + static T* Cast(TableUpdate* aThat) { + return (T::TAG == aThat->Tag() ? reinterpret_cast<T*>(aThat) : nullptr); + } + template <typename T> + static const T* Cast(const TableUpdate* aThat) { + return (T::TAG == aThat->Tag() ? reinterpret_cast<const T*>(aThat) + : nullptr); + } + + protected: + virtual ~TableUpdate() = default; + + private: + virtual int Tag() const = 0; + + const nsCString mTable; +}; + +typedef nsTArray<RefPtr<TableUpdate>> TableUpdateArray; +typedef nsTArray<RefPtr<const TableUpdate>> ConstTableUpdateArray; + +// A table update is built from a single update chunk from the server. As the +// protocol parser processes each chunk, it constructs a table update with the +// new hashes. +class TableUpdateV2 : public TableUpdate { + public: + explicit TableUpdateV2(const nsACString& aTable) : TableUpdate(aTable) {} + + bool Empty() const override { + return mAddChunks.Length() == 0 && mSubChunks.Length() == 0 && + mAddExpirations.Length() == 0 && mSubExpirations.Length() == 0 && + mAddPrefixes.Length() == 0 && mSubPrefixes.Length() == 0 && + mAddCompletes.Length() == 0 && mSubCompletes.Length() == 0 && + mMissPrefixes.Length() == 0; + } + + // Throughout, uint32_t aChunk refers only to the chunk number. Chunk data is + // stored in the Prefix structures. + [[nodiscard]] nsresult NewAddChunk(uint32_t aChunk) { + return mAddChunks.Set(aChunk); + }; + [[nodiscard]] nsresult NewSubChunk(uint32_t aChunk) { + return mSubChunks.Set(aChunk); + }; + [[nodiscard]] nsresult NewAddExpiration(uint32_t aChunk) { + return mAddExpirations.Set(aChunk); + }; + [[nodiscard]] nsresult NewSubExpiration(uint32_t aChunk) { + return mSubExpirations.Set(aChunk); + }; + [[nodiscard]] nsresult NewAddPrefix(uint32_t aAddChunk, + const Prefix& aPrefix); + [[nodiscard]] nsresult NewSubPrefix(uint32_t aAddChunk, const Prefix& aPrefix, + uint32_t aSubChunk); + [[nodiscard]] nsresult NewAddComplete(uint32_t aChunk, + const Completion& aCompletion); + [[nodiscard]] nsresult NewSubComplete(uint32_t aAddChunk, + const Completion& aCompletion, + uint32_t aSubChunk); + [[nodiscard]] nsresult NewMissPrefix(const Prefix& aPrefix); + + const ChunkSet& AddChunks() const { return mAddChunks; } + const ChunkSet& SubChunks() const { return mSubChunks; } + + // Expirations for chunks. + const ChunkSet& AddExpirations() const { return mAddExpirations; } + const ChunkSet& SubExpirations() const { return mSubExpirations; } + + // Hashes associated with this chunk. + AddPrefixArray& AddPrefixes() { return mAddPrefixes; } + SubPrefixArray& SubPrefixes() { return mSubPrefixes; } + const AddCompleteArray& AddCompletes() const { return mAddCompletes; } + AddCompleteArray& AddCompletes() { return mAddCompletes; } + SubCompleteArray& SubCompletes() { return mSubCompletes; } + + // Entries that cannot be completed. + const MissPrefixArray& MissPrefixes() const { return mMissPrefixes; } + + // For downcasting. + static const int TAG = 2; + + private: + // The list of chunk numbers that we have for each of the type of chunks. + ChunkSet mAddChunks; + ChunkSet mSubChunks; + ChunkSet mAddExpirations; + ChunkSet mSubExpirations; + + // 4-byte sha256 prefixes. + AddPrefixArray mAddPrefixes; + SubPrefixArray mSubPrefixes; + + // This is only used by gethash so don't add this to Header. + MissPrefixArray mMissPrefixes; + + // 32-byte hashes. + AddCompleteArray mAddCompletes; + SubCompleteArray mSubCompletes; + + virtual int Tag() const override { return TAG; } +}; + +// Structure for DBService/HashStore/Classifiers to update. +// It would contain the prefixes (both fixed and variable length) +// for addition and indices to removal. See Bug 1283009. +class TableUpdateV4 : public TableUpdate { + public: + typedef nsTArray<int32_t> RemovalIndiceArray; + + public: + explicit TableUpdateV4(const nsACString& aTable) + : TableUpdate(aTable), mFullUpdate(false) {} + + bool Empty() const override { + return mPrefixesMap.IsEmpty() && mRemovalIndiceArray.IsEmpty() && + mFullHashResponseMap.IsEmpty(); + } + + bool IsFullUpdate() const { return mFullUpdate; } + const PrefixStringMap& Prefixes() const { return mPrefixesMap; } + const RemovalIndiceArray& RemovalIndices() const { + return mRemovalIndiceArray; + } + const nsACString& ClientState() const { return mClientState; } + const nsACString& SHA256() const { return mSHA256; } + const FullHashResponseMap& FullHashResponse() const { + return mFullHashResponseMap; + } + + // For downcasting. + static const int TAG = 4; + + void SetFullUpdate(bool aIsFullUpdate) { mFullUpdate = aIsFullUpdate; } + void NewPrefixes(int32_t aSize, const nsACString& aPrefixes); + void SetNewClientState(const nsACString& aState) { mClientState = aState; } + void SetSHA256(const std::string& aSHA256); + + nsresult NewRemovalIndices(const uint32_t* aIndices, size_t aNumOfIndices); + nsresult NewFullHashResponse(const Prefix& aPrefix, + const CachedFullHashResponse& aResponse); + + // Clear Prefixes & Removal indice. + void Clear(); + + private: + virtual int Tag() const override { return TAG; } + + bool mFullUpdate; + PrefixStringMap mPrefixesMap; + RemovalIndiceArray mRemovalIndiceArray; + nsCString mClientState; + nsCString mSHA256; + + // This is used to store response from fullHashes.find. + FullHashResponseMap mFullHashResponseMap; +}; + +// There is one hash store per table. +class HashStore { + public: + HashStore(const nsACString& aTableName, const nsACString& aProvider, + nsIFile* aRootStoreFile); + ~HashStore(); + + const nsCString& TableName() const { return mTableName; } + + // Version is set to 0 by default, it is only used when we want to open + // a specific version of HashStore. Note that the intention of aVersion + // is only to pass SanityCheck, reading data from older version should + // be handled additionally. + nsresult Open(uint32_t aVersion = 0); + + // Add Prefixes/Completes are stored partly in the PrefixSet (contains the + // Prefix data organized for fast lookup/low RAM usage) and partly in the + // HashStore (Add Chunk numbers - only used for updates, slow retrieval). + // AugmentAdds function joins the separate datasets into one complete + // prefixes+chunknumbers dataset. + nsresult AugmentAdds(const nsTArray<uint32_t>& aPrefixes, + const nsTArray<nsCString>& aCompletes); + + ChunkSet& AddChunks(); + ChunkSet& SubChunks(); + AddPrefixArray& AddPrefixes() { return mAddPrefixes; } + SubPrefixArray& SubPrefixes() { return mSubPrefixes; } + AddCompleteArray& AddCompletes() { return mAddCompletes; } + SubCompleteArray& SubCompletes() { return mSubCompletes; } + + // ======= + // Updates + // ======= + // Begin the update process. Reads the store into memory. + nsresult BeginUpdate(); + + // Imports the data from a TableUpdate. + nsresult ApplyUpdate(RefPtr<TableUpdateV2> aUpdate); + + // Process expired chunks + nsresult Expire(); + + // Rebuild the store, Incorporating all the applied updates. + nsresult Rebuild(); + + // Write the current state of the store to disk. + // If you call between ApplyUpdate() and Rebuild(), you'll + // have a mess on your hands. + nsresult WriteFile(); + + nsresult ReadCompletionsLegacyV3(AddCompleteArray& aCompletes); + + nsresult Reset(); + + private: + nsresult ReadHeader(); + nsresult SanityCheck(uint32_t aVersion = 0) const; + nsresult CalculateChecksum(nsAutoCString& aChecksum, uint32_t aFileSize, + bool aChecksumPresent); + nsresult CheckChecksum(uint32_t aFileSize); + void UpdateHeader(); + + nsresult ReadCompletions(); + nsresult ReadChunkNumbers(); + nsresult ReadHashes(); + + nsresult ReadAddPrefixes(); + nsresult ReadSubPrefixes(); + nsresult ReadAddCompletes(); + + nsresult WriteAddPrefixChunks(nsIOutputStream* aOut); + nsresult WriteSubPrefixes(nsIOutputStream* aOut); + nsresult WriteAddCompleteChunks(nsIOutputStream* aOut); + + nsresult ProcessSubs(); + + nsresult PrepareForUpdate(); + + // This is used for checking that the database is correct and for figuring out + // the number of chunks, etc. to read from disk on restart. + struct Header { + uint32_t magic; + uint32_t version; + uint32_t numAddChunks; + uint32_t numSubChunks; + uint32_t numAddPrefixes; + uint32_t numSubPrefixes; + uint32_t numAddCompletes; + uint32_t numSubCompletes; + }; + + Header mHeader; + + // The name of the table (must end in -shavar or -digest256, or evidently + // -simple for unittesting. + const nsCString mTableName; + nsCOMPtr<nsIFile> mStoreDirectory; + + bool mInUpdate; + + nsCOMPtr<nsIInputStream> mInputStream; + + // Chunk numbers, stored as uint32_t arrays. + ChunkSet mAddChunks; + ChunkSet mSubChunks; + + ChunkSet mAddExpirations; + ChunkSet mSubExpirations; + + // Chunk data for shavar tables. See Entries.h for format. + AddPrefixArray mAddPrefixes; + SubPrefixArray mSubPrefixes; + + // See bug 806422 for background. We must be able to distinguish between + // updates from the completion server and updates from the regular server. + AddCompleteArray mAddCompletes; + SubCompleteArray mSubCompletes; + + uint32_t mFileSize; + + // For gtest to inspect private members. + friend class PerProviderDirectoryTestUtils; +}; + +} // namespace safebrowsing +} // namespace mozilla + +#endif |