Test for WebSocketJob being deleted on the stack

SocketStreamDispatcherHost can delete the WebSocketJob while it is still
on the stack. Add tests to ensure that WebSocketJob does not attempt to
access its own members after being deleted.

Also fix two cases where WebSocketJob attempted to access its members after
being deleted.

BUG=358038
TEST=net_unittests --gtest_filter=WebSocketJobDeleteTest*

Review URL: https://codereview.chromium.org/221833002

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@261707 0039d316-1c4b-4281-b951-d872f2087c98

2 files changed, 174 insertions, 6 deletions

diff --git a/net/websockets/websocket_job.cc b/net/websockets/websocket_job.cc
index 9ee38b7..f63259f 100644
--- a/net/websockets/websocket_job.cc
+++ b/net/websockets/websocket_job.cc
@@ -304,9 +304,10 @@ void WebSocketJob::OnSentSpdyHeaders() {
   DCHECK_NE(INITIALIZED, state_);
   if (state_ != CONNECTING)
     return;
-  if (delegate_)
-    delegate_->OnSentData(socket_.get(), handshake_request_->original_length());
+  size_t original_length = handshake_request_->original_length();
   handshake_request_.reset();
+  if (delegate_)
+    delegate_->OnSentData(socket_.get(), original_length);
 }
 
 void WebSocketJob::OnSpdyResponseHeadersUpdated(
@@ -423,11 +424,12 @@ void WebSocketJob::OnSentHandshakeRequest(
   if (handshake_request_sent_ >= handshake_request_->raw_length()) {
     // handshake request has been sent.
     // notify original size of handshake request to delegate.
-    if (delegate_)
-      delegate_->OnSentData(
-          socket,
-          handshake_request_->original_length());
+    // Reset the handshake_request_ first in case this object is deleted by the
+    // delegate.
+    size_t original_length = handshake_request_->original_length();
     handshake_request_.reset();
+    if (delegate_)
+      delegate_->OnSentData(socket, original_length);
   }
 }
 
diff --git a/net/websockets/websocket_job_test.cc b/net/websockets/websocket_job_test.cc
index 1363ff7..f519b86 100644
--- a/net/websockets/websocket_job_test.cc
+++ b/net/websockets/websocket_job_test.cc
@@ -316,6 +316,74 @@ class MockHttpTransactionFactory : public HttpTransactionFactory {
   SpdySessionKey spdy_session_key_;
 };
 
+class DeletingSocketStreamDelegate : public SocketStream::Delegate {
+ public:
+  DeletingSocketStreamDelegate()
+      : delete_next_(false) {}
+
+  // Since this class needs to be able to delete |job_|, it must be the only
+  // reference holder (except for temporary references). Provide access to the
+  // pointer for tests to use.
+  WebSocketJob* job() { return job_.get(); }
+
+  void set_job(WebSocketJob* job) { job_ = job; }
+
+  // After calling this, the next call to a method on this delegate will delete
+  // the WebSocketJob object.
+  void set_delete_next(bool delete_next) { delete_next_ = delete_next; }
+
+  void DeleteJobMaybe() {
+    if (delete_next_) {
+      job_->DetachContext();
+      job_->DetachDelegate();
+      job_ = NULL;
+    }
+  }
+
+  // SocketStream::Delegate implementation
+
+  // OnStartOpenConnection() is not implemented by SocketStreamDispatcherHost
+
+  virtual void OnConnected(SocketStream* socket,
+                           int max_pending_send_allowed) OVERRIDE {
+    DeleteJobMaybe();
+  }
+
+  virtual void OnSentData(SocketStream* socket, int amount_sent) OVERRIDE {
+    DeleteJobMaybe();
+  }
+
+  virtual void OnReceivedData(SocketStream* socket,
+                              const char* data,
+                              int len) OVERRIDE {
+    DeleteJobMaybe();
+  }
+
+  virtual void OnClose(SocketStream* socket) OVERRIDE { DeleteJobMaybe(); }
+
+  virtual void OnAuthRequired(SocketStream* socket,
+                              AuthChallengeInfo* auth_info) OVERRIDE {
+    DeleteJobMaybe();
+  }
+
+  virtual void OnSSLCertificateError(SocketStream* socket,
+                                     const SSLInfo& ssl_info,
+                                     bool fatal) OVERRIDE {
+    DeleteJobMaybe();
+  }
+
+  virtual void OnError(const SocketStream* socket, int error) OVERRIDE {
+    DeleteJobMaybe();
+  }
+
+  // CanGetCookies() and CanSetCookies() do not appear to be able to delete the
+  // WebSocketJob object.
+
+ private:
+  scoped_refptr<WebSocketJob> job_;
+  bool delete_next_;
+};
+
 }  // namespace
 
 class WebSocketJobTest : public PlatformTest,
@@ -480,6 +548,34 @@ class WebSocketJobTest : public PlatformTest,
   static const size_t kDataWorldLength;
 };
 
+// Tests using this fixture verify that the WebSocketJob can handle being
+// deleted while calling back to the delegate correctly. These tests need to be
+// run under AddressSanitizer or other systems for detecting use-after-free
+// errors in order to find problems.
+class WebSocketJobDeleteTest : public ::testing::Test {
+ protected:
+  WebSocketJobDeleteTest()
+      : delegate_(new DeletingSocketStreamDelegate),
+        cookie_store_(new MockCookieStore),
+        context_(new MockURLRequestContext(cookie_store_.get())) {
+    WebSocketJob* websocket = new WebSocketJob(delegate_.get());
+    delegate_->set_job(websocket);
+
+    socket_ = new MockSocketStream(
+        GURL("ws://127.0.0.1/"), websocket, context_.get(), NULL);
+
+    websocket->InitSocketStream(socket_.get());
+  }
+
+  void SetDeleteNext() { return delegate_->set_delete_next(true); }
+  WebSocketJob* job() { return delegate_->job(); }
+
+  scoped_ptr<DeletingSocketStreamDelegate> delegate_;
+  scoped_refptr<MockCookieStore> cookie_store_;
+  scoped_ptr<MockURLRequestContext> context_;
+  scoped_refptr<SocketStream> socket_;
+};
+
 const char WebSocketJobTest::kHandshakeRequestWithoutCookie[] =
     "GET /demo HTTP/1.1\r\n"
     "Host: example.com\r\n"
@@ -1122,6 +1218,76 @@ TEST_P(WebSocketJobTest, ThrottlingSpdySpdyEnabled) {
   TestConnectBySpdy(SPDY_ON, THROTTLING_ON);
 }
 
+TEST_F(WebSocketJobDeleteTest, OnClose) {
+  SetDeleteNext();
+  job()->OnClose(socket_.get());
+  // OnClose() sets WebSocketJob::_socket to NULL before we can detach it, so
+  // socket_->delegate is still set at this point. Clear it to avoid hitting
+  // DCHECK(!delegate_) in the SocketStream destructor. SocketStream::Finish()
+  // is the only caller of this method in real code, and it also sets delegate_
+  // to NULL.
+  socket_->DetachDelegate();
+  EXPECT_FALSE(job());
+}
+
+TEST_F(WebSocketJobDeleteTest, OnAuthRequired) {
+  SetDeleteNext();
+  job()->OnAuthRequired(socket_.get(), NULL);
+  EXPECT_FALSE(job());
+}
+
+TEST_F(WebSocketJobDeleteTest, OnSSLCertificateError) {
+  SSLInfo ssl_info;
+  SetDeleteNext();
+  job()->OnSSLCertificateError(socket_.get(), ssl_info, true);
+  EXPECT_FALSE(job());
+}
+
+TEST_F(WebSocketJobDeleteTest, OnError) {
+  SetDeleteNext();
+  job()->OnError(socket_.get(), ERR_CONNECTION_RESET);
+  EXPECT_FALSE(job());
+}
+
+TEST_F(WebSocketJobDeleteTest, OnSentSpdyHeaders) {
+  job()->Connect();
+  SetDeleteNext();
+  job()->OnSentSpdyHeaders();
+  EXPECT_FALSE(job());
+}
+
+TEST_F(WebSocketJobDeleteTest, OnSentHandshakeRequest) {
+  static const char kMinimalRequest[] =
+      "GET /demo HTTP/1.1\r\n"
+      "Host: example.com\r\n"
+      "Upgrade: WebSocket\r\n"
+      "Connection: Upgrade\r\n"
+      "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n"
+      "Origin: http://example.com\r\n"
+      "Sec-WebSocket-Version: 13\r\n"
+      "\r\n";
+  const size_t kMinimalRequestSize = arraysize(kMinimalRequest) - 1;
+  job()->Connect();
+  job()->SendData(kMinimalRequest, kMinimalRequestSize);
+  SetDeleteNext();
+  job()->OnSentData(socket_.get(), kMinimalRequestSize);
+  EXPECT_FALSE(job());
+}
+
+TEST_F(WebSocketJobDeleteTest, NotifyHeadersComplete) {
+  static const char kMinimalResponse[] =
+      "HTTP/1.1 101 Switching Protocols\r\n"
+      "Upgrade: websocket\r\n"
+      "Connection: Upgrade\r\n"
+      "Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n"
+      "\r\n";
+  job()->Connect();
+  SetDeleteNext();
+  job()->OnReceivedData(
+      socket_.get(), kMinimalResponse, arraysize(kMinimalResponse) - 1);
+  EXPECT_FALSE(job());
+}
+
 // TODO(toyoshim): Add tests to verify throttling, SPDY stream limitation.
 // TODO(toyoshim,yutak): Add tests to verify closing handshake.
 }  // namespace net