Fix two race conditions in ImageDecoder.

chrome/browser/image_decoder.cc

Index: chrome/browser/image_decoder.cc
diff --git a/chrome/browser/image_decoder.cc b/chrome/browser/image_decoder.cc
index 9e21662088f5ab09a5cd61b71057da235bb6fe24..be81d5c0704d752870d5841219d0df530bf0bf33 100644
--- a/chrome/browser/image_decoder.cc
+++ b/chrome/browser/image_decoder.cc
@@ -4,6 +4,8 @@
 
 #include "chrome/browser/image_decoder.h"
 
+#include <set>
+
 #include "base/bind.h"
 #include "chrome/browser/browser_process.h"
 #include "chrome/common/chrome_utility_messages.h"
@@ -24,6 +26,84 @@ const int kBatchModeTimeoutSeconds = 5;
 
 }  // namespace
 
+class ImageDecoder::Job : public base::RefCountedThreadSafe<ImageDecoder::Job> {
+ public:
+  explicit Job(ImageDecoder::ImageRequest* request);
+
+  void Cancel();
+  void NotifyDecodeSucceeded(const SkBitmap& decoded_image);
+  void NotifyDecodeFailed();
+
+  ImageDecoder::ImageRequest* image_request() const { return request_; }
+
+ private:
+  friend class base::RefCountedThreadSafe<ImageDecoder::Job>;
+  ~Job() = default;
+
+  // Runs on the requests's task runner.

[Theresa, 2015/04/13 18:19:30]

nit: sed/requests's/request's

[Anand Mistry (off Chromium), 2015/04/14 00:52:13]

Done.

+  void OnSuccess(const SkBitmap& decoded_image);
+  void OnFailed();
+
+  void RemoveTask();
+
+  // Must only be dereferenced while |is_cancelled_| == false.
+  ImageDecoder::ImageRequest* const request_;
+
+  base::Lock lock_;
+  bool is_cancelled_;
+};
+
+ImageDecoder::Job::Job(ImageDecoder::ImageRequest* request)
+    : request_(request), is_cancelled_(false) {
+}
+
+void ImageDecoder::Job::Cancel() {
+  {
+    base::AutoLock lock(lock_);
+    is_cancelled_ = true;
+  }
+  g_decoder.Pointer()->RemoveJob(this);
+}
+
+void ImageDecoder::Job::NotifyDecodeSucceeded(const SkBitmap& decoded_image) {
+  base::AutoLock lock(lock_);
+  if (is_cancelled_)
+    return;
+
+  request_->task_runner()->PostTask(
+      FROM_HERE,
+      base::Bind(&ImageDecoder::Job::OnSuccess, this, decoded_image));

[Theresa, 2015/04/13 18:19:30]

Why are NotifyDecodeSucceeded and OnSuccess both needed? Couldn't this method do
something like this?

g_decoder.Pointer()->RemoveJob(this);
request_->task_runner()->PostTask(	
    FROM_HERE, base::Bind(&ImageRequest::OnImageDecoded,	
    base::Unretained(request_), decoded_image));
	

[Anand Mistry (off Chromium), 2015/04/14 00:52:13]

This is race condition #2. If you post ImageRequest::OnImageDecoded, it's
possible for the ImageRequest to be deleted before the task is run, leading to a
use-after-free. Job::OnSuccess needs to run on the task runner to ensure the
ImageRequest exists when ImageRequest::OnImageDecoded run.

[Theresa, 2015/04/14 01:20:47]

Acknowledged.

+}
+
+void ImageDecoder::Job::NotifyDecodeFailed() {
+  base::AutoLock lock(lock_);
+  if (is_cancelled_)
+    return;
+
+  request_->task_runner()->PostTask(
+      FROM_HERE, base::Bind(&ImageDecoder::Job::OnFailed, this));
+}
+
+void ImageDecoder::Job::OnSuccess(const SkBitmap& decoded_image) {
+  base::AutoLock lock(lock_);
+  if (is_cancelled_)
+    return;
+
+  DCHECK(request_->task_runner()->RunsTasksOnCurrentThread());
+  g_decoder.Pointer()->RemoveJob(this);
+  request_->OnImageDecoded(decoded_image);
+}
+
+void ImageDecoder::Job::OnFailed() {
+  base::AutoLock lock(lock_);
+  if (is_cancelled_)
+    return;
+
+  DCHECK(request_->task_runner()->RunsTasksOnCurrentThread());
+  g_decoder.Pointer()->RemoveJob(this);
+  request_->OnDecodeImageFailed();
+}
+
 ImageDecoder::ImageDecoder()
     : image_request_id_counter_(0), last_request_(base::TimeTicks::Now()) {
   // A single ImageDecoder instance should live for the life of the program.
@@ -54,13 +134,28 @@ void ImageDecoder::StartWithOptions(ImageRequest* image_request,
                                     const std::string& image_data,
                                     ImageCodec image_codec,
                                     bool shrink_to_fit) {
+  g_decoder.Pointer()->StartWithOptionsImpl(image_request, image_data,
+                                            image_codec, shrink_to_fit);
+}
+
+void ImageDecoder::StartWithOptionsImpl(ImageRequest* image_request,
+                                        const std::string& image_data,
+                                        ImageCodec image_codec,
+                                        bool shrink_to_fit) {
   DCHECK(image_request);
   DCHECK(image_request->task_runner());
+
+  scoped_refptr<Job> job(new Job(image_request));
+  int request_id;
+  {
+    base::AutoLock lock(map_lock_);
+    request_id = image_request_id_counter_++;
+    image_request_id_map_.insert(std::make_pair(request_id, job));
+  }
   BrowserThread::PostTask(
       BrowserThread::IO, FROM_HERE,
       base::Bind(
-          &ImageDecoder::DecodeImageInSandbox,
-          g_decoder.Pointer(), image_request,
+          &ImageDecoder::DecodeImageInSandbox, this, request_id,
           std::vector<unsigned char>(image_data.begin(), image_data.end()),
           image_codec, shrink_to_fit));
 }
@@ -72,7 +167,7 @@ void ImageDecoder::Cancel(ImageRequest* image_request) {
 }
 
 void ImageDecoder::DecodeImageInSandbox(
-    ImageRequest* image_request,
+    int request_id,
     const std::vector<unsigned char>& image_data,
     ImageCodec image_codec,
     bool shrink_to_fit) {
@@ -85,41 +180,42 @@ void ImageDecoder::DecodeImageInSandbox(
     // Without this check, we were seeing crashes on startup. Further
     // investigation is needed to determine why the utility process
     // is failing to start. See crbug.com/472272
-    image_request->task_runner()->PostTask(
-        FROM_HERE, base::Bind(&ImageRequest::OnDecodeImageFailed,
-                              base::Unretained(image_request)));
+    OnDecodeImageFailed(request_id);
     return;
   }
 
   last_request_ = base::TimeTicks::Now();
-  base::AutoLock lock(map_lock_);
-  image_request_id_map_.insert(
-      std::make_pair(image_request_id_counter_, image_request));
 
   switch (image_codec) {
     case ROBUST_JPEG_CODEC:
-      utility_process_host_->Send(new ChromeUtilityMsg_RobustJPEGDecodeImage(
-          image_data, image_request_id_counter_));
+      utility_process_host_->Send(
+          new ChromeUtilityMsg_RobustJPEGDecodeImage(image_data, request_id));
       break;
     case DEFAULT_CODEC:
       utility_process_host_->Send(new ChromeUtilityMsg_DecodeImage(
-          image_data, shrink_to_fit, image_request_id_counter_));
+          image_data, shrink_to_fit, request_id));
       break;
   }
-
-  ++image_request_id_counter_;
 }
 
 void ImageDecoder::CancelImpl(ImageRequest* image_request) {
-  base::AutoLock lock(map_lock_);
-  for (auto it = image_request_id_map_.begin();
-       it != image_request_id_map_.end();) {
-    if (it->second == image_request) {
-      image_request_id_map_.erase(it++);
-    } else {
-      ++it;
+  std::set<scoped_refptr<Job>> jobs;
+
+  {
+    base::AutoLock lock(map_lock_);
+    for (const auto& request : image_request_id_map_) {
+      if (request.second->image_request() == image_request) {
+        jobs.insert(request.second);
+        // Don't erase the job from the map here, but let Job::Cancel do it
+        // instead. This is mainly to maintain consistency with how decode
+        // success/fail are handled.
+      }
     }
   }
+
+  // Will block if |OnImageDecoded| or |OnDecodeImageFailed| is running.
+  for (auto& job : jobs)
+    job->Cancel();
 }
 
 void ImageDecoder::StartBatchMode() {
@@ -167,28 +263,47 @@ void ImageDecoder::OnDecodeImageSucceeded(
     const SkBitmap& decoded_image,
     int request_id) {
   DCHECK_CURRENTLY_ON(BrowserThread::IO);
-  base::AutoLock lock(map_lock_);
-  auto it = image_request_id_map_.find(request_id);
-  if (it != image_request_id_map_.end()) {
-    ImageRequest* image_request = it->second;
-    image_request->task_runner()->PostTask(
-        FROM_HERE, base::Bind(&ImageRequest::OnImageDecoded,
-                              base::Unretained(image_request), decoded_image));
-
-    image_request_id_map_.erase(it);
+
+  scoped_refptr<Job> job;
+  {
+    base::AutoLock lock(map_lock_);
+    auto it = image_request_id_map_.find(request_id);
+    if (it != image_request_id_map_.end()) {
+      job = it->second;
+      // NOTE: The job isn't removed from the map because it can still be
+      // cancelled.
+    }
   }
+  // Call Job's functions outside |map_lock_| to avoid a potential deadlock with
+  // |Job.lock_|.
+  if (job)
+    job->NotifyDecodeSucceeded(decoded_image);
 }
 
 void ImageDecoder::OnDecodeImageFailed(int request_id) {
   DCHECK_CURRENTLY_ON(BrowserThread::IO);
+
+  scoped_refptr<Job> job;
+  {
+    base::AutoLock lock(map_lock_);
+    auto it = image_request_id_map_.find(request_id);
+    if (it != image_request_id_map_.end()) {
+      job = it->second;
+    }
+  }
+  if (job)
+    job->NotifyDecodeFailed();
+}
+
+void ImageDecoder::RemoveJob(const scoped_refptr<Job>& job) {
   base::AutoLock lock(map_lock_);
-  auto it = image_request_id_map_.find(request_id);
-  if (it != image_request_id_map_.end()) {
-    ImageRequest* image_request = it->second;
-    image_request->task_runner()->PostTask(
-        FROM_HERE, base::Bind(&ImageRequest::OnDecodeImageFailed,
-                              base::Unretained(image_request)));
-
-    image_request_id_map_.erase(it);
+  for (auto it = image_request_id_map_.begin();
+       it != image_request_id_map_.end();) {
+    if (it->second == job) {
+      image_request_id_map_.erase(it++);

[Theresa, 2015/04/13 18:19:30]

nit: does it++ need to be incremented here since there's a break right after?

[Anand Mistry (off Chromium), 2015/04/14 00:52:13]

Done.

+      break;
+    } else {
+      ++it;
+    }
   }
 }