fune/toolkit/components/url-classifier/tests/unit/test_listmanager.js

Cu.import("resource://gre/modules/XPCOMUtils.jsm");
Cu.import("resource://gre/modules/Services.jsm");

XPCOMUtils.defineLazyModuleGetter(this, "NetUtil",
                                  "resource://gre/modules/NetUtil.jsm");

// These tables share the same updateURL.
const TEST_TABLE_DATA_LIST = [
  // 0:
  {
    tableName: "test-listmanager0-digest256",
    providerName: "google",
    updateUrl: "http://localhost:4444/safebrowsing/update",
    gethashUrl: "http://localhost:4444/safebrowsing/gethash0",
  },

  // 1:
  {
    tableName: "test-listmanager1-digest256",
    providerName: "google",
    updateUrl: "http://localhost:4444/safebrowsing/update",
    gethashUrl: "http://localhost:4444/safebrowsing/gethash1",
  },

  // 2.
  {
    tableName: "test-listmanager2-digest256",
    providerName: "google",
    updateUrl: "http://localhost:4444/safebrowsing/update",
    gethashUrl: "http://localhost:4444/safebrowsing/gethash2",
  }
];

// This table has a different update URL (for v4).
const TEST_TABLE_DATA_V4 = {
  tableName: "test-phish-proto",
  providerName: "google4",
  updateUrl: "http://localhost:5555/safebrowsing/update?",
  gethashUrl: "http://localhost:5555/safebrowsing/gethash-v4",
};

const PREF_NEXTUPDATETIME = "browser.safebrowsing.provider.google.nextupdatetime";
const PREF_NEXTUPDATETIME_V4 = "browser.safebrowsing.provider.google4.nextupdatetime";

let gListManager = Cc["@mozilla.org/url-classifier/listmanager;1"]
                     .getService(Ci.nsIUrlListManager);

let gUrlUtils = Cc["@mozilla.org/url-classifier/utils;1"]
                   .getService(Ci.nsIUrlClassifierUtils);

// Global test server for serving safebrowsing updates.
let gHttpServ = null;
let gUpdateResponse = "";
let gExpectedUpdateRequest = "";
let gExpectedQueryV4 = "";

// Handles request for TEST_TABLE_DATA_V4.
let gHttpServV4 = null;

// These two variables are used to synchronize the last two racing updates
// (in terms of "update URL") in test_update_all_tables().
let gUpdatedCntForTableData = 0; // For TEST_TABLE_DATA_LIST.
let gIsV4Updated = false;   // For TEST_TABLE_DATA_V4.

prefBranch.setBoolPref("browser.safebrowsing.debug", true);

// Register tables.
TEST_TABLE_DATA_LIST.forEach(function(t) {
  gListManager.registerTable(t.tableName,
                             t.providerName,
                             t.updateUrl,
                             t.gethashUrl);
});

gListManager.registerTable(TEST_TABLE_DATA_V4.tableName,
                           TEST_TABLE_DATA_V4.providerName,
                           TEST_TABLE_DATA_V4.updateUrl,
                           TEST_TABLE_DATA_V4.gethashUrl);

const SERVER_INVOLVED_TEST_CASE_LIST = [
  // - Do table0 update.
  // - Server would respond "a:5:32:32\n[DATA]".
  function test_update_table0() {
    disableAllUpdates();

    gListManager.enableUpdate(TEST_TABLE_DATA_LIST[0].tableName);
    gExpectedUpdateRequest = TEST_TABLE_DATA_LIST[0].tableName + ";\n";

    gUpdateResponse = "n:1000\ni:" + TEST_TABLE_DATA_LIST[0].tableName + "\n";
    gUpdateResponse += readFileToString("data/digest2.chunk");

    forceTableUpdate();
  },

  // - Do table0 update again. Since chunk 5 was added to table0 in the last
  //   update, the expected request contains "a:5".
  // - Server would respond "s;2-12\n[DATA]".
  function test_update_table0_with_existing_chunks() {
    disableAllUpdates();

    gListManager.enableUpdate(TEST_TABLE_DATA_LIST[0].tableName);
    gExpectedUpdateRequest = TEST_TABLE_DATA_LIST[0].tableName + ";a:5\n";

    gUpdateResponse = "n:1000\ni:" + TEST_TABLE_DATA_LIST[0].tableName + "\n";
    gUpdateResponse += readFileToString("data/digest1.chunk");

    forceTableUpdate();
  },

  // - Do all-table update.
  // - Server would respond no chunk control.
  //
  // Note that this test MUST be the last one in the array since we rely on
  // the number of sever-involved test case to synchronize the racing last
  // two udpates for different URL.
  function test_update_all_tables() {
    disableAllUpdates();

    // Enable all tables including TEST_TABLE_DATA_V4!
    TEST_TABLE_DATA_LIST.forEach(function(t) {
      gListManager.enableUpdate(t.tableName);
    });
    gListManager.enableUpdate(TEST_TABLE_DATA_V4.tableName);

    // Expected results for v2.
    gExpectedUpdateRequest = TEST_TABLE_DATA_LIST[0].tableName + ";a:5:s:2-12\n" +
                             TEST_TABLE_DATA_LIST[1].tableName + ";\n" +
                             TEST_TABLE_DATA_LIST[2].tableName + ";\n";
    gUpdateResponse = "n:1000\n";

    // We test the request against the query string since v4 request
    // would be appened to the query string. The request is generated
    // by protobuf API (binary) then encoded to base64 format.
    let requestV4 = gUrlUtils.makeUpdateRequestV4([TEST_TABLE_DATA_V4.tableName],
                                                  [""],
                                                  1);
    gExpectedQueryV4 = "&$req=" + btoa(requestV4);

    forceTableUpdate();
  },

];

SERVER_INVOLVED_TEST_CASE_LIST.forEach(t => add_test(t));

// Tests nsIUrlListManager.getGethashUrl.
add_test(function test_getGethashUrl() {
  TEST_TABLE_DATA_LIST.forEach(function (t) {
    equal(gListManager.getGethashUrl(t.tableName), t.gethashUrl);
  });
  equal(gListManager.getGethashUrl(TEST_TABLE_DATA_V4.tableName),
        TEST_TABLE_DATA_V4.gethashUrl);
  run_next_test();
});

function run_test() {
  // Setup primary testing server.
  gHttpServ = new HttpServer();
  gHttpServ.registerDirectory("/", do_get_cwd());

  gHttpServ.registerPathHandler("/safebrowsing/update", function(request, response) {
    let body = NetUtil.readInputStreamToString(request.bodyInputStream,
                                               request.bodyInputStream.available());

    // Verify if the request is as expected.
    equal(body, gExpectedUpdateRequest);

    // Respond the update which is controlled by the test case.
    response.setHeader("Content-Type",
                       "application/vnd.google.safebrowsing-update", false);
    response.setStatusLine(request.httpVersion, 200, "OK");
    response.bodyOutputStream.write(gUpdateResponse, gUpdateResponse.length);

    gUpdatedCntForTableData++;

    if (gUpdatedCntForTableData !== SERVER_INVOLVED_TEST_CASE_LIST.length) {
      // This is not the last test case so run the next once upon the
      // the update success.
      waitForUpdateSuccess(run_next_test);
      return;
    }

    if (gIsV4Updated) {
      run_next_test();  // All tests are done. Just finish.
      return;
    }

    do_print("Waiting for TEST_TABLE_DATA_V4 to be tested ...");
  });

  gHttpServ.start(4444);

  // Setup v4 testing server for the different update URL.
  gHttpServV4 = new HttpServer();
  gHttpServV4.registerDirectory("/", do_get_cwd());

  gHttpServV4.registerPathHandler("/safebrowsing/update", function(request, response) {
    // V4 update request body should be empty.
    equal(request.bodyInputStream.available(), 0);

    // Not on the spec. Found in Chromium source code...
    equal(request.getHeader("X-HTTP-Method-Override"), "POST");

    // V4 update request uses GET.
    equal(request.method, "GET");

    // V4 append the base64 encoded request to the query string.
    equal(request.queryString, gExpectedQueryV4);

    // Respond a V2 compatible content for now. In the future we can
    // send a meaningful response to test Bug 1284178 to see if the
    // update is successfully stored to database.
    response.setHeader("Content-Type",
                       "application/vnd.google.safebrowsing-update", false);
    response.setStatusLine(request.httpVersion, 200, "OK");
    let content = "n:1000\n";
    response.bodyOutputStream.write(content, content.length);

    gIsV4Updated = true;

    if (gUpdatedCntForTableData === SERVER_INVOLVED_TEST_CASE_LIST.length) {
      // All tests are done!
      run_next_test();
      return;
    }

    do_print("Wait for all sever-involved tests to be done ...");
  });

  gHttpServV4.start(5555);

  run_next_test();
}

// A trick to force updating tables. However, before calling this, we have to
// call disableAllUpdates() first to clean up the updateCheckers in listmanager.
function forceTableUpdate() {
  prefBranch.setCharPref(PREF_NEXTUPDATETIME, "1");
  prefBranch.setCharPref(PREF_NEXTUPDATETIME_V4, "1");
  gListManager.maybeToggleUpdateChecking();
}

function disableAllUpdates() {
  TEST_TABLE_DATA_LIST.forEach(t => gListManager.disableUpdate(t.tableName));
  gListManager.disableUpdate(TEST_TABLE_DATA_V4.tableName);
}

// Since there's no public interface on listmanager to know the update success,
// we could only rely on the refresh of "nextupdatetime".
function waitForUpdateSuccess(callback) {
  let nextupdatetime = parseInt(prefBranch.getCharPref(PREF_NEXTUPDATETIME));
  do_print("nextupdatetime: " + nextupdatetime);
  if (nextupdatetime !== 1) {
    callback();
    return;
  }
  do_timeout(1000, waitForUpdateSuccess.bind(null, callback));
}

// Construct an update from a file.
function readFileToString(aFilename) {
  let f = do_get_file(aFilename);
  let stream = Cc["@mozilla.org/network/file-input-stream;1"]
    .createInstance(Ci.nsIFileInputStream);
  stream.init(f, -1, 0, 0);
  let buf = NetUtil.readInputStreamToString(stream, stream.available());
  return buf;
}

function buildUpdateRequestV4InBase64() {

  let request =  urlUtils.makeUpdateRequestV4([TEST_TABLE_DATA_V4.tableName],
                                              [""],
                                              1);
  return btoa(request);
}