3 files changed, 27 insertions, 0 deletions
diff --git a/WebCore/websockets/WebSocket.cpp b/WebCore/websockets/WebSocket.cpp
index 0435490..f50dc5c 100644
--- a/WebCore/websockets/WebSocket.cpp
+++ b/WebCore/websockets/WebSocket.cpp
@@ -151,6 +151,7 @@ void WebSocket::connect(const KURL& url, const String& protocol, ExceptionCode&
 
     m_channel = ThreadableWebSocketChannel::create(scriptExecutionContext(), this, m_url, m_protocol);
     m_channel->connect();
+    ActiveDOMObject::setPendingActivity(this);
 }
 
 bool WebSocket::send(const String& message, ExceptionCode& ec)
@@ -198,6 +199,26 @@ ScriptExecutionContext* WebSocket::scriptExecutionContext() const
     return ActiveDOMObject::scriptExecutionContext();
 }
 
+void WebSocket::contextDestroyed()
+{
+    LOG(Network, "WebSocket %p scriptExecutionContext destroyed", this);
+    ASSERT(!m_channel);
+    ASSERT(m_state == CLOSED);
+    ActiveDOMObject::contextDestroyed();
+}
+
+void WebSocket::stop()
+{
+    bool pending = hasPendingActivity();
+    if (m_channel)
+        m_channel->disconnect();
+    m_channel = 0;
+    m_state = CLOSED;
+    ActiveDOMObject::stop();
+    if (pending)
+        ActiveDOMObject::unsetPendingActivity(this);
+}
+
 void WebSocket::didConnect()
 {
     LOG(Network, "WebSocket %p didConnect", this);
@@ -224,6 +245,9 @@ void WebSocket::didClose()
     LOG(Network, "WebSocket %p didClose", this);
     m_state = CLOSED;
     dispatchEvent(Event::create(eventNames().closeEvent, false, false));
+    m_channel = 0;
+    if (hasPendingActivity())
+        ActiveDOMObject::unsetPendingActivity(this);
 }
 
 EventTargetData* WebSocket::eventTargetData()
diff --git a/WebCore/websockets/WebSocket.h b/WebCore/websockets/WebSocket.h
index 18e2b25..c72dbbd 100644
--- a/WebCore/websockets/WebSocket.h
+++ b/WebCore/websockets/WebSocket.h
@@ -82,6 +82,8 @@ namespace WebCore {
         virtual WebSocket* toWebSocket() { return this; }
 
         virtual ScriptExecutionContext* scriptExecutionContext() const;
+        virtual void contextDestroyed();
+        virtual void stop();
 
         using RefCounted<WebSocket>::ref;
         using RefCounted<WebSocket>::deref;
diff --git a/WebCore/websockets/WebSocketChannel.cpp b/WebCore/websockets/WebSocketChannel.cpp
index 5c0f4c3..df66c14 100644
--- a/WebCore/websockets/WebSocketChannel.cpp
+++ b/WebCore/websockets/WebSocketChannel.cpp
@@ -141,6 +141,7 @@ void WebSocketChannel::didClose(SocketStreamHandle* handle)
 void WebSocketChannel::didReceiveData(SocketStreamHandle* handle, const char* data, int len)
 {
     LOG(Network, "WebSocketChannel %p didReceiveData %d", this, len);
+    RefPtr<WebSocketChannel> protect(this); // The client can close the channel, potentially removing the last reference.
     ASSERT(handle == m_handle);
     if (!appendToBuffer(data, len)) {
         handle->close();