Bug 1378010 - screenshot from command line with headless; r=mossop
authorBenjamin Dahse <bdahse@mozilla.com>
Tue, 12 Sep 2017 16:55:27 -0700
changeset 430223 2ce3e68d72429226b6fb79bbebc4b19b3d124d76
parent 430222 f70469dfd87a13d1dec503f51f463ce05fd637ea
child 430224 2b74ec6bb1d84093e3367e8903aa6d94817cfa00
push id7761
push userjlund@mozilla.com
push dateFri, 15 Sep 2017 00:19:52 +0000
treeherdermozilla-beta@c38455951db4 [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersmossop
bugs1378010
milestone57.0a1
first release with
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
last release without
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
Bug 1378010 - screenshot from command line with headless; r=mossop Add a `--screenshot` argument that implies `--headless` and is used to take a screenshot of a page from the command line. Default is a full page screenshot, but `--window-size=width[,height]` can change this. A path for the screenshot can be supplied with `--screenshot=/path/to/file`. MozReview-Commit-ID: 13tUjk2Yrsl
browser/components/nsBrowserContentHandler.js
browser/components/shell/HeadlessShell.jsm
browser/components/shell/moz.build
browser/components/shell/test/chrome.ini
browser/components/shell/test/headless.html
browser/components/shell/test/test_headless_screenshot.html
toolkit/xre/nsAppRunner.cpp
--- a/browser/components/nsBrowserContentHandler.js
+++ b/browser/components/nsBrowserContentHandler.js
@@ -1,16 +1,18 @@
 /* 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/. */
 
 Components.utils.import("resource://gre/modules/XPCOMUtils.jsm");
 Components.utils.import("resource://gre/modules/Services.jsm");
 Components.utils.import("resource://gre/modules/AppConstants.jsm");
 
+XPCOMUtils.defineLazyModuleGetter(this, "HeadlessShell",
+                                  "resource:///modules/HeadlessShell.jsm");
 XPCOMUtils.defineLazyModuleGetter(this, "LaterRun",
                                   "resource:///modules/LaterRun.jsm");
 XPCOMUtils.defineLazyModuleGetter(this, "PrivateBrowsingUtils",
                                   "resource://gre/modules/PrivateBrowsingUtils.jsm");
 XPCOMUtils.defineLazyModuleGetter(this, "RecentWindow",
                                   "resource:///modules/RecentWindow.jsm");
 XPCOMUtils.defineLazyModuleGetter(this, "ShellService",
                                   "resource:///modules/ShellService.jsm");
@@ -461,16 +463,20 @@ nsBrowserContentHandler.prototype = {
               "  --new-window <url> Open <url> in a new window.\n" +
               "  --new-tab <url>    Open <url> in a new tab.\n" +
               "  --private-window <url> Open <url> in a new private window.\n";
     if (AppConstants.platform == "win") {
       info += "  --preferences      Open Options dialog.\n";
     } else {
       info += "  --preferences      Open Preferences dialog.\n";
     }
+    if (AppConstants.platform == "win" || AppConstants.MOZ_WIDGET_GTK) {
+      info += "  --screenshot [<path>] Save screenshot to <path> or in working directory.\n";
+      info += "  --window-size width[,height] Width and optionally height of screenshot.\n";
+    }
     info += "  --search <term>    Search <term> with your default search engine.\n";
     return info;
   },
 
   /* nsIBrowserHandler */
 
   get defaultArgs() {
     var prefb = Services.prefs;
@@ -736,16 +742,21 @@ nsDefaultCommandLineHandler.prototype = 
       while ((ar = cmdLine.handleFlagWithParam("url", false))) {
         var uri = resolveURIInternal(cmdLine, ar);
         urilist.push(uri);
       }
     } catch (e) {
       Components.utils.reportError(e);
     }
 
+    if (cmdLine.findFlag("screenshot", true) != -1) {
+      HeadlessShell.handleCmdLineArgs(cmdLine, urilist.filter(shouldLoadURI).map(u => u.spec));
+      return;
+    }
+
     for (let i = 0; i < cmdLine.length; ++i) {
       var curarg = cmdLine.getArgument(i);
       if (curarg.match(/^-/)) {
         Components.utils.reportError("Warning: unrecognized command line flag " + curarg + "\n");
         // To emulate the pre-nsICommandLine behavior, we ignore
         // the argument after an unrecognized flag.
         ++i;
       } else {
new file mode 100644
--- /dev/null
+++ b/browser/components/shell/HeadlessShell.jsm
@@ -0,0 +1,161 @@
+/* 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/. */
+
+"use strict";
+
+let EXPORTED_SYMBOLS = ["HeadlessShell"];
+
+Components.utils.import("resource://gre/modules/Services.jsm");
+Components.utils.import("resource://gre/modules/XPCOMUtils.jsm");
+Components.utils.import("resource://gre/modules/osfile.jsm");
+
+const Ci = Components.interfaces;
+
+function loadContentWindow(webNavigation, uri) {
+  return new Promise((resolve, reject) => {
+    webNavigation.loadURI(uri, Ci.nsIWebNavigation.LOAD_FLAGS_NONE, null, null, null);
+    let docShell = webNavigation.QueryInterface(Ci.nsIInterfaceRequestor)
+                                .getInterface(Ci.nsIDocShell);
+    let webProgress = docShell.QueryInterface(Ci.nsIInterfaceRequestor)
+                              .getInterface(Ci.nsIWebProgress);
+    let progressListener = {
+      onLocationChange(progress, request, location, flags) {
+        // Ignore inner-frame events
+        if (progress != webProgress) {
+          return;
+        }
+        // Ignore events that don't change the document
+        if (flags & Ci.nsIWebProgressListener.LOCATION_CHANGE_SAME_DOCUMENT) {
+          return;
+        }
+        let contentWindow = docShell.QueryInterface(Ci.nsIInterfaceRequestor)
+                                    .getInterface(Ci.nsIDOMWindow);
+        webProgress.removeProgressListener(progressListener);
+        contentWindow.addEventListener("load", (event) => {
+          resolve(contentWindow);
+        }, { once: true });
+      },
+      QueryInterface: XPCOMUtils.generateQI(["nsIWebProgressListener",
+                                             "nsISupportsWeakReference"])
+    };
+    webProgress.addProgressListener(progressListener,
+                                    Ci.nsIWebProgress.NOTIFY_LOCATION);
+  });
+}
+
+async function takeScreenshot(fullWidth, fullHeight, contentWidth, contentHeight, path, url) {
+  try {
+    let windowlessBrowser = Services.appShell.createWindowlessBrowser(false);
+    var webNavigation = windowlessBrowser.QueryInterface(Ci.nsIWebNavigation);
+    let contentWindow = await loadContentWindow(webNavigation, url);
+    contentWindow.resizeTo(contentWidth, contentHeight);
+
+    let canvas = contentWindow.document.createElementNS("http://www.w3.org/1999/xhtml", "html:canvas");
+    let context = canvas.getContext("2d");
+    let width = fullWidth ? contentWindow.innerWidth + contentWindow.scrollMaxX - contentWindow.scrollMinX
+                          : contentWindow.innerWidth;
+    let height = fullHeight ? contentWindow.innerHeight + contentWindow.scrollMaxY - contentWindow.scrollMinY
+                            : contentWindow.innerHeight;
+    canvas.width = width;
+    canvas.height = height;
+    context.drawWindow(contentWindow, 0, 0, width, height, "rgb(255, 255, 255)");
+
+    function getBlob() {
+      return new Promise(resolve => canvas.toBlob(resolve));
+    }
+
+    function readBlob(blob) {
+      return new Promise(resolve => {
+        let reader = new FileReader();
+        reader.onloadend = () => resolve(reader);
+        reader.readAsArrayBuffer(blob);
+      });
+    }
+
+    let blob = await getBlob();
+    let reader = await readBlob(blob);
+    await OS.File.writeAtomic(path, new Uint8Array(reader.result), {flush: true});
+    dump("Screenshot saved to: " + path + "\n");
+  } catch (e) {
+    dump("Failure taking screenshot: " + e + "\n");
+  } finally {
+    if (webNavigation) {
+      webNavigation.close();
+    }
+  }
+}
+
+let HeadlessShell = {
+  async handleCmdLineArgs(cmdLine, URLlist) {
+    try {
+      // Don't quit even though we don't create a window
+      Services.startup.enterLastWindowClosingSurvivalArea();
+
+      // Default options
+      let fullWidth = true;
+      let fullHeight = true;
+      // Most common screen resolution of Firefox users
+      let contentWidth = 1366;
+      let contentHeight = 768;
+
+      // Parse `window-size`
+      try {
+        var dimensionsStr = cmdLine.handleFlagWithParam("window-size", true);
+      } catch (e) {
+        dump("expected format: --window-size width[,height]\n");
+        return;
+      }
+      if (dimensionsStr) {
+        let success;
+        let dimensions = dimensionsStr.split(",", 2);
+        if (dimensions.length == 1) {
+          success = dimensions[0] > 0;
+          if (success) {
+            fullWidth = false;
+            fullHeight = true;
+            contentWidth = dimensions[0];
+          }
+        } else {
+          success = dimensions[0] > 0 && dimensions[1] > 0;
+          if (success) {
+            fullWidth = false;
+            fullHeight = false;
+            contentWidth = dimensions[0];
+            contentHeight = dimensions[1];
+          }
+        }
+
+        if (!success) {
+          dump("expected format: --window-size width[,height]\n");
+          return;
+        }
+      }
+
+      // Only command line argument left should be `screenshot`
+      // There could still be URLs however
+      try {
+        var path = cmdLine.handleFlagWithParam("screenshot", true);
+        if (!cmdLine.length && !URLlist.length) {
+          URLlist.push(path); // Assume the user wanted to specify a URL
+          path = OS.Path.join(cmdLine.workingDirectory.path, "screenshot.png");
+        }
+      } catch (e) {
+        path = OS.Path.join(cmdLine.workingDirectory.path, "screenshot.png");
+        cmdLine.handleFlag("screenshot", true); // Remove `screenshot`
+      }
+
+      for (let i = 0; i < cmdLine.length; ++i) {
+        URLlist.push(cmdLine.getArgument(i)); // Assume that all remaining arguments are URLs
+      }
+
+      if (URLlist.length == 1) {
+        await takeScreenshot(fullWidth, fullHeight, contentWidth, contentHeight, path, URLlist[0]);
+      } else {
+        dump("expected exactly one URL when using `screenshot`\n");
+      }
+    } finally {
+      Services.startup.exitLastWindowClosingSurvivalArea();
+    }
+  }
+};
--- a/browser/components/shell/moz.build
+++ b/browser/components/shell/moz.build
@@ -6,16 +6,17 @@
 
 # For BinaryPath::GetLong for Windows
 LOCAL_INCLUDES += [
     '/xpcom/build'
 ]
 
 XPCSHELL_TESTS_MANIFESTS += ['test/unit/xpcshell.ini']
 BROWSER_CHROME_MANIFESTS += ['test/browser.ini']
+MOCHITEST_CHROME_MANIFESTS += ['test/chrome.ini']
 
 JAR_MANIFESTS += ['jar.mn']
 
 XPIDL_SOURCES += [
     'nsIShellService.idl',
 ]
 
 if CONFIG['MOZ_WIDGET_TOOLKIT'] == 'cocoa':
@@ -49,16 +50,17 @@ if SOURCES:
     FINAL_LIBRARY = 'browsercomps'
 
 EXTRA_COMPONENTS += [
     'nsSetDefaultBrowser.js',
     'nsSetDefaultBrowser.manifest',
 ]
 
 EXTRA_JS_MODULES += [
+    'HeadlessShell.jsm',
     'ShellService.jsm',
 ]
 
 for var in ('MOZ_APP_NAME', 'MOZ_APP_VERSION'):
     DEFINES[var] = '"%s"' % CONFIG[var]
 
 CXXFLAGS += CONFIG['TK_CFLAGS']
 
new file mode 100644
--- /dev/null
+++ b/browser/components/shell/test/chrome.ini
@@ -0,0 +1,5 @@
+[DEFAULT]
+support-files = headless.html
+
+[test_headless_screenshot.html]
+skip-if = (os != 'win' && os != 'linux')
new file mode 100644
--- /dev/null
+++ b/browser/components/shell/test/headless.html
@@ -0,0 +1,6 @@
+<html>
+<head><meta content="text/html; charset=utf-8" http-equiv="Content-Type"></head>
+<body style="background-color: rgb(0, 255, 0); color: rgb(0, 0, 255)">
+Hi
+</body>
+</html>
new file mode 100644
--- /dev/null
+++ b/browser/components/shell/test/test_headless_screenshot.html
@@ -0,0 +1,124 @@
+<!DOCTYPE HTML>
+<html>
+<!--
+https://bugzilla.mozilla.org/show_bug.cgi?id=1378010
+-->
+<head>
+  <meta charset="utf-8">
+  <title>Test for Bug 1378010</title>
+  <script type="application/javascript" src="chrome://mochikit/content/tests/SimpleTest/SimpleTest.js"></script>
+  <link rel="stylesheet" type="text/css" href="chrome://mochikit/content/tests/SimpleTest/test.css"/>
+  <script type="application/javascript">
+  Components.utils.import("resource://gre/modules/Services.jsm");
+  Components.utils.import("resource://gre/modules/Subprocess.jsm");
+  Components.utils.import("resource://gre/modules/osfile.jsm");
+
+  const screenshotPath = OS.Path.join(OS.Constants.Path.tmpDir, "headless_test_screenshot.png");
+
+  async function runFirefox(args) {
+    const Ci = Components.interfaces;
+    const XRE_EXECUTABLE_FILE = "XREExeF";
+    const firefoxExe = Services.dirsvc.get(XRE_EXECUTABLE_FILE, Ci.nsIFile).path;
+    const NS_APP_PREFS_50_FILE = "PrefF";
+    const mochiPrefsFile = Services.dirsvc.get(NS_APP_PREFS_50_FILE, Ci.nsIFile);
+    const mochiPrefsPath = mochiPrefsFile.path;
+    const mochiPrefsName = mochiPrefsFile.leafName;
+    const profilePath = OS.Path.join(OS.Constants.Path.tmpDir, "headless_test_screenshot_profile");
+    const prefsPath = OS.Path.join(profilePath, mochiPrefsName);
+    const firefoxArgs = ["-profile", profilePath, "-no-remote"];
+
+    await OS.File.makeDir(profilePath);
+    await OS.File.copy(mochiPrefsPath, prefsPath);
+    let proc = await Subprocess.call({
+      command: firefoxExe,
+      arguments: firefoxArgs.concat(args),
+    });
+    let stdout;
+    while ((stdout = await proc.stdout.readString())) {
+      dump(">>> " + stdout + "\n");
+    }
+    let {exitCode} = await proc.wait();
+    is(exitCode, 0, "Firefox process should exit with code 0");
+    await OS.File.removeDir(profilePath);
+  }
+
+  async function testFileCreationPositive(args, path) {
+    await runFirefox(args);
+
+    let saved = await OS.File.exists(path);
+    ok(saved, "A screenshot should be saved as " + path);
+    if (!saved) {
+      return;
+    }
+
+    let info = await OS.File.stat(path);
+    ok(info.size > 0, "Screenshot should not be an empty file");
+    await OS.File.remove(path);
+  }
+
+  async function testFileCreationNegative(args, path) {
+    await runFirefox(args);
+
+    let saved = await OS.File.exists(path);
+    ok(!saved, "A screenshot should not be saved");
+    await OS.File.remove(path, { ignoreAbsent: true });
+  }
+
+  async function testWindowSizePositive(width, height) {
+    let size = width + "";
+    if (height) {
+      size += "," + height;
+    }
+
+    await runFirefox(["-url", "http://mochi.test:8888/headless.html", "-screenshot", screenshotPath, "-window-size", size]);
+
+    let saved = await OS.File.exists(screenshotPath);
+    ok(saved, "A screenshot should be saved in the tmp directory");
+    if (!saved) {
+      return;
+    }
+
+    let data = await OS.File.read(screenshotPath);
+    await new Promise((resolve, reject) => {
+      let blob = new Blob([data], { type: "image/png" });
+      let reader = new FileReader();
+      reader.onloadend = function() {
+        let screenshot = new Image();
+        screenshot.onloadend = function() {
+          is(screenshot.width, width, "Screenshot should be " + width + " pixels wide");
+          if (height) {
+            is(screenshot.height, height, "Screenshot should be " + height + " pixels tall");
+          }
+          resolve();
+        };
+        screenshot.src = reader.result;
+      };
+      reader.readAsDataURL(blob);
+    });
+    await OS.File.remove(screenshotPath);
+  }
+
+  (async function() {
+    SimpleTest.waitForExplicitFinish();
+    await testFileCreationPositive(["-url", "http://mochi.test:8888/headless.html", "-screenshot", screenshotPath], screenshotPath);
+    await testFileCreationPositive(["-screenshot", "http://mochi.test:8888/headless.html"], "screenshot.png");
+    await testFileCreationPositive(["http://mochi.test:8888/headless.html", "-screenshot"], "screenshot.png");
+    await testFileCreationNegative(["-screenshot"], "screenshot.png");
+    await testFileCreationNegative(["http://mochi.test:8888/headless.html", "http://mochi.test:8888/headless.html", "-screenshot"], "screenshot.png");
+    await testWindowSizePositive(800, 600);
+    await testWindowSizePositive(1234);
+    await testFileCreationNegative(["-url", "http://mochi.test:8888/headless.html", "-screenshot", screenshotPath, "-window-size", "hello"], screenshotPath);
+    await testFileCreationNegative(["-url", "http://mochi.test:8888/headless.html", "-screenshot", screenshotPath, "-window-size", "800,"], screenshotPath);
+    SimpleTest.finish();
+  })();
+  </script>
+</head>
+<body>
+<a target="_blank" href="https://bugzilla.mozilla.org/show_bug.cgi?id=1378010">Mozilla Bug 1378010</a>
+<p id="display"></p>
+<div id="content" style="display: none">
+</div>
+<pre id="test">
+</pre>
+</body>
+</html>
--- a/toolkit/xre/nsAppRunner.cpp
+++ b/toolkit/xre/nsAppRunner.cpp
@@ -565,16 +565,60 @@ CheckArg(const char* aArg, bool aCheckOS
       ar = ARG_BAD;
       PR_fprintf(PR_STDERR, "Error: argument --osint is invalid\n");
     }
   }
 
   return ar;
 }
 
+/**
+ * Check for a commandline flag. Ignore data that's passed in with the flag.
+ * Flags may be in the form -arg or --arg (or /arg on win32).
+ * Will not remove flag if found.
+ *
+ * @param aArg the parameter to check. Must be lowercase.
+ */
+static ArgResult
+CheckArgExists(const char* aArg)
+{
+  char **curarg = gArgv + 1; // skip argv[0]
+  while (*curarg) {
+    char *arg = curarg[0];
+
+    if (arg[0] == '-'
+#if defined(XP_WIN)
+        || *arg == '/'
+#endif
+        ) {
+      ++arg;
+      if (*arg == '-')
+        ++arg;
+
+      char delimiter = '=';
+#if defined(XP_WIN)
+      delimiter = ':';
+#endif
+      int i;
+      for (i = 0; arg[i] && arg[i] != delimiter; i++) {}
+      char tmp = arg[i];
+      arg[i] = '\0';
+      bool found = strimatch(aArg, arg);
+      arg[i] = tmp;
+      if (found) {
+        return ARG_FOUND;
+      }
+    }
+
+    ++curarg;
+  }
+
+  return ARG_NONE;
+}
+
 #if defined(XP_WIN)
 /**
  * Check for a commandline flag from the windows shell and remove it from the
  * argv used when restarting. Flags MUST be in the form -arg.
  *
  * @param aArg the parameter to check. Must be lowercase.
  */
 static ArgResult
@@ -3175,17 +3219,17 @@ XREMain::XRE_mainInit(bool* aExitFlag)
     }
     ChaosMode::SetChaosFeature(feature);
   }
 
   if (ChaosMode::isActive(ChaosFeature::Any)) {
     printf_stderr("*** You are running in chaos test mode. See ChaosMode.h. ***\n");
   }
 
-  if (CheckArg("headless")) {
+  if (CheckArg("headless") || CheckArgExists("screenshot")) {
     PR_SetEnv("MOZ_HEADLESS=1");
   }
 
   if (gfxPlatform::IsHeadless()) {
 #if defined(XP_WIN) || defined(MOZ_WIDGET_GTK) || defined(XP_MACOSX)
     printf_stderr("*** You are running in headless mode.\n");
 #else
     Output(true, "Error: headless mode is not currently supported on this platform.\n");