1 files changed, 47 insertions, 14 deletions

diff --git a/extensions/renderer/display_source_custom_bindings.cc b/extensions/renderer/display_source_custom_bindings.cc
index 50752c0..ab68553 100644
--- a/extensions/renderer/display_source_custom_bindings.cc
+++ b/extensions/renderer/display_source_custom_bindings.cc
@@ -71,8 +71,10 @@ v8::Local<v8::Value> GetChildValue(v8::Local<v8::Object> value,
 
 void DisplaySourceCustomBindings::StartSession(
     const v8::FunctionCallbackInfo<v8::Value>& args) {
-  CHECK_EQ(1, args.Length());
+  CHECK_EQ(2, args.Length());
   CHECK(args[0]->IsObject());
+  CHECK(args[1]->IsFunction());
+
   v8::Isolate* isolate = context()->isolate();
   v8::Local<v8::Object> start_info = args[0].As<v8::Object>();
 
@@ -164,14 +166,20 @@ void DisplaySourceCustomBindings::StartSession(
   session->SetCallbacks(on_started_callback,
                         on_terminated_callback,
                         on_error_callback);
+
+  CallbackInfo cb_info = GetCallbackInfo(kStarted, sink_id);
+  args.GetReturnValue().Set(static_cast<int32_t>(cb_info.call_id));
+  callbacks_.push_back(cb_info);
+
   session->Start();
   session_map_.insert(std::make_pair(sink_id, std::move(session)));
 }
 
 void DisplaySourceCustomBindings::TerminateSession(
     const v8::FunctionCallbackInfo<v8::Value>& args) {
-  CHECK_EQ(1, args.Length());
+  CHECK_EQ(2, args.Length());
   CHECK(args[0]->IsInt32());
+  CHECK(args[1]->IsFunction());
 
   v8::Isolate* isolate = context()->isolate();
   int sink_id = args[0]->ToInt32(args.GetIsolate())->Value();
@@ -187,17 +195,40 @@ void DisplaySourceCustomBindings::TerminateSession(
         isolate, kSessionAlreadyTerminating)));
     return;
   }
+
+  CallbackInfo cb_info = GetCallbackInfo(kTerminated, sink_id);
+  args.GetReturnValue().Set(static_cast<int32_t>(cb_info.call_id));
+  callbacks_.push_back(cb_info);
+
   // The session will get removed from session_map_ in OnSessionTerminated.
   session->Terminate();
 }
 
-void DisplaySourceCustomBindings::DispatchSessionStarted(int sink_id) const {
+void DisplaySourceCustomBindings::CallCompletionCallback(
+    int sink_id,
+    CallbackType type,
+    const std::string& error_message) {
+  auto predicate = [sink_id, type](const CallbackInfo& info) -> bool {
+    return info.sink_id == sink_id && info.type == type;
+  };
+  auto it = std::find_if(callbacks_.begin(), callbacks_.end(), predicate);
+  if (it == callbacks_.end())
+    return;
+
   v8::Isolate* isolate = context()->isolate();
+  ModuleSystem* module_system = context()->module_system();
   v8::HandleScope handle_scope(isolate);
   v8::Context::Scope context_scope(context()->v8_context());
-  v8::Local<v8::Array> event_args = v8::Array::New(isolate, 1);
-  event_args->Set(0, v8::Integer::New(isolate, sink_id));
-  context()->DispatchEvent("displaySource.onSessionStarted", event_args);
+
+  v8::Local<v8::Value> callback_args[2];
+  callback_args[0] = v8::Integer::New(isolate, it->call_id);
+  if (error_message.empty())
+    callback_args[1] = v8::Null(isolate);
+  else
+    callback_args[1] = v8::String::NewFromUtf8(isolate, error_message.c_str());
+
+  module_system->CallModuleMethod("displaySource", "callCompletionCallback", 2,
+                                  callback_args);
 }
 
 void DisplaySourceCustomBindings::DispatchSessionTerminated(int sink_id) const {
@@ -242,7 +273,7 @@ DisplaySourceSession* DisplaySourceCustomBindings::GetDisplaySession(
 }
 
 void DisplaySourceCustomBindings::OnSessionStarted(int sink_id) {
-  DispatchSessionStarted(sink_id);
+  CallCompletionCallback(sink_id, kStarted);
 }
 
 void DisplaySourceCustomBindings::OnSessionTerminated(int sink_id) {
@@ -250,6 +281,7 @@ void DisplaySourceCustomBindings::OnSessionTerminated(int sink_id) {
   CHECK(session);
   session_map_.erase(sink_id);
   DispatchSessionTerminated(sink_id);
+  CallCompletionCallback(sink_id, kTerminated);
 }
 
 void DisplaySourceCustomBindings::OnSessionError(int sink_id,
@@ -257,14 +289,15 @@ void DisplaySourceCustomBindings::OnSessionError(int sink_id,
                                                  const std::string& message) {
   DisplaySourceSession* session = GetDisplaySession(sink_id);
   CHECK(session);
-  if (session->state() != DisplaySourceSession::Established &&
-      session->state() != DisplaySourceSession::Terminating) {
-    // Error has occured before the session has actually started,
-    // no need to wait for session termination notification.
-    session_map_.erase(sink_id);
-  }
-
   DispatchSessionError(sink_id, type, message);
 }
 
+DisplaySourceCustomBindings::CallbackInfo
+DisplaySourceCustomBindings::GetCallbackInfo(
+    DisplaySourceCustomBindings::CallbackType type,
+    int sink_id) const {
+  static int sCallId = 0;
+  return {type, sink_id, ++sCallId};
+}
+
 }  // extensions