michael@0: Cu.import("resource://gre/modules/XPCOMUtils.jsm");
michael@0: 
michael@0: XPCOMUtils.defineLazyModuleGetter(this, "NetUtil",
michael@0:                                   "resource://gre/modules/NetUtil.jsm");
michael@0: XPCOMUtils.defineLazyModuleGetter(this, "Promise",
michael@0:                                   "resource://gre/modules/Promise.jsm");
michael@0: // Global test server for serving safebrowsing updates.
michael@0: let gHttpServ = null;
michael@0: // Global nsIUrlClassifierDBService
michael@0: let gDbService = Cc["@mozilla.org/url-classifier/dbservice;1"]
michael@0:   .getService(Ci.nsIUrlClassifierDBService);
michael@0: // Security manager for creating nsIPrincipals from URIs
michael@0: let gSecMan = Cc["@mozilla.org/scriptsecuritymanager;1"]
michael@0:   .getService(Ci.nsIScriptSecurityManager);
michael@0: 
michael@0: // A map of tables to arrays of update redirect urls.
michael@0: let gTables = {};
michael@0: 
michael@0: // Construct an update from a file.
michael@0: function readFileToString(aFilename) {
michael@0:   let f = do_get_file(aFilename);
michael@0:   let stream = Cc["@mozilla.org/network/file-input-stream;1"]
michael@0:     .createInstance(Ci.nsIFileInputStream);
michael@0:   stream.init(f, -1, 0, 0);
michael@0:   let buf = NetUtil.readInputStreamToString(stream, stream.available());
michael@0:   return buf;
michael@0: }
michael@0: 
michael@0: // Registers a table for which to serve update chunks. Returns a promise that
michael@0: // resolves when that chunk has been downloaded.
michael@0: function registerTableUpdate(aTable, aFilename) {
michael@0:   let deferred = Promise.defer();
michael@0:   // If we haven't been given an update for this table yet, add it to the map
michael@0:   if (!(aTable in gTables)) {
michael@0:     gTables[aTable] = [];
michael@0:   }
michael@0: 
michael@0:   // The number of chunks associated with this table.
michael@0:   let numChunks = gTables[aTable].length + 1;
michael@0:   let redirectPath = "/" + aTable + "-" + numChunks;
michael@0:   let redirectUrl = "localhost:4444" + redirectPath;
michael@0: 
michael@0:   // Store redirect url for that table so we can return it later when we
michael@0:   // process an update request.
michael@0:   gTables[aTable].push(redirectUrl);
michael@0: 
michael@0:   gHttpServ.registerPathHandler(redirectPath, function(request, response) {
michael@0:     do_print("Mock safebrowsing server handling request for " + redirectPath);
michael@0:     let contents = readFileToString(aFilename);
michael@0:     response.setHeader("Content-Type",
michael@0:                        "application/vnd.google.safebrowsing-update", false);
michael@0:     response.setStatusLine(request.httpVersion, 200, "OK");
michael@0:     response.bodyOutputStream.write(contents, contents.length);
michael@0:     deferred.resolve(contents);
michael@0:   });
michael@0:   return deferred.promise;
michael@0: }
michael@0: 
michael@0: // Construct a response with redirect urls.
michael@0: function processUpdateRequest() {
michael@0:   let response = "n:1000\n";
michael@0:   for (let table in gTables) {
michael@0:     response += "i:" + table + "\n";
michael@0:     for (let i = 0; i < gTables[table].length; ++i) {
michael@0:       response += "u:" + gTables[table][i] + "\n";
michael@0:     }
michael@0:   }
michael@0:   do_print("Returning update response: " + response);
michael@0:   return response;
michael@0: }
michael@0: 
michael@0: // Set up our test server to handle update requests.
michael@0: function run_test() {
michael@0:   gHttpServ = new HttpServer();
michael@0:   gHttpServ.registerDirectory("/", do_get_cwd());
michael@0: 
michael@0:   gHttpServ.registerPathHandler("/downloads", function(request, response) {
michael@0:     let buf = NetUtil.readInputStreamToString(request.bodyInputStream,
michael@0:       request.bodyInputStream.available());
michael@0:     let blob = processUpdateRequest();
michael@0:     response.setHeader("Content-Type",
michael@0:                        "application/vnd.google.safebrowsing-update", false);
michael@0:     response.setStatusLine(request.httpVersion, 200, "OK");
michael@0:     response.bodyOutputStream.write(blob, blob.length);
michael@0:   });
michael@0: 
michael@0:   gHttpServ.start(4444);
michael@0:   run_next_test();
michael@0: }
michael@0: 
michael@0: function createURI(s) {
michael@0:   let service = Cc["@mozilla.org/network/io-service;1"]
michael@0:     .getService(Ci.nsIIOService);
michael@0:   return service.newURI(s, null, null);
michael@0: }
michael@0: 
michael@0: // Just throw if we ever get an update or download error.
michael@0: function handleError(aEvent) {
michael@0:   do_throw("We didn't download or update correctly: " + aEvent);
michael@0: }
michael@0: 
michael@0: add_test(function test_update() {
michael@0:   let streamUpdater = Cc["@mozilla.org/url-classifier/streamupdater;1"]
michael@0:     .getService(Ci.nsIUrlClassifierStreamUpdater);
michael@0:   streamUpdater.updateUrl = "http://localhost:4444/downloads";
michael@0: 
michael@0:   // Load up some update chunks for the safebrowsing server to serve.
michael@0:   registerTableUpdate("goog-downloadwhite-digest256", "data/digest1.chunk");
michael@0:   registerTableUpdate("goog-downloadwhite-digest256", "data/digest2.chunk");
michael@0: 
michael@0:   // Download some updates, and don't continue until the downloads are done.
michael@0:   function updateSuccess(aEvent) {
michael@0:     // Timeout of n:1000 is constructed in processUpdateRequest above and
michael@0:     // passed back in the callback in nsIUrlClassifierStreamUpdater on success.
michael@0:     do_check_eq("1000", aEvent);
michael@0:     do_print("All data processed");
michael@0:     run_next_test();
michael@0:   }
michael@0:   streamUpdater.downloadUpdates(
michael@0:     "goog-downloadwhite-digest256",
michael@0:     "goog-downloadwhite-digest256;\n",
michael@0:     updateSuccess, handleError, handleError);
michael@0: });
michael@0: 
michael@0: add_test(function test_url_not_whitelisted() {
michael@0:   let uri = createURI("http://example.com");
michael@0:   let principal = gSecMan.getNoAppCodebasePrincipal(uri);
michael@0:   gDbService.lookup(principal, "goog-downloadwhite-digest256",
michael@0:     function handleEvent(aEvent) {
michael@0:       // This URI is not on any lists.
michael@0:       do_check_eq("", aEvent);
michael@0:       run_next_test();
michael@0:     });
michael@0: });
michael@0: 
michael@0: add_test(function test_url_whitelisted() {
michael@0:   // Hash of "whitelisted.com/" (canonicalized URL) is:
michael@0:   // 93CA5F48E15E9861CD37C2D95DB43D23CC6E6DE5C3F8FA6E8BE66F97CC518907
michael@0:   let uri = createURI("http://whitelisted.com");
michael@0:   let principal = gSecMan.getNoAppCodebasePrincipal(uri);
michael@0:   gDbService.lookup(principal, "goog-downloadwhite-digest256",
michael@0:     function handleEvent(aEvent) {
michael@0:       do_check_eq("goog-downloadwhite-digest256", aEvent);
michael@0:       run_next_test();
michael@0:     });
michael@0: });