[chromedriver] Fallback to devtools for screenshotting of forbidden urls.

BUG=none
R=chrisgao@chromium.org

Review URL: https://codereview.chromium.org/23542015

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@222200 0039d316-1c4b-4281-b951-d872f2087c98

chrome/test/chromedriver/chrome/automation_extension.cc

@@ -24,7 +24,7 @@ Status AutomationExtension::CaptureScreenshot(std::string* screenshot) {
       base::TimeDelta::FromSeconds(10),
       &result);
   if (status.IsError())
-    return Status(kUnknownError, "cannot take screenshot", status);
+    return Status(status.code(), "cannot take screenshot", status);
   if (!result->GetAsString(screenshot))
     return Status(kUnknownError, "screenshot is not a string");
   return Status(kOk);

chrome/test/chromedriver/chrome/automation_extension.h

@@ -24,6 +24,7 @@ class AutomationExtension {
   ~AutomationExtension();
 
   // Captures the visible part of the current tab as a base64-encoded PNG.
+  // Returns |kForbidden| for security restricted pages.
   Status CaptureScreenshot(std::string* screenshot);
 
   // Gets the position of the current window.

chrome/test/chromedriver/chrome/status.cc

@@ -55,6 +55,8 @@ const char* DefaultMessageForStatusCode(StatusCode code) {
       return "chrome not reachable";
     case kDisconnected:
       return "disconnected";
+    case kForbidden:
+      return "forbidden";
     default:
       return "<unknown>";
   }

chrome/test/chromedriver/chrome/status.h

@@ -32,6 +32,7 @@ enum StatusCode {
   kChromeNotReachable = 100,
   kNoSuchExecutionContext,
   kDisconnected,
+  kForbidden = 103,
 };
 
 // Represents a WebDriver status, which may be an error or ok.

chrome/test/chromedriver/extension/background.js

@@ -27,6 +27,13 @@ function checkForExtensionError(errCallback) {
  */
 function captureScreenshot(callback, errCallback) {
   chrome.tabs.captureVisibleTab({format:'png'}, function(dataUrl) {
+    if (chrome.extension.lastError &&
+        chrome.extension.lastError.message.indexOf('permission') != -1) {
+      var error = new Error(chrome.extension.lastError.message);
+      error.code = 103;  // kForbidden
+      errCallback(error);
+      return;
+    }
     checkForExtensionError(errCallback);
     var base64 = ';base64,';
     callback(dataUrl.substr(dataUrl.indexOf(base64) + base64.length))

chrome/test/chromedriver/window_commands.cc

@@ -743,13 +743,15 @@ Status ExecuteScreenshot(
     if (status.IsError())
       return status;
     status = extension->CaptureScreenshot(&screenshot);
-    if (status.IsError())
-      return status;
-  } else {
-    Status status = web_view->CaptureScreenshot(&screenshot);
-    if (status.IsError())
+    // If the screenshot was forbidden, fallback to DevTools.
+    if (status.code() != kForbidden)
       return status;
   }
+
+  status = web_view->CaptureScreenshot(&screenshot);
+  if (status.IsError())
+    return status;
+
   value->reset(new base::StringValue(screenshot));
   return Status(kOk);
 }