Add timeout to ImageFetchTabHelper::GetImageData

Add a timeout parameter for ImageFetchTabHelper::GetImageData in case the
JavaScript does not response.

Bug: 163201
Cq-Include-Trybots: luci.chromium.try:ios-simulator-full-configs;master.tryserver.chromium.mac:ios-simulator-cronet
Change-Id: Ibee9284d937f020ae8deed55eab9eb3aabdba5c1
Reviewed-on: https://chromium-review.googlesource.com/1169168
Commit-Queue: Yi Su <mrsuyi@chromium.org>
Reviewed-by: Eugene But <eugenebut@chromium.org>
Cr-Commit-Position: refs/heads/master@{#582092}

ios/chrome/browser/web/image_fetch_tab_helper.h

@@ -8,6 +8,7 @@
 #include <unordered_map>
 
 #include "base/macros.h"
+#include "base/memory/weak_ptr.h"
 #include "ios/web/public/web_state/web_state_observer.h"
 #import "ios/web/public/web_state/web_state_user_data.h"
 
@@ -29,8 +30,11 @@ class ImageFetchTabHelper : public web::WebStateObserver,
   //   cache.
   // This method should only be called from UI thread, and |url| should be equal
   // to the resolved "src" attribute of <img>, otherwise the method 1 would
-  // fail. |callback| will be called on UI thread.
-  void GetImageData(const GURL& url, ImageDataCallback&& callback);
+  // fail. |callback| will be called on UI thread. If the JavaScript does not
+  // response after |timeout|, the |callback| will be invoked with nullptr.
+  void GetImageData(const GURL& url,
+                    base::TimeDelta timeout,
+                    ImageDataCallback&& callback);
 
  private:
   friend class web::WebStateUserData<ImageFetchTabHelper>;
@@ -43,10 +47,10 @@ class ImageFetchTabHelper : public web::WebStateObserver,
   void WebStateDestroyed(web::WebState* web_state) override;
 
   // Handler for messages sent back from injected JavaScript.
-  bool OnImageDataReceived(const base::DictionaryValue& message,
-                           const GURL& page_url,
-                           bool has_user_gesture,
-                           bool form_in_main_frame);
+  bool OnImageDataReceived(const base::DictionaryValue& message);
+
+  // Handler for timeout on GetImageData.
+  void OnImageDataTimeout(int call_id);
 
   // WebState this tab helper is attached to.
   web::WebState* web_state_ = nullptr;
@@ -60,6 +64,8 @@ class ImageFetchTabHelper : public web::WebStateObserver,
   // |OnImageDataReceived| and used to invoke the corresponding callback.
   int call_id_ = 0;
 
+  base::WeakPtrFactory<ImageFetchTabHelper> weak_ptr_factory_;
+
   DISALLOW_COPY_AND_ASSIGN(ImageFetchTabHelper);
 };
 

ios/chrome/browser/web/image_fetch_tab_helper.mm

@@ -9,6 +9,7 @@
 #include "base/strings/utf_string_conversions.h"
 #include "base/values.h"
 #import "ios/web/public/web_state/navigation_context.h"
+#include "ios/web/public/web_thread.h"
 
 #if !defined(__has_feature) || !__has_feature(objc_arc)
 #error "This file requires ARC support."
@@ -22,11 +23,18 @@ const char kCommandPrefix[] = "imageFetch";
 }
 
 ImageFetchTabHelper::ImageFetchTabHelper(web::WebState* web_state)
-    : web_state_(web_state) {
+    : web_state_(web_state), weak_ptr_factory_(this) {
   web_state->AddObserver(this);
+  // BindRepeating cannot work on WeakPtr and function with return value, use
+  // lambda as mediator.
   web_state->AddScriptCommandCallback(
-      base::BindRepeating(&ImageFetchTabHelper::OnImageDataReceived,
-                          base::Unretained(this)),
+      base::BindRepeating(
+          [](base::WeakPtr<ImageFetchTabHelper> ptr,
+             const base::DictionaryValue& message, const GURL& page_url,
+             bool has_user_gesture, bool form_in_main_frame) {
+            return ptr ? ptr->OnImageDataReceived(message) : true;
+          },
+          weak_ptr_factory_.GetWeakPtr()),
       kCommandPrefix);
 }
 
@@ -52,15 +60,22 @@ void ImageFetchTabHelper::WebStateDestroyed(web::WebState* web_state) {
 }
 
 void ImageFetchTabHelper::GetImageData(const GURL& url,
+                                       base::TimeDelta timeout,
                                        ImageDataCallback&& callback) {
   ++call_id_;
   DCHECK_EQ(callbacks_.count(call_id_), 0UL);
   callbacks_.insert({call_id_, std::move(callback)});
+
+  web::WebThread::PostDelayedTask(
+      web::WebThread::UI, FROM_HERE,
+      base::BindRepeating(&ImageFetchTabHelper::OnImageDataTimeout,
+                          weak_ptr_factory_.GetWeakPtr(), call_id_),
+      timeout);
+
   std::string js =
       base::StringPrintf("__gCrWeb.imageFetch.getImageData(%d, '%s')", call_id_,
                          url.spec().c_str());
-  // TODO(crbug.com/163201): Add timeout for callback in case JavaScript does
-  // not return.
+
   web_state_->ExecuteJavaScript(base::UTF8ToUTF16(js));
 }
 
@@ -75,17 +90,14 @@ void ImageFetchTabHelper::GetImageData(const GURL& url,
 //   {'command': 'image.getImageData',
 //    'id': id_sent_to_gCrWeb_image_getImageData}
 bool ImageFetchTabHelper::OnImageDataReceived(
-    const base::DictionaryValue& message,
-    const GURL& page_url,
-    bool has_user_gesture,
-    bool form_in_main_frame) {
+    const base::DictionaryValue& message) {
   const base::Value* id_key = message.FindKey("id");
   if (!id_key || !id_key->is_double()) {
     return false;
   }
   int id_value = static_cast<int>(id_key->GetDouble());
   if (!callbacks_.count(id_value)) {
-    return false;
+    return true;
   }
   ImageDataCallback callback = std::move(callbacks_[id_value]);
   callbacks_.erase(id_value);
@@ -100,3 +112,11 @@ bool ImageFetchTabHelper::OnImageDataReceived(
   }
   return true;
 }
+
+void ImageFetchTabHelper::OnImageDataTimeout(int call_id) {
+  if (callbacks_.count(call_id)) {
+    ImageDataCallback callback = std::move(callbacks_[call_id]);
+    callbacks_.erase(call_id);
+    std::move(callback).Run(nullptr);
+  }
+}

ios/chrome/browser/web/image_fetch_tab_helper_unittest.mm

@@ -59,6 +59,7 @@ TEST_F(ImageFetchTabHelperTest, GetImageDataJsSucceed) {
 
   image_fetch_tab_helper()->GetImageData(
       GURL("http://a.com/"),
+      base::TimeDelta::FromSeconds(kWaitForJSCompletionTimeout),
       base::BindOnce(&ImageFetchTabHelperTest::OnImageData,
                      base::Unretained(this)));
 
@@ -82,6 +83,7 @@ TEST_F(ImageFetchTabHelperTest, GetImageDataJsFail) {
 
   image_fetch_tab_helper()->GetImageData(
       GURL("http://a.com/"),
+      base::TimeDelta::FromSeconds(kWaitForJSCompletionTimeout),
       base::BindOnce(&ImageFetchTabHelperTest::OnImageData,
                      base::Unretained(this)));
 
@@ -92,11 +94,35 @@ TEST_F(ImageFetchTabHelperTest, GetImageDataJsFail) {
   EXPECT_FALSE(image_data_);
 }
 
+// Tests that ImageFetchTabHelper::GetImageData returns nullptr in callback when
+// JavaScript does not send a message back.
+TEST_F(ImageFetchTabHelperTest, GetImageDataJsTimeout) {
+  // Replaces __gCrWeb.imageFetch.getImageData with empty function to trigger a
+  // timeout.
+  id script_result = ExecuteJavaScript(
+      @"__gCrWeb.imageFetch = {}; __gCrWeb.imageFetch.getImageData = "
+      @"function(id, url){}; true;");
+  ASSERT_NSEQ(@YES, script_result);
+
+  image_fetch_tab_helper()->GetImageData(
+      GURL("http://a.com/"), base::TimeDelta::FromMilliseconds(100),
+      base::BindOnce(&ImageFetchTabHelperTest::OnImageData,
+                     base::Unretained(this)));
+
+  ASSERT_TRUE(WaitUntilConditionOrTimeout(kWaitForJSCompletionTimeout, ^{
+    base::RunLoop().RunUntilIdle();
+    return on_image_data_called_;
+  }));
+
+  EXPECT_FALSE(image_data_);
+}
+
 // Tests that ImageFetchTabHelper::GetImageData returns nullptr in callback when
 // WebState is destroyed.
 TEST_F(ImageFetchTabHelperTest, GetImageDataWebStateDestroy) {
   image_fetch_tab_helper()->GetImageData(
       GURL("http://a.com/"),
+      base::TimeDelta::FromSeconds(kWaitForJSCompletionTimeout),
       base::BindOnce(&ImageFetchTabHelperTest::OnImageData,
                      base::Unretained(this)));
 
@@ -114,6 +140,7 @@ TEST_F(ImageFetchTabHelperTest, GetImageDataWebStateDestroy) {
 TEST_F(ImageFetchTabHelperTest, GetImageDataWebStateNavigate) {
   image_fetch_tab_helper()->GetImageData(
       GURL("http://a.com/"),
+      base::TimeDelta::FromSeconds(kWaitForJSCompletionTimeout),
       base::BindOnce(&ImageFetchTabHelperTest::OnImageData,
                      base::Unretained(this)));