Add unittests for BufferedSocketWriter and fix some bugs in that code.

remoting/protocol/buffered_socket_writer.cc

 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "remoting/protocol/buffered_socket_writer.h"
 
 #include "base/bind.h"
 #include "base/location.h"
+#include "base/single_thread_task_runner.h"
 #include "base/stl_util.h"
+#include "base/thread_task_runner_handle.h"
 #include "net/base/net_errors.h"
 
 namespace remoting {
 namespace protocol {
 
 class BufferedSocketWriterBase::PendingPacket {

[simonmorris, 2012/07/31 20:19:50]

Maybe just make this a struct?

[Sergey Ulanov, 2012/07/31 21:37:53]

yes, that was my intent

  public:
   PendingPacket(scoped_refptr<net::IOBufferWithSize> data,
                 const base::Closure& done_task)
-      : data_(data),
-        done_task_(done_task) {
-  }
-  ~PendingPacket() {
-    if (!done_task_.is_null())
-      done_task_.Run();
+      : data(data),
+        done_task(done_task) {
   }
 
-  net::IOBufferWithSize* data() {
-    return data_;
-  }
-
- private:
-  scoped_refptr<net::IOBufferWithSize> data_;
-  base::Closure done_task_;
-
-  DISALLOW_COPY_AND_ASSIGN(PendingPacket);
+  scoped_refptr<net::IOBufferWithSize> data;
+  base::Closure done_task;
 };
 
 BufferedSocketWriterBase::BufferedSocketWriterBase()
     : buffer_size_(0),
       socket_(NULL),
       write_pending_(false),
-      closed_(false) {
+      closed_(false),
+      destroyed_flag_(NULL) {
 }
 
 void BufferedSocketWriterBase::Init(net::Socket* socket,
                                     const WriteFailedCallback& callback) {
   DCHECK(CalledOnValidThread());
   DCHECK(socket);
   socket_ = socket;
   write_failed_callback_ = callback;
 }
 
 bool BufferedSocketWriterBase::Write(
     scoped_refptr<net::IOBufferWithSize> data, const base::Closure& done_task) {
   DCHECK(CalledOnValidThread());
   DCHECK(socket_);
+  DCHECK(data.get());
 
   // Don't write after Close().
   if (closed_)
     return false;
 
   queue_.push_back(new PendingPacket(data, done_task));
   buffer_size_ += data->size();
 
   DoWrite();
   return true;
@@ skipping 17 matching lines @@
     GetNextPacket(&current_packet, &current_packet_size);
 
     // Return if the queue is empty.
     if (!current_packet)
       return;
 
     int result = socket_->Write(
         current_packet, current_packet_size,
         base::Bind(&BufferedSocketWriterBase::OnWritten,
                    base::Unretained(this)));
-    if (result >= 0) {
-      AdvanceBufferPosition(result);
+    bool write_again = false;
+    HandleWriteResult(result, &write_again);
+    if (!write_again)
+      return;
+  }
+}
+
+void BufferedSocketWriterBase::HandleWriteResult(int result,
+                                                 bool* write_again) {

[simonmorris, 2012/07/31 20:19:50]

Can't HandleWriteResult just return write_again?

[Sergey Ulanov, 2012/07/31 21:37:53]

It can by, it's more readable this way because the result is named. Otherwise I
would have to add comment in the header to explain the meaning of that result
and that comment would have to be read every time somebody reads this code.

[simonmorris, 2012/07/31 22:20:13]

I think this way makes the call sites unnecessarily hard to read, but I don't
mind if you still prefer it.

+  *write_again = false;
+  if (result < 0) {
+    if (result == net::ERR_IO_PENDING) {
+      write_pending_ = true;
     } else {
-      if (result == net::ERR_IO_PENDING) {
-        write_pending_ = true;
-      } else {
-        HandleError(result);
-        if (!write_failed_callback_.is_null())
-          write_failed_callback_.Run(result);
-      }
+      HandleError(result);
+      if (!write_failed_callback_.is_null())
+        write_failed_callback_.Run(result);
+    }
+    return;
+  }
+
+  base::Closure done_task = AdvanceBufferPosition(result);
+  if (!done_task.is_null()) {
+    bool destroyed = false;
+    destroyed_flag_ = &destroyed;
+    done_task.Run();
+    destroyed_flag_ = NULL;

[simonmorris, 2012/07/31 20:19:50]

Isn't this use-after-free if the BufferedSocketWriterBase has been destroyed?

[Sergey Ulanov, 2012/07/31 21:37:53]

Good catch. Fixed.

+
+    if (destroyed) {
+      // Stop doing anything if we've been destroyed by the callback.
       return;
     }
   }
+
+  *write_again = true;
 }
 
 void BufferedSocketWriterBase::OnWritten(int result) {
   DCHECK(CalledOnValidThread());
+  DCHECK(write_pending_);
   write_pending_ = false;
 
-  if (result < 0) {
-    HandleError(result);
-    if (!write_failed_callback_.is_null())
-      write_failed_callback_.Run(result);
-    return;
-  }
-
-  AdvanceBufferPosition(result);
-
-  DoWrite();
+  bool write_again;
+  HandleWriteResult(result, &write_again);
+  if (write_again)
+    DoWrite();
 }
 
 void BufferedSocketWriterBase::HandleError(int result) {
   DCHECK(CalledOnValidThread());
 
   closed_ = true;
 
   STLDeleteElements(&queue_);
 
   // Notify subclass that an error is received.
   OnError(result);
 }
 
 int BufferedSocketWriterBase::GetBufferSize() {
   return buffer_size_;
 }
 
 int BufferedSocketWriterBase::GetBufferChunks() {
   return queue_.size();
 }
 
 void BufferedSocketWriterBase::Close() {
   DCHECK(CalledOnValidThread());
   closed_ = true;
 }
 
-BufferedSocketWriterBase::~BufferedSocketWriterBase() {}
+BufferedSocketWriterBase::~BufferedSocketWriterBase() {
+  if (destroyed_flag_)
+    *destroyed_flag_ = true;
 
-void BufferedSocketWriterBase::PopQueue() {
-  // This also calls |done_task|.
+  STLDeleteElements(&queue_);
+}
+
+base::Closure BufferedSocketWriterBase::PopQueue() {
+  base::Closure result = queue_.front()->done_task;
   delete queue_.front();
   queue_.pop_front();
+  return result;
 }
 
 BufferedSocketWriter::BufferedSocketWriter() {
 }
 
 void BufferedSocketWriter::GetNextPacket(
     net::IOBuffer** buffer, int* size) {
   if (!current_buf_) {
     if (queue_.empty()) {
       *buffer = NULL;
       return;  // Nothing to write.
     }
     current_buf_ = new net::DrainableIOBuffer(
-        queue_.front()->data(), queue_.front()->data()->size());
+        queue_.front()->data, queue_.front()->data->size());
   }
 
   *buffer = current_buf_;
   *size = current_buf_->BytesRemaining();
 }
 
-void BufferedSocketWriter::AdvanceBufferPosition(int written) {
+base::Closure BufferedSocketWriter::AdvanceBufferPosition(int written) {
   buffer_size_ -= written;
   current_buf_->DidConsume(written);
 
   if (current_buf_->BytesRemaining() == 0) {
-    PopQueue();
     current_buf_ = NULL;
+    return PopQueue();
   }
+  return base::Closure();
 }
 
 void BufferedSocketWriter::OnError(int result) {
   current_buf_ = NULL;
 }
 
 BufferedSocketWriter::~BufferedSocketWriter() {
-  STLDeleteElements(&queue_);
 }
 
 BufferedDatagramWriter::BufferedDatagramWriter() {
 }
 
 void BufferedDatagramWriter::GetNextPacket(
     net::IOBuffer** buffer, int* size) {
   if (queue_.empty()) {
     *buffer = NULL;
     return;  // Nothing to write.
   }
-  *buffer = queue_.front()->data();
-  *size = queue_.front()->data()->size();
+  *buffer = queue_.front()->data;
+  *size = queue_.front()->data->size();
 }
 
-void BufferedDatagramWriter::AdvanceBufferPosition(int written) {
-  DCHECK_EQ(written, queue_.front()->data()->size());
-  buffer_size_ -= queue_.front()->data()->size();
-  PopQueue();
+base::Closure BufferedDatagramWriter::AdvanceBufferPosition(int written) {
+  DCHECK_EQ(written, queue_.front()->data->size());
+  buffer_size_ -= queue_.front()->data->size();
+  return PopQueue();
 }
 
 void BufferedDatagramWriter::OnError(int result) {
   // Nothing to do here.
 }
 
-BufferedDatagramWriter::~BufferedDatagramWriter() {}
+BufferedDatagramWriter::~BufferedDatagramWriter() {
+}
 
 }  // namespace protocol
 }  // namespace remoting