--- a/testing/marionette/driver.js
+++ b/testing/marionette/driver.js
@@ -2389,39 +2389,58 @@ GeckoDriver.prototype.clearImportedScrip
  *
  * If called in the chrome context, the screenshot will always represent the
  * entire viewport.
  *
  * @param {string} id
  *     Reference to a web element.
  * @param {string} highlights
  *     List of web elements to highlight.
+ * @param {boolean} full
+ *     True to take a screenshot of the entire document element. Is not
+ *     considered if {@code id} is not defined. Defaults to true.
  * @param {boolean} hash
  *     True if the user requests a hash of the image data.
  *
  * @return {string}
  *     If {@code hash} is false, PNG image encoded as base64 encoded string. If
  *     'hash' is True, hex digest of the SHA-256 hash of the base64 encoded
  *     string.
  */
 GeckoDriver.prototype.takeScreenshot = function (cmd, resp) {
   let {id, highlights, full, hash} = cmd.parameters;
   highlights = highlights || [];
 
   switch (this.context) {
     case Context.CHROME:
+      let canvas;
       let container = {frame: this.getCurrentWindow()};
       let highlightEls = [];
 
       for (let h of highlights) {
         let el = this.curBrowser.seenEls.get(h, container);
         highlightEls.push(el);
       }
 
-      let canvas = capture.viewport(this.getCurrentWindow(), highlightEls);
+      // viewport
+      if (!id && !full) {
+        canvas = capture.viewport(container.frame, highlightEls);
+
+      // element or full document element
+      } else {
+        let node;
+        if (id) {
+          node = this.curBrowser.seenEls.get(id, container);
+        } else {
+          node = container.frame.document.documentElement;
+        }
+
+        canvas = capture.element(node, highlightEls);
+      }
+
       if (hash) {
         return capture.toHash(canvas);
       } else {
         return capture.toBase64(canvas);
       }
 
     case Context.CONTENT:
       if (hash) {