More UDP mock support.

BUG=60149
TEST=net_unittests

Review URL: http://codereview.chromium.org/7276008

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

2 files changed, 123 insertions, 33 deletions

diff --git a/net/socket/socket_test_util.cc b/net/socket/socket_test_util.cc
index 474a739..c32d6bb 100644
--- a/net/socket/socket_test_util.cc
+++ b/net/socket/socket_test_util.cc
@@ -595,8 +595,7 @@ DatagramClientSocket* MockClientSocketFactory::CreateDatagramClientSocket(
     const NetLog::Source& source) {
   SocketDataProvider* data_provider = mock_data_.GetNext();
   MockUDPClientSocket* socket = new MockUDPClientSocket(data_provider);
-
-  // TODO(agayev): figure out how to enable data_provider->set_socket(socket).
+  data_provider->set_socket(socket);
   udp_client_sockets_.push_back(socket);
   return socket;
 }
@@ -1166,6 +1165,12 @@ void MockSSLClientSocket::OnReadComplete(const MockRead& data) {
 MockUDPClientSocket::MockUDPClientSocket(SocketDataProvider* data)
     : connected_(false),
       data_(data),
+      read_offset_(0),
+      read_data_(false, net::ERR_UNEXPECTED),
+      need_read_data_(true),
+      pending_buf_(NULL),
+      pending_buf_len_(0),
+      pending_callback_(NULL),
       ALLOW_THIS_IN_INITIALIZER_LIST(method_factory_(this)) {
   DCHECK(data_);
   data_->Reset();
@@ -1175,26 +1180,29 @@ MockUDPClientSocket::~MockUDPClientSocket() {}
 
 int MockUDPClientSocket::Read(net::IOBuffer* buf, int buf_len,
                               net::CompletionCallback* callback) {
-  DCHECK(buf);
-  DCHECK_GT(buf_len, 0);
-
   if (!connected_)
-    return ERR_UNEXPECTED;
+    return net::ERR_UNEXPECTED;
 
-  MockRead read_data = data_->GetNextRead();
-  DCHECK_LE(read_data.data_len, buf_len);
+  // If the buffer is already in use, a read is already in progress!
+  DCHECK(pending_buf_ == NULL);
 
-  int result = read_data.result;
-  if (result == 0) {
-    result = read_data.data_len;
-    memcpy(buf->data(), read_data.data, read_data.data_len);
-  }
+  // Store our async IO data.
+  pending_buf_ = buf;
+  pending_buf_len_ = buf_len;
+  pending_callback_ = callback;
 
-  if (read_data.async) {
-    RunCallbackAsync(callback, result);
-    return ERR_IO_PENDING;
+  if (need_read_data_) {
+    read_data_ = data_->GetNextRead();
+    // ERR_IO_PENDING means that the SocketDataProvider is taking responsibility
+    // to complete the async IO manually later (via OnReadComplete).
+    if (read_data_.result == ERR_IO_PENDING) {
+      DCHECK(callback);  // We need to be using async IO in this case.
+      return ERR_IO_PENDING;
+    }
+    need_read_data_ = false;
   }
-  return result;
+
+  return CompleteRead();
 }
 
 int MockUDPClientSocket::Write(net::IOBuffer* buf, int buf_len,
@@ -1242,6 +1250,63 @@ int MockUDPClientSocket::Connect(const IPEndPoint& address) {
   return OK;
 }
 
+void MockUDPClientSocket::OnReadComplete(const MockRead& data) {
+  // There must be a read pending.
+  DCHECK(pending_buf_);
+  // You can't complete a read with another ERR_IO_PENDING status code.
+  DCHECK_NE(ERR_IO_PENDING, data.result);
+  // Since we've been waiting for data, need_read_data_ should be true.
+  DCHECK(need_read_data_);
+
+  read_data_ = data;
+  need_read_data_ = false;
+
+  // The caller is simulating that this IO completes right now.  Don't
+  // let CompleteRead() schedule a callback.
+  read_data_.async = false;
+
+  net::CompletionCallback* callback = pending_callback_;
+  int rv = CompleteRead();
+  RunCallback(callback, rv);
+}
+
+int MockUDPClientSocket::CompleteRead() {
+  DCHECK(pending_buf_);
+  DCHECK(pending_buf_len_ > 0);
+
+  // Save the pending async IO data and reset our |pending_| state.
+  net::IOBuffer* buf = pending_buf_;
+  int buf_len = pending_buf_len_;
+  net::CompletionCallback* callback = pending_callback_;
+  pending_buf_ = NULL;
+  pending_buf_len_ = 0;
+  pending_callback_ = NULL;
+
+  int result = read_data_.result;
+  DCHECK(result != ERR_IO_PENDING);
+
+  if (read_data_.data) {
+    if (read_data_.data_len - read_offset_ > 0) {
+      result = std::min(buf_len, read_data_.data_len - read_offset_);
+      memcpy(buf->data(), read_data_.data + read_offset_, result);
+      read_offset_ += result;
+      if (read_offset_ == read_data_.data_len) {
+        need_read_data_ = true;
+        read_offset_ = 0;
+      }
+    } else {
+      result = 0;  // EOF
+    }
+  }
+
+  if (read_data_.async) {
+    DCHECK(callback);
+    RunCallbackAsync(callback, result);
+    return net::ERR_IO_PENDING;
+  }
+  return result;
+}
+
 void MockUDPClientSocket::RunCallbackAsync(net::CompletionCallback* callback,
                                            int result) {
   MessageLoop::current()->PostTask(FROM_HERE,
diff --git a/net/socket/socket_test_util.h b/net/socket/socket_test_util.h
index 52bfdcd..3de7c75 100644
--- a/net/socket/socket_test_util.h
+++ b/net/socket/socket_test_util.h
@@ -46,6 +46,7 @@ enum {
   ERR_TEST_PEER_CLOSE_AFTER_NEXT_MOCK_READ = -10000,
 };
 
+class AsyncSocket;
 class MockClientSocket;
 class SSLClientSocket;
 class SSLHostInfo;
@@ -140,26 +141,38 @@ class SocketDataProvider {
 
   // Returns the buffer and result code for the next simulated read.
   // If the |MockRead.result| is ERR_IO_PENDING, it informs the caller
-  // that it will be called via the MockClientSocket::OnReadComplete()
+  // that it will be called via the AsyncSocket::OnReadComplete()
   // function at a later time.
   virtual MockRead GetNextRead() = 0;
   virtual MockWriteResult OnWrite(const std::string& data) = 0;
   virtual void Reset() = 0;
 
   // Accessor for the socket which is using the SocketDataProvider.
-  MockClientSocket* socket() { return socket_; }
-  void set_socket(MockClientSocket* socket) { socket_ = socket; }
+  AsyncSocket* socket() { return socket_; }
+  void set_socket(AsyncSocket* socket) { socket_ = socket; }
 
   MockConnect connect_data() const { return connect_; }
   void set_connect_data(const MockConnect& connect) { connect_ = connect; }
 
  private:
   MockConnect connect_;
-  MockClientSocket* socket_;
+  AsyncSocket* socket_;
 
   DISALLOW_COPY_AND_ASSIGN(SocketDataProvider);
 };
 
+// The AsyncSocket is an interface used by the SocketDataProvider to
+// complete the asynchronous read operation.
+class AsyncSocket {
+ public:
+  // If an async IO is pending because the SocketDataProvider returned
+  // ERR_IO_PENDING, then the AsyncSocket waits until this OnReadComplete
+  // is called to complete the asynchronous read operation.
+  // data.async is ignored, and this read is completed synchronously as
+  // part of this call.
+  virtual void OnReadComplete(const MockRead& data) = 0;
+};
+
 // SocketDataProvider which responds based on static tables of mock reads and
 // writes.
 class StaticSocketDataProvider : public SocketDataProvider {
@@ -533,6 +546,9 @@ class MockClientSocketFactory : public ClientSocketFactory {
   std::vector<MockTCPClientSocket*>& tcp_client_sockets() {
     return tcp_client_sockets_;
   }
+  std::vector<MockUDPClientSocket*>& udp_client_sockets() {
+    return udp_client_sockets_;
+  }
 
   // ClientSocketFactory
   virtual DatagramClientSocket* CreateDatagramClientSocket(
@@ -567,13 +583,6 @@ class MockClientSocket : public net::SSLClientSocket {
  public:
   explicit MockClientSocket(net::NetLog* net_log);
 
-  // If an async IO is pending because the SocketDataProvider returned
-  // ERR_IO_PENDING, then the MockClientSocket waits until this OnReadComplete
-  // is called to complete the asynchronous read operation.
-  // data.async is ignored, and this read is completed synchronously as
-  // part of this call.
-  virtual void OnReadComplete(const MockRead& data) = 0;
-
   // Socket methods:
   virtual int Read(net::IOBuffer* buf, int buf_len,
                    net::CompletionCallback* callback) = 0;
@@ -612,7 +621,7 @@ class MockClientSocket : public net::SSLClientSocket {
   net::BoundNetLog net_log_;
 };
 
-class MockTCPClientSocket : public MockClientSocket {
+class MockTCPClientSocket : public MockClientSocket, public AsyncSocket {
  public:
   MockTCPClientSocket(const net::AddressList& addresses, net::NetLog* net_log,
                       net::SocketDataProvider* socket);
@@ -636,7 +645,7 @@ class MockTCPClientSocket : public MockClientSocket {
   virtual int64 NumBytesRead() const;
   virtual base::TimeDelta GetConnectTimeMicros() const;
 
-  // MockClientSocket:
+  // AsyncSocket:
   virtual void OnReadComplete(const MockRead& data);
 
  private:
@@ -663,6 +672,7 @@ class MockTCPClientSocket : public MockClientSocket {
 };
 
 class DeterministicMockTCPClientSocket : public MockClientSocket,
+    public AsyncSocket,
     public base::SupportsWeakPtr<DeterministicMockTCPClientSocket> {
  public:
   DeterministicMockTCPClientSocket(net::NetLog* net_log,
@@ -691,7 +701,7 @@ class DeterministicMockTCPClientSocket : public MockClientSocket,
   virtual int64 NumBytesRead() const;
   virtual base::TimeDelta GetConnectTimeMicros() const;
 
-  // MockClientSocket:
+  // AsyncSocket:
   virtual void OnReadComplete(const MockRead& data);
 
  private:
@@ -709,7 +719,7 @@ class DeterministicMockTCPClientSocket : public MockClientSocket,
   bool was_used_to_convey_data_;
 };
 
-class MockSSLClientSocket : public MockClientSocket {
+class MockSSLClientSocket : public MockClientSocket, public AsyncSocket {
  public:
   MockSSLClientSocket(
       net::ClientSocketHandle* transport_socket,
@@ -755,7 +765,8 @@ class MockSSLClientSocket : public MockClientSocket {
   bool was_used_to_convey_data_;
 };
 
-class MockUDPClientSocket : public DatagramClientSocket {
+class MockUDPClientSocket : public DatagramClientSocket,
+    public AsyncSocket {
  public:
   explicit MockUDPClientSocket(SocketDataProvider* data);
   virtual ~MockUDPClientSocket();
@@ -776,12 +787,26 @@ class MockUDPClientSocket : public DatagramClientSocket {
   // DatagramClientSocket interface
   virtual int Connect(const IPEndPoint& address);
 
+  // AsyncSocket interface
+  virtual void OnReadComplete(const MockRead& data);
+
  private:
+  int CompleteRead();
+
   void RunCallbackAsync(net::CompletionCallback* callback, int result);
   void RunCallback(net::CompletionCallback* callback, int result);
 
   bool connected_;
   SocketDataProvider* data_;
+  int read_offset_;
+  net::MockRead read_data_;
+  bool need_read_data_;
+
+  // While an asynchronous IO is pending, we save our user-buffer state.
+  net::IOBuffer* pending_buf_;
+  int pending_buf_len_;
+  net::CompletionCallback* pending_callback_;
+
   ScopedRunnableMethodFactory<MockUDPClientSocket> method_factory_;
 
   DISALLOW_COPY_AND_ASSIGN(MockUDPClientSocket);