Bug 1378010 - screenshot from command line with headless; r=mossop draft
authorBenjamin Dahse <bdahse@mozilla.com>
Tue, 12 Sep 2017 16:55:27 -0700
changeset 664076 d26ed0856af882ad52a606941685b2c91577d3bf
parent 661689 ea7b55d65d76214f97aaae502d65cb26fc6f5659
child 731370 e10f8b31abd4bc8fe4886da8b1efa1422f5c352f
push id79619
push userbmo:ronoueb@gmail.com
push dateWed, 13 Sep 2017 18:06:45 +0000
reviewersmossop
bugs1378010
milestone57.0a1
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
@@ -3165,17 +3209,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");