[Media Router] 2nd take on fix route response callback lifetime in UI.

chrome/browser/ui/webui/media_router/media_router_dialog_controller.cc

Index: chrome/browser/ui/webui/media_router/media_router_dialog_controller.cc
diff --git a/chrome/browser/ui/webui/media_router/media_router_dialog_controller.cc b/chrome/browser/ui/webui/media_router/media_router_dialog_controller.cc
index 392ebfc891d04c8db732f31a3a90a99b26a880db..0f142c2f650555c34e90eb7f1cddcc9b2ef8425a 100644
--- a/chrome/browser/ui/webui/media_router/media_router_dialog_controller.cc
+++ b/chrome/browser/ui/webui/media_router/media_router_dialog_controller.cc
@@ -170,6 +170,85 @@ class MediaRouterDialogController::InitiatorWebContentsObserver
   MediaRouterDialogController* const dialog_controller_;
 };
 
+MediaRouterDialogCallbacks::MediaRouterDialogCallbacks(
+    const base::WeakPtr<PresentationServiceDelegateImpl>& delegate)
+    : MediaRouterDialogCallbacks(nullptr, delegate) {
+}
+
+MediaRouterDialogCallbacks::MediaRouterDialogCallbacks(
+    scoped_ptr<CreateSessionRequest> presentation_request,
+    const base::WeakPtr<PresentationServiceDelegateImpl>& delegate)
+    : presentation_request_(presentation_request.Pass()),
+      pending_presentation_response_after_dialog_destroyed_(false),
+      presentation_service_delegate_(delegate),
+      weak_factory_(this) {
+}
+
+MediaRouterDialogCallbacks::~MediaRouterDialogCallbacks() {
+  if (presentation_request_) {
+    presentation_request_->MaybeInvokeErrorCallback(content::PresentationError(
+        content::PRESENTATION_ERROR_UNKNOWN, "Unknown error."));
+  }
+}
+
+void MediaRouterDialogCallbacks::OnRouteResponseReceived(
+    bool is_presentation,
+    scoped_ptr<MediaRoute> route,
+    const std::string& error) {
+  if (is_presentation) {
+    if (presentation_request_) {
+      // Remove the presentation request after responding. Subsequent
+      // responses from the same dialog will be treated as if they came from
+      // browser-initiated requests.
+      HandleRouteResponseForPresentation(route.get(), error);
+      presentation_request_.reset();
+    } else if (route && presentation_service_delegate_) {
+      // Response was due to a browser-initiated request. Let
+      // PresentationServiceDelegateImpl perform the match against the default
+      // presentation URL.
+      presentation_service_delegate_->OnRouteCreated(*route);
+    }
+  }
+
+  if (!route_response_cb_.is_null())
+    route_response_cb_.Run(route.get(), error);
+
+  if (pending_presentation_response_after_dialog_destroyed_)
+    delete this;
+}
+
+void MediaRouterDialogCallbacks::HandleRouteResponseForPresentation(
+    const MediaRoute* route,
+    const std::string& error) {
+  if (!route) {
+    presentation_request_->MaybeInvokeErrorCallback(
+        content::PresentationError(content::PRESENTATION_ERROR_UNKNOWN, error));
+  } else {
+    presentation_request_->MaybeInvokeSuccessCallback(route->media_route_id());
+  }
+}
+
+void MediaRouterDialogCallbacks::OnDialogDestroyed(
+    bool has_pending_presentation_response) {
+  if (has_pending_presentation_response) {
+    pending_presentation_response_after_dialog_destroyed_ = true;
+  } else {
+    if (presentation_request_) {
+      presentation_request_->MaybeInvokeErrorCallback(
+          content::PresentationError(
+              content::PRESENTATION_ERROR_SESSION_REQUEST_CANCELLED,
+              "Dialog closed."));
+      presentation_request_.reset();
+    }
+    delete this;
+  }
+}
+
+base::WeakPtr<MediaRouterDialogCallbacks>
+MediaRouterDialogCallbacks::GetWeakPtr() {
+  return weak_factory_.GetWeakPtr();
+}
+
 MediaRouterDialogController::MediaRouterDialogController(
     WebContents* web_contents)
     : initiator_(web_contents),
@@ -331,15 +410,22 @@ void MediaRouterDialogController::PopulateDialog(
     return;
   }
 
+  base::WeakPtr<PresentationServiceDelegateImpl> delegate =
+      PresentationServiceDelegateImpl::GetOrCreateForWebContents(initiator)
+          ->GetWeakPtr();
+  CHECK(delegate);
+  // TODO(imcheng): Currently it is assumed that MediaRouter requests will
+  // always be eventually resolved, but this may not be true due to
+  // implementation errors, in which case DialogCallbacks objects may be leaked.
+  // We should have some timeout mechanism to make sure they self-destruct
+  // after certain conditions and timeout are met.
   if (!presentation_request_.get()) {
-    PresentationServiceDelegateImpl::CreateForWebContents(initiator);
-    PresentationServiceDelegateImpl* delegate =
-        PresentationServiceDelegateImpl::FromWebContents(initiator);
-    CHECK(delegate);
-    media_router_ui->InitWithDefaultMediaSource(delegate);
+    media_router_ui->InitWithDefaultMediaSource(
+        delegate,
+        new MediaRouterDialogCallbacks(presentation_request_.Pass(), delegate));
   } else {
     media_router_ui->InitWithPresentationSessionRequest(
-        initiator, presentation_request_.Pass());
+        initiator, new MediaRouterDialogCallbacks(delegate));
   }
 }