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 "HashStore.h"
7 #include "mozilla/Unused.h"
8 #include "nsPrintfCString.h"
9 #include "string.h"
10 
11 #include "Common.h"
12 
13 static const char* kFilesInV2[] = {".vlpset", ".sbstore"};
14 static const char* kFilesInV4[] = {".vlpset", ".metadata"};
15 
16 #define GTEST_MALWARE_TABLE_V4 "goog-malware-proto"_ns
17 #define GTEST_PHISH_TABLE_V4 "goog-phish-proto"_ns
18 
19 #define ROOT_DIR u"safebrowsing"_ns
20 #define SB_FILE(x, y) NS_ConvertUTF8toUTF16(nsPrintfCString("%s%s", x, y))
21 
22 template <typename T, size_t N>
CheckFileExist(const nsCString & aTable,const T (& aFiles)[N],bool aExpectExists,const char * aMsg=nullptr)23 static void CheckFileExist(const nsCString& aTable, const T (&aFiles)[N],
24                            bool aExpectExists, const char* aMsg = nullptr) {
25   for (uint32_t i = 0; i < N; i++) {
26     // This is just a quick way to know if this is v4 table
27     NS_ConvertUTF8toUTF16 SUB_DIR(strstr(aTable.get(), "-proto") ? "google4"
28                                                                  : "");
29     nsCOMPtr<nsIFile> file = GetFile(nsTArray<nsString>{
30         ROOT_DIR, SUB_DIR, SB_FILE(aTable.get(), aFiles[i])});
31 
32     bool exists;
33     file->Exists(&exists);
34 
35     if (aMsg) {
36       ASSERT_EQ(aExpectExists, exists)
37           << file->HumanReadablePath().get() << " " << aMsg;
38     } else {
39       ASSERT_EQ(aExpectExists, exists) << file->HumanReadablePath().get();
40     }
41   }
42 }
43 
TEST(UrlClassifierFailUpdate,CheckTableReset)44 TEST(UrlClassifierFailUpdate, CheckTableReset)
45 {
46   const bool FULL_UPDATE = true;
47   const bool PARTIAL_UPDATE = false;
48 
49   // Apply V2 update
50   {
51     RefPtr<TableUpdateV2> update = new TableUpdateV2(GTEST_TABLE_V2);
52     Unused << update->NewAddChunk(1);
53 
54     ApplyUpdate(update);
55 
56     // A successful V2 update should create .vlpset & .sbstore files
57     CheckFileExist(GTEST_TABLE_V2, kFilesInV2, true,
58                    "V2 update doesn't create vlpset or sbstore");
59   }
60 
61   // Helper function to generate table update data
62   auto func = [](RefPtr<TableUpdateV4> update, bool full, const char* str) {
63     update->SetFullUpdate(full);
64     nsCString prefix(str);
65     update->NewPrefixes(prefix.Length(), prefix);
66   };
67 
68   // Apply V4 update for table1
69   {
70     RefPtr<TableUpdateV4> update = new TableUpdateV4(GTEST_MALWARE_TABLE_V4);
71     func(update, FULL_UPDATE, "test_prefix");
72 
73     ApplyUpdate(update);
74 
75     // A successful V4 update should create .vlpset & .metadata files
76     CheckFileExist(GTEST_MALWARE_TABLE_V4, kFilesInV4, true,
77                    "v4 update doesn't create vlpset or metadata");
78   }
79 
80   // Apply V4 update for table2
81   {
82     RefPtr<TableUpdateV4> update = new TableUpdateV4(GTEST_PHISH_TABLE_V4);
83     func(update, FULL_UPDATE, "test_prefix");
84 
85     ApplyUpdate(update);
86 
87     CheckFileExist(GTEST_PHISH_TABLE_V4, kFilesInV4, true,
88                    "v4 update doesn't create vlpset or metadata");
89   }
90 
91   // Apply V4 update with the same prefix in previous full udpate
92   // This should cause an update error.
93   {
94     RefPtr<TableUpdateV4> update = new TableUpdateV4(GTEST_MALWARE_TABLE_V4);
95     func(update, PARTIAL_UPDATE, "test_prefix");
96 
97     ApplyUpdate(update);
98 
99     // A fail update should remove files for that table
100     CheckFileExist(GTEST_MALWARE_TABLE_V4, kFilesInV4, false,
101                    "a fail v4 update doesn't remove the tables");
102 
103     // A fail update should NOT remove files for the other tables
104     CheckFileExist(GTEST_TABLE_V2, kFilesInV2, true,
105                    "a fail v4 update removes a v2 table");
106     CheckFileExist(GTEST_PHISH_TABLE_V4, kFilesInV4, true,
107                    "a fail v4 update removes the other v4 table");
108   }
109 }
110