1 files changed, 438 insertions, 0 deletions
diff --git a/toolkit/components/url-classifier/tests/unit/test_hashcompleter.js b/toolkit/components/url-classifier/tests/unit/test_hashcompleter.js
new file mode 100644
index 0000000000..5837753635
--- /dev/null
+++ b/toolkit/components/url-classifier/tests/unit/test_hashcompleter.js
@@ -0,0 +1,438 @@
+/* Any copyright is dedicated to the Public Domain.
+   http://creativecommons.org/publicdomain/zero/1.0/ */
+
+// This test ensures that the nsIUrlClassifierHashCompleter works as expected
+// and simulates an HTTP server to provide completions.
+//
+// In order to test completions, each group of completions sent as one request
+// to the HTTP server is called a completion set. There is currently not
+// support for multiple requests being sent to the server at once, in this test.
+// This tests makes a request for each element of |completionSets|, waits for
+// a response and then moves to the next element.
+// Each element of |completionSets| is an array of completions, and each
+// completion is an object with the properties:
+//   hash: complete hash for the completion. Automatically right-padded
+//         to be COMPLETE_LENGTH.
+//   expectCompletion: boolean indicating whether the server should respond
+//                     with a full hash.
+//   forceServerError: boolean indicating whether the server should respond
+//                     with a 503.
+//   table: name of the table that the hash corresponds to. Only needs to be set
+//          if a completion is expected.
+//   chunkId: positive integer corresponding to the chunk that the hash belongs
+//            to. Only needs to be set if a completion is expected.
+//   multipleCompletions: boolean indicating whether the server should respond
+//                        with more than one full hash. If this is set to true
+//                        then |expectCompletion| must also be set to true and
+//                        |hash| must have the same prefix as all |completions|.
+//   completions: an array of completions (objects with a hash, table and
+//                chunkId property as described above). This property is only
+//                used when |multipleCompletions| is set to true.
+
+// Basic prefixes with 2/3 completions.
+var basicCompletionSet = [
+  {
+    hash: "abcdefgh",
+    expectCompletion: true,
+    table: "test",
+    chunkId: 1234,
+  },
+  {
+    hash: "1234",
+    expectCompletion: false,
+  },
+  {
+    hash: "\u0000\u0000\u000012312",
+    expectCompletion: true,
+    table: "test",
+    chunkId: 1234,
+  },
+];
+
+// 3 prefixes with 0 completions to test HashCompleter handling a 204 status.
+var falseCompletionSet = [
+  {
+    hash: "1234",
+    expectCompletion: false,
+  },
+  {
+    hash: "",
+    expectCompletion: false,
+  },
+  {
+    hash: "abc",
+    expectCompletion: false,
+  },
+];
+
+// The current implementation (as of Mar 2011) sometimes sends duplicate
+// entries to HashCompleter and even expects responses for duplicated entries.
+var dupedCompletionSet = [
+  {
+    hash: "1234",
+    expectCompletion: true,
+    table: "test",
+    chunkId: 1,
+  },
+  {
+    hash: "5678",
+    expectCompletion: false,
+    table: "test2",
+    chunkId: 2,
+  },
+  {
+    hash: "1234",
+    expectCompletion: true,
+    table: "test",
+    chunkId: 1,
+  },
+  {
+    hash: "5678",
+    expectCompletion: false,
+    table: "test2",
+    chunkId: 2,
+  },
+];
+
+// It is possible for a hash completion request to return with multiple
+// completions, the HashCompleter should return all of these.
+var multipleResponsesCompletionSet = [
+  {
+    hash: "1234",
+    expectCompletion: true,
+    multipleCompletions: true,
+    completions: [
+      {
+        hash: "123456",
+        table: "test1",
+        chunkId: 3,
+      },
+      {
+        hash: "123478",
+        table: "test2",
+        chunkId: 4,
+      },
+    ],
+  },
+];
+
+function buildCompletionRequest(aCompletionSet) {
+  let prefixes = [];
+  let prefixSet = new Set();
+  aCompletionSet.forEach(s => {
+    let prefix = s.hash.substring(0, 4);
+    if (prefixSet.has(prefix)) {
+      return;
+    }
+    prefixSet.add(prefix);
+    prefixes.push(prefix);
+  });
+  return 4 + ":" + 4 * prefixes.length + "\n" + prefixes.join("");
+}
+
+function parseCompletionRequest(aRequest) {
+  // Format: [partial_length]:[num_of_prefix * partial_length]\n[prefixes_data]
+
+  let tokens = /(\d):(\d+)/.exec(aRequest);
+  if (tokens.length < 3) {
+    dump("Request format error.");
+    return null;
+  }
+
+  let partialLength = parseInt(tokens[1]);
+
+  let payloadStart =
+    tokens[1].length + // partial length
+    1 + // ':'
+    tokens[2].length + // payload length
+    1; // '\n'
+
+  let prefixSet = [];
+  for (let i = payloadStart; i < aRequest.length; i += partialLength) {
+    let prefix = aRequest.substr(i, partialLength);
+    if (prefix.length !== partialLength) {
+      dump("Header info not correct: " + aRequest.substr(0, payloadStart));
+      return null;
+    }
+    prefixSet.push(prefix);
+  }
+  prefixSet.sort();
+
+  return prefixSet;
+}
+
+// Compare the requests in string format.
+function compareCompletionRequest(aRequest1, aRequest2) {
+  let prefixSet1 = parseCompletionRequest(aRequest1);
+  let prefixSet2 = parseCompletionRequest(aRequest2);
+
+  return equal(JSON.stringify(prefixSet1), JSON.stringify(prefixSet2));
+}
+
+// The fifth completion set is added at runtime by getRandomCompletionSet.
+// Each completion in the set only has one response and its purpose is to
+// provide an easy way to test the HashCompleter handling an arbitrarily large
+// completion set (determined by SIZE_OF_RANDOM_SET).
+const SIZE_OF_RANDOM_SET = 16;
+function getRandomCompletionSet(forceServerError) {
+  let completionSet = [];
+  let hashPrefixes = [];
+
+  let seed = Math.floor(Math.random() * Math.pow(2, 32));
+  dump("Using seed of " + seed + " for random completion set.\n");
+  let rand = new LFSRgenerator(seed);
+
+  for (let i = 0; i < SIZE_OF_RANDOM_SET; i++) {
+    let completion = {
+      expectCompletion: false,
+      forceServerError: false,
+      _finished: false,
+    };
+
+    // Generate a random 256 bit hash. First we get a random number and then
+    // convert it to a string.
+    let hash;
+    let prefix;
+    do {
+      hash = "";
+      let length = 1 + rand.nextNum(5);
+      for (let j = 0; j < length; j++) {
+        hash += String.fromCharCode(rand.nextNum(8));
+      }
+      prefix = hash.substring(0, 4);
+    } while (hashPrefixes.includes(prefix));
+
+    hashPrefixes.push(prefix);
+    completion.hash = hash;
+
+    if (!forceServerError) {
+      completion.expectCompletion = rand.nextNum(1) == 1;
+    } else {
+      completion.forceServerError = true;
+    }
+    if (completion.expectCompletion) {
+      // Generate a random alpha-numeric string of length start with "test" for the
+      // table name.
+      completion.table = "test" + rand.nextNum(31).toString(36);
+
+      completion.chunkId = rand.nextNum(16);
+    }
+    completionSet.push(completion);
+  }
+
+  return completionSet;
+}
+
+var completionSets = [
+  basicCompletionSet,
+  falseCompletionSet,
+  dupedCompletionSet,
+  multipleResponsesCompletionSet,
+];
+var currentCompletionSet = -1;
+var finishedCompletions = 0;
+
+const SERVER_PATH = "/hash-completer";
+var server;
+
+// Completion hashes are automatically right-padded with null chars to have a
+// length of COMPLETE_LENGTH.
+// Taken from nsUrlClassifierDBService.h
+const COMPLETE_LENGTH = 32;
+
+var completer = Cc["@mozilla.org/url-classifier/hashcompleter;1"].getService(
+  Ci.nsIUrlClassifierHashCompleter
+);
+
+var gethashUrl;
+
+// Expected highest completion set for which the server sends a response.
+var expectedMaxServerCompletionSet = 0;
+var maxServerCompletionSet = 0;
+
+function run_test() {
+  // This test case exercises the backoff functionality so we can't leave it disabled.
+  Services.prefs.setBoolPref(
+    "browser.safebrowsing.provider.test.disableBackoff",
+    false
+  );
+  // Generate a random completion set that return successful responses.
+  completionSets.push(getRandomCompletionSet(false));
+  // We backoff after receiving an error, so requests shouldn't reach the
+  // server after that.
+  expectedMaxServerCompletionSet = completionSets.length;
+  // Generate some completion sets that return 503s.
+  for (let j = 0; j < 10; ++j) {
+    completionSets.push(getRandomCompletionSet(true));
+  }
+
+  // Fix up the completions before running the test.
+  for (let completionSet of completionSets) {
+    for (let completion of completionSet) {
+      // Pad the right of each |hash| so that the length is COMPLETE_LENGTH.
+      if (completion.multipleCompletions) {
+        for (let responseCompletion of completion.completions) {
+          let numChars = COMPLETE_LENGTH - responseCompletion.hash.length;
+          responseCompletion.hash += new Array(numChars + 1).join("\u0000");
+        }
+      } else {
+        let numChars = COMPLETE_LENGTH - completion.hash.length;
+        completion.hash += new Array(numChars + 1).join("\u0000");
+      }
+    }
+  }
+  do_test_pending();
+
+  server = new HttpServer();
+  server.registerPathHandler(SERVER_PATH, hashCompleterServer);
+
+  server.start(-1);
+  const SERVER_PORT = server.identity.primaryPort;
+
+  gethashUrl = "http://localhost:" + SERVER_PORT + SERVER_PATH;
+
+  runNextCompletion();
+}
+
+function runNextCompletion() {
+  // The server relies on currentCompletionSet to send the correct response, so
+  // don't increment it until we start the new set of callbacks.
+  currentCompletionSet++;
+  if (currentCompletionSet >= completionSets.length) {
+    finish();
+    return;
+  }
+
+  dump(
+    "Now on completion set index " +
+      currentCompletionSet +
+      ", length " +
+      completionSets[currentCompletionSet].length +
+      "\n"
+  );
+  // Number of finished completions for this set.
+  finishedCompletions = 0;
+  for (let completion of completionSets[currentCompletionSet]) {
+    completer.complete(
+      completion.hash.substring(0, 4),
+      gethashUrl,
+      "test-phish-shavar", // Could be arbitrary v2 table name.
+      new callback(completion)
+    );
+  }
+}
+
+function hashCompleterServer(aRequest, aResponse) {
+  let stream = aRequest.bodyInputStream;
+  let wrapperStream = Cc["@mozilla.org/binaryinputstream;1"].createInstance(
+    Ci.nsIBinaryInputStream
+  );
+  wrapperStream.setInputStream(stream);
+
+  let len = stream.available();
+  let data = wrapperStream.readBytes(len);
+
+  // Check if we got the expected completion request.
+  let expectedRequest = buildCompletionRequest(
+    completionSets[currentCompletionSet]
+  );
+  compareCompletionRequest(data, expectedRequest);
+
+  // To avoid a response with duplicate hash completions, we keep track of all
+  // completed hash prefixes so far.
+  let completedHashes = [];
+  let responseText = "";
+
+  function responseForCompletion(x) {
+    return x.table + ":" + x.chunkId + ":" + x.hash.length + "\n" + x.hash;
+  }
+  // As per the spec, a server should response with a 204 if there are no
+  // full-length hashes that match the prefixes.
+  let httpStatus = 204;
+  for (let completion of completionSets[currentCompletionSet]) {
+    if (
+      completion.expectCompletion &&
+      !completedHashes.includes(completion.hash)
+    ) {
+      completedHashes.push(completion.hash);
+
+      if (completion.multipleCompletions) {
+        responseText += completion.completions
+          .map(responseForCompletion)
+          .join("");
+      } else {
+        responseText += responseForCompletion(completion);
+      }
+    }
+    if (completion.forceServerError) {
+      httpStatus = 503;
+    }
+  }
+
+  dump("Server sending response for " + currentCompletionSet + "\n");
+  maxServerCompletionSet = currentCompletionSet;
+  if (responseText && httpStatus != 503) {
+    aResponse.write(responseText);
+  } else {
+    aResponse.setStatusLine(null, httpStatus, null);
+  }
+}
+
+function callback(completion) {
+  this._completion = completion;
+}
+
+callback.prototype = {
+  completionV2: function completionV2(hash, table, chunkId, trusted) {
+    Assert.ok(this._completion.expectCompletion);
+    if (this._completion.multipleCompletions) {
+      for (let completion of this._completion.completions) {
+        if (completion.hash == hash) {
+          Assert.equal(JSON.stringify(hash), JSON.stringify(completion.hash));
+          Assert.equal(table, completion.table);
+          Assert.equal(chunkId, completion.chunkId);
+
+          completion._completed = true;
+
+          if (this._completion.completions.every(x => x._completed)) {
+            this._completed = true;
+          }
+
+          break;
+        }
+      }
+    } else {
+      // Hashes are not actually strings and can contain arbitrary data.
+      Assert.equal(JSON.stringify(hash), JSON.stringify(this._completion.hash));
+      Assert.equal(table, this._completion.table);
+      Assert.equal(chunkId, this._completion.chunkId);
+
+      this._completed = true;
+    }
+  },
+
+  completionFinished: function completionFinished(status) {
+    finishedCompletions++;
+    Assert.equal(!!this._completion.expectCompletion, !!this._completed);
+    this._completion._finished = true;
+
+    // currentCompletionSet can mutate before all of the callbacks are complete.
+    if (
+      currentCompletionSet < completionSets.length &&
+      finishedCompletions == completionSets[currentCompletionSet].length
+    ) {
+      runNextCompletion();
+    }
+  },
+};
+
+function finish() {
+  Services.prefs.clearUserPref(
+    "browser.safebrowsing.provider.test.disableBackoff"
+  );
+
+  Assert.equal(expectedMaxServerCompletionSet, maxServerCompletionSet);
+  server.stop(function() {
+    do_test_finished();
+  });
+}