1 files changed, 10 insertions, 16 deletions
diff --git a/cast/standalone_receiver/cast_socket_message_port.cc b/cast/standalone_receiver/cast_socket_message_port.cc
index d5540f28..6f3c55c8 100644
--- a/cast/standalone_receiver/cast_socket_message_port.cc
+++ b/cast/standalone_receiver/cast_socket_message_port.cc
@@ -19,27 +19,17 @@ CastSocketMessagePort::~CastSocketMessagePort() = default;
 // since sockets should map one to one with receiver sessions, we reset our
 // client. The consumer of this message port should call SetClient with the new
 // message port client after setting the socket.
-void CastSocketMessagePort::SetSocket(std::unique_ptr<CastSocket> socket) {
+void CastSocketMessagePort::SetSocket(WeakPtr<CastSocket> socket) {
   client_ = nullptr;
-  socket_ = std::move(socket);
+  socket_ = socket;
 }
 
-void CastSocketMessagePort::SetClient(MessagePort::Client* client) {
-  client_ = client;
-}
-
-void CastSocketMessagePort::OnError(CastSocket* socket, Error error) {
-  if (client_) {
-    client_->OnError(error);
-  }
+int CastSocketMessagePort::GetSocketId() {
+  return socket_ ? socket_->socket_id() : -1;
 }
 
-void CastSocketMessagePort::OnMessage(CastSocket* socket,
-                                      ::cast::channel::CastMessage message) {
-  if (client_) {
-    client_->OnMessage(message.source_id(), message.namespace_(),
-                       message.payload_utf8());
-  }
+void CastSocketMessagePort::SetClient(MessagePort::Client* client) {
+  client_ = client;
 }
 
 void CastSocketMessagePort::PostMessage(absl::string_view sender_id,
@@ -51,6 +41,10 @@ void CastSocketMessagePort::PostMessage(absl::string_view sender_id,
                               message_namespace.size());
   cast_message.set_payload_utf8(message.data(), message.size());
 
+  if (!socket_) {
+    client_->OnError(Error::Code::kAlreadyClosed);
+    return;
+  }
   Error error = socket_->Send(cast_message);
   if (!error.ok()) {
     client_->OnError(error);