Multicast DNS implementation (initial)

net/dns/mdns_client_unittest.cc

 // Copyright (c) 2013 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 "base/memory/ref_counted.h"
 #include "base/message_loop.h"
 #include "net/base/rand_callback.h"
 #include "net/base/test_completion_callback.h"
 #include "net/dns/mdns_client_impl.h"
 #include "net/dns/record_rdata.h"
 #include "net/udp/udp_client_socket.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 using ::testing::Invoke;
 using ::testing::InvokeWithoutArgs;
 using ::testing::StrictMock;
 using ::testing::Exactly;
 using ::testing::Return;
+using ::testing::SaveArg;
 using ::testing::_;
 
 namespace net {
 
 namespace {
 
 const char kSamplePacket1[] = {
     // Header
     0x00, 0x00,               // ID is zeroed out
     0x81, 0x80,               // Standard query response, RA, no error
@@ skipping 149 matching lines @@
     0x00, 0x01,        // TTL (4 bytes) is 20 hours, 47 minutes, 48 seconds.
     0x24, 0x74,
     0x00, 0x08,        // RDLENGTH is 8 bytes.
     0x05, 'z', 'z', 'z', 'z', 'z',
     0xc0, 0x32
   };
 
 const char kQueryPacketPrivet[] = {
   // Header
   0x00, 0x00,               // ID is zeroed out
-  0x00, 0x00,               // Standard query.
+  0x01, 0x00,               // RD flag set
   0x00, 0x01,               // One question.
   0x00, 0x00,               // 0 RRs (answers)
   0x00, 0x00,               // 0 authority RRs
   0x00, 0x00,               // 0 additional RRs
 
   // Question
   // This part is echoed back from the respective query.
   0x07, '_', 'p', 'r', 'i', 'v', 'e', 't',
   0x04, '_', 't', 'c', 'p',
   0x05, 'l', 'o', 'c', 'a', 'l',
@@ skipping 30 matching lines @@
 class MockMDnsConnection : public MDnsConnection {
  public:
   MockMDnsConnection(MDnsConnection::Delegate* delegate,
                      MockMDnsConnectionFactory* factory)
       : delegate_(delegate), factory_(factory) {
   }
 
   virtual ~MockMDnsConnection() {
   }
 
-  virtual bool Init() OVERRIDE {
-    return true;
+  virtual int Init() OVERRIDE {
+    return OK;
   }
 
-  virtual bool Send(IOBuffer* buffer, unsigned size) OVERRIDE;
+  virtual int Send(IOBuffer* buffer, unsigned size) OVERRIDE;
 
   void SimulateRecieve(const char* packet, int size);
 
  private:
   DnsResponse response;
   MDnsConnection::Delegate* delegate_;
   MockMDnsConnectionFactory* factory_;
 };
 
 class MockMDnsConnectionFactory : public MDnsConnectionFactory {
  public:
   MockMDnsConnectionFactory() {}
   virtual ~MockMDnsConnectionFactory() {}
   virtual scoped_ptr<MDnsConnection> CreateConnection(
       MDnsConnection::Delegate* delegate) OVERRIDE;
 
-  MOCK_METHOD1(OnSend, void(std::string));
+  MOCK_METHOD1(OnSend, void(std::string packet));
 
   MockMDnsConnection* latest_connection() { return latest_connection_; }
 
  private:
   MockMDnsConnection* latest_connection_;
 };
 
 class MockMDnsConnectionDelegate : public MDnsConnection::Delegate {
  public:
   virtual void HandlePacket(DnsResponse* response, int size) {
-    OnHandlePacket(std::string(response->io_buffer()->data(), size));
+    HandlePacketInternal(std::string(response->io_buffer()->data(), size));
   }
 
-  MOCK_METHOD1(OnHandlePacket, void(std::string));
+  MOCK_METHOD1(HandlePacketInternal, void(std::string packet));
+
+  MOCK_METHOD1(OnConnectionError, void(int error));
 };
 
-bool MockMDnsConnection::Send(IOBuffer* buffer, unsigned size) {
+int MockMDnsConnection::Send(IOBuffer* buffer, unsigned size) {
   factory_->OnSend(std::string(buffer->data(), size));
-  return true;
+  return OK;
 }
 
 void MockMDnsConnection::SimulateRecieve(const char* packet, int size) {
   memcpy(response.io_buffer()->data(), packet, size);
   delegate_->HandlePacket(&response, size);
 }
 
 scoped_ptr<MDnsConnection> MockMDnsConnectionFactory::CreateConnection(
     MDnsConnection::Delegate* delegate) {
   latest_connection_ = new MockMDnsConnection(delegate, this);
@@ skipping 27 matching lines @@
 class MDnsTest : public ::testing::Test {
  public:
   MDnsTest();
   virtual ~MDnsTest();
   void DeleteTransaction();
   void DeleteBothListeners();
   void RunUntilIdle();
   void RunFor(base::TimeDelta time_period);
   void Stop();
 
-  MOCK_METHOD2(MockableRecordCallback, void(MDnsTransactionResult,
-                                            const RecordParsed*));
+  MOCK_METHOD2(MockableRecordCallback, void(MDnsTransactionResult result,
+                                            const RecordParsed* record));
 
  protected:
   void ExpectPacket(const char* packet, unsigned size);
   void SendPacket(const char* packet, unsigned size);
   void ExpectPtrRecord(
       const std::string& name,
       const std::string& ptrdomain,
       const RecordParsedCopyContainer& record);
 
   base::MessageLoop* message_loop_;
 
   scoped_ptr<MDnsClientImpl> test_client_;
   IPEndPoint mdns_ipv4_endpoint_;
   StrictMock<MockMDnsConnectionFactory> connection_factory_;
 
   // Transactions and listeners that can be deleted by class methods for
   // reentrancy tests.
   scoped_ptr<MDnsTransaction> transaction_;
   scoped_ptr<MDnsListener> listener1_;
   scoped_ptr<MDnsListener> listener2_;
 };
 
 class MockListenerDelegate : public MDnsListener::Delegate {
  public:
   MOCK_METHOD2(OnRecordUpdate,
-               void(MDnsUpdateType, const RecordParsed*));
+               void(MDnsUpdateType update, const RecordParsed* records));
   MOCK_METHOD2(OnNsecRecord, void(const std::string&, unsigned));
+  MOCK_METHOD0(OnCachePurged, void());
 };
 
 MDnsTest::MDnsTest()
     : message_loop_(base::MessageLoop::current()) {
   test_client_.reset(new MDnsClientImpl(&connection_factory_));
 }
 
 MDnsTest::~MDnsTest() {
 }
 
@@ skipping 49 matching lines @@
 
 TEST_F(MDnsTest, PassiveListeners) {
   StrictMock<MockListenerDelegate> delegate_privet;
   StrictMock<MockListenerDelegate> delegate_printer;
   StrictMock<MockListenerDelegate> delegate_ptr;
 
   RecordParsedCopyContainer record_privet;
   RecordParsedCopyContainer record_printer;
 
   scoped_ptr<MDnsListener> listener_privet = test_client_->CreateListener(
-      dns_protocol::kTypePTR, "_privet._tcp.local", false, false,
-      &delegate_privet);
+      dns_protocol::kTypePTR, "_privet._tcp.local", &delegate_privet);
   scoped_ptr<MDnsListener> listener_printer = test_client_->CreateListener(
-      dns_protocol::kTypePTR, "_printer._tcp.local", false, false,
-      &delegate_printer);
+      dns_protocol::kTypePTR, "_printer._tcp.local", &delegate_printer);
   scoped_ptr<MDnsListener> listener_ptr = test_client_->CreateListener(
-      dns_protocol::kTypePTR, "", false, false, &delegate_ptr);
+      dns_protocol::kTypePTR, "", &delegate_ptr);
 
   listener_privet->Start();
   listener_printer->Start();
   listener_ptr->Start();
 
   ASSERT_TRUE(test_client_->IsListeningForTests());
 
   // Send the same packet twice to ensure no records are double-counted.
 
   EXPECT_CALL(delegate_privet, OnRecordUpdate(kMDnsRecordAdded, _))
@@ skipping 42 matching lines @@
   ASSERT_FALSE(test_client_->IsListeningForTests());
 }
 
 TEST_F(MDnsTest, PassiveListenersCleanup) {
   StrictMock<MockListenerDelegate> delegate_privet;
 
   RecordParsedCopyContainer record_privet;
   RecordParsedCopyContainer record_privet2;
 
   scoped_ptr<MDnsListener> listener_privet = test_client_->CreateListener(
-      dns_protocol::kTypePTR, "_privet._tcp.local", false, false,
-      &delegate_privet);
+      dns_protocol::kTypePTR, "_privet._tcp.local", &delegate_privet);
 
   listener_privet->Start();
 
   ASSERT_TRUE(test_client_->IsListeningForTests());
 
   EXPECT_CALL(delegate_privet, OnRecordUpdate(kMDnsRecordAdded, _))
       .Times(Exactly(1))
       .WillOnce(Invoke(
           &record_privet,
           &RecordParsedCopyContainer::SaveWithDummyArg));
@@ skipping 12 matching lines @@
                              &RecordParsedCopyContainer::SaveWithDummyArg)));
 
   RunFor(base::TimeDelta::FromSeconds(record_privet->ttl() + 1));
 
   RunUntilIdle();
 
   ExpectPtrRecord("_privet._tcp.local", "hello._privet._tcp.local",
                   record_privet2);
 }
 
-TEST_F(MDnsTest, PassiveListenersNotifyExisting) {
-  StrictMock<MockListenerDelegate> delegate_privet;
-  StrictMock<MockListenerDelegate> delegate_privet2;
-
-  RecordParsedCopyContainer record_privet;
-  RecordParsedCopyContainer record_privet2;
-
-  scoped_ptr<MDnsListener> listener_privet = test_client_->CreateListener(
-      dns_protocol::kTypePTR, "_privet._tcp.local", false, false,
-      &delegate_privet);
-
-  listener_privet->Start();
-
-  ASSERT_TRUE(test_client_->IsListeningForTests());
-
-  EXPECT_CALL(delegate_privet, OnRecordUpdate(kMDnsRecordAdded, _))
-      .Times(Exactly(1))
-      .WillOnce(Invoke(
-          &record_privet,
-          &RecordParsedCopyContainer::SaveWithDummyArg));
-
-  SendPacket(kSamplePacket1, sizeof(kSamplePacket1));
-
-  RunUntilIdle();
-
-  ExpectPtrRecord("_privet._tcp.local", "hello._privet._tcp.local",
-                  record_privet);
-
-  EXPECT_CALL(delegate_privet2, OnRecordUpdate(kMDnsRecordAdded, _))
-      .Times(Exactly(1))
-      .WillOnce(Invoke(
-          &record_privet2,
-          &RecordParsedCopyContainer::SaveWithDummyArg));
-
-  scoped_ptr<MDnsListener> listener_privet2 = test_client_->CreateListener(
-      dns_protocol::kTypePTR, "_privet._tcp.local", false, true,
-      &delegate_privet2);
-
-  listener_privet2->Start();
-
-  RunUntilIdle();
-
-  ExpectPtrRecord("_privet._tcp.local", "hello._privet._tcp.local",
-                  record_privet2);
-}
-
 TEST_F(MDnsTest, MalformedPacket) {
   StrictMock<MockListenerDelegate> delegate_printer;
 
   RecordParsedCopyContainer record_printer;
 
   scoped_ptr<MDnsListener> listener_printer = test_client_->CreateListener(
-      dns_protocol::kTypePTR, "_printer._tcp.local", false, false,
-      &delegate_printer);
+      dns_protocol::kTypePTR, "_printer._tcp.local", &delegate_printer);
 
   listener_printer->Start();
 
   ASSERT_TRUE(test_client_->IsListeningForTests());
 
   EXPECT_CALL(delegate_printer, OnRecordUpdate(kMDnsRecordAdded, _))
       .Times(Exactly(1))
       .WillOnce(Invoke(
           &record_printer,
           &RecordParsedCopyContainer::SaveWithDummyArg));
 
   // First, send unsalvagable packet to ensure we can deal with it.
   SendPacket(kCorruptedPacketUnsalvagable,
              sizeof(kCorruptedPacketUnsalvagable));
 
   // Regression test: send a packet where the question cannot be read.
   SendPacket(kCorruptedPacketBadQuestion,
              sizeof(kCorruptedPacketBadQuestion));
 
   // Then send salvagable packet to ensure we can extract useful records.
   SendPacket(kCorruptedPacketSalvagable,
              sizeof(kCorruptedPacketSalvagable));
 
   RunUntilIdle();
 
   ExpectPtrRecord("_printer._tcp.local", "hello._printer._tcp.local",
                   record_printer);
 }
 
-TEST_F(MDnsTest, QueryCache) {
-  StrictMock<MockListenerDelegate> delegate_privet;
-
-  RecordParsedCopyContainer record_privet;
-  std::vector<const RecordParsed*> records_from_cache;
-
-  scoped_ptr<MDnsListener> listener_privet = test_client_->CreateListener(
-      dns_protocol::kTypePTR, "_privet._tcp.local", false, false,
-      &delegate_privet);
-
-  listener_privet->Start();
-
-  ASSERT_TRUE(test_client_->IsListeningForTests());
-
-  EXPECT_CALL(delegate_privet, OnRecordUpdate(kMDnsRecordAdded, _))
-      .Times(Exactly(1))
-      .WillOnce(Invoke(
-          &record_privet,
-          &RecordParsedCopyContainer::SaveWithDummyArg));
-
-  SendPacket(kSamplePacket1, sizeof(kSamplePacket1));
-
-  RunUntilIdle();
-
-  ExpectPtrRecord("_privet._tcp.local", "hello._privet._tcp.local",
-                  record_privet);
-
-  listener_privet->QueryCache(&records_from_cache);
-
-  EXPECT_EQ(1u, records_from_cache.size());
-  EXPECT_EQ("_privet._tcp.local", records_from_cache.front()->name());
-  EXPECT_EQ(dns_protocol::kTypePTR, records_from_cache.front()->type());
-  EXPECT_EQ(dns_protocol::kClassIN, records_from_cache.front()->klass());
-
-  const PtrRecordRdata* ptr_rdata =
-      records_from_cache.front()->rdata<PtrRecordRdata>();
-
-  EXPECT_TRUE(ptr_rdata != NULL);
-
-  EXPECT_EQ("hello._privet._tcp.local",
-            ptr_rdata->ptrdomain());
-}
-
-TEST_F(MDnsTest, Query) {
-  StrictMock<MockListenerDelegate> delegate_privet;
-  ExpectPacket(kQueryPacketPrivet, sizeof(kQueryPacketPrivet));
-
-  scoped_ptr<MDnsListener> listener_privet = test_client_->CreateListener(
-      dns_protocol::kTypePTR, "_privet._tcp.local", false, false,
-      &delegate_privet);
-
-  listener_privet->Start();
-
-  ASSERT_TRUE(listener_privet->SendQuery(false));
-
-  // Active listeners should send queries when created.
-  ExpectPacket(kQueryPacketPrivet, sizeof(kQueryPacketPrivet));
-
-  scoped_ptr<MDnsListener> listener_privet2 = test_client_->CreateListener(
-      dns_protocol::kTypePTR, "_privet._tcp.local", true /*active*/,
-      false /*existing*/, &delegate_privet);
-
-  listener_privet2->Start();
-}
-
 TEST_F(MDnsTest, TransactionNoCache) {
   ExpectPacket(kQueryPacketPrivet, sizeof(kQueryPacketPrivet));
 
   scoped_ptr<MDnsTransaction> transaction_privet =
       test_client_->CreateTransaction(
           dns_protocol::kTypePTR, "_privet._tcp.local",
+          kMDnsTransactionQueryNetwork |
+          kMDnsTransactionQueryCache |
+          kMDnsTransactionSingleResult,
           base::Bind(&MDnsTest::MockableRecordCallback,
                      base::Unretained(this)));
 
   transaction_privet->Start();
 
   EXPECT_TRUE(test_client_->IsListeningForTests());
 
   RecordParsedCopyContainer record_privet;
 
-  EXPECT_CALL(*this, MockableRecordCallback(kMDnsTransactionSuccess, _))
+  EXPECT_CALL(*this, MockableRecordCallback(kMDnsTransactionRecord, _))
       .Times(Exactly(1))
       .WillOnce(Invoke(&record_privet,
                        &RecordParsedCopyContainer::SaveWithDummyArg));
 
   SendPacket(kSamplePacket1, sizeof(kSamplePacket1));
 
   RunUntilIdle();
 
   ExpectPtrRecord("_privet._tcp.local", "hello._privet._tcp.local",
                   record_privet);
 
   EXPECT_FALSE(test_client_->IsListeningForTests());
 }
 
+
+TEST_F(MDnsTest, TransactionCacheOnlyNoResult) {
+  scoped_ptr<MDnsTransaction> transaction_privet =
+      test_client_->CreateTransaction(
+          dns_protocol::kTypePTR, "_privet._tcp.local",
+          kMDnsTransactionQueryCache |
+          kMDnsTransactionSingleResult,
+          base::Bind(&MDnsTest::MockableRecordCallback,
+                     base::Unretained(this)));
+
+  EXPECT_CALL(*this, MockableRecordCallback(kMDnsTransactionNoResults, _))
+      .Times(Exactly(1));
+
+  transaction_privet->Start();
+
+  EXPECT_FALSE(test_client_->IsListeningForTests());
+
+  RunUntilIdle();
+}
+
+
 TEST_F(MDnsTest, TransactionWithCache) {
   // Listener to force the client to listen
   StrictMock<MockListenerDelegate> delegate_irrelevant;
   scoped_ptr<MDnsListener> listener_irrelevant = test_client_->CreateListener(
-      dns_protocol::kTypeA, "codereview.chromium.local", false, false,
+      dns_protocol::kTypeA, "codereview.chromium.local",
       &delegate_irrelevant);
 
   listener_irrelevant->Start();
 
   EXPECT_TRUE(test_client_->IsListeningForTests());
 
   SendPacket(kSamplePacket1, sizeof(kSamplePacket1));
 
   RunUntilIdle();
 
   RecordParsedCopyContainer record_privet;
 
-  EXPECT_CALL(*this, MockableRecordCallback(kMDnsTransactionSuccess, _))
+  EXPECT_CALL(*this, MockableRecordCallback(kMDnsTransactionRecord, _))
       .WillOnce(Invoke(&record_privet,
                        &RecordParsedCopyContainer::SaveWithDummyArg));
 
   scoped_ptr<MDnsTransaction> transaction_privet =
       test_client_->CreateTransaction(
           dns_protocol::kTypePTR, "_privet._tcp.local",
+          kMDnsTransactionQueryNetwork |
+          kMDnsTransactionQueryCache |
+          kMDnsTransactionSingleResult,
           base::Bind(&MDnsTest::MockableRecordCallback,
                      base::Unretained(this)));
 
   transaction_privet->Start();
 
   RunUntilIdle();
 
   ExpectPtrRecord("_privet._tcp.local", "hello._privet._tcp.local",
                   record_privet);
 }
 
 TEST_F(MDnsTest, AdditionalRecords) {
   StrictMock<MockListenerDelegate> delegate_privet;
 
   RecordParsedCopyContainer record_privet;
 
   scoped_ptr<MDnsListener> listener_privet = test_client_->CreateListener(
-      dns_protocol::kTypePTR, "_privet._tcp.local", false, false,
+      dns_protocol::kTypePTR, "_privet._tcp.local",
       &delegate_privet);
 
   listener_privet->Start();
 
   ASSERT_TRUE(test_client_->IsListeningForTests());
 
   EXPECT_CALL(delegate_privet, OnRecordUpdate(kMDnsRecordAdded, _))
       .Times(Exactly(1))
       .WillOnce(Invoke(
           &record_privet,
           &RecordParsedCopyContainer::SaveWithDummyArg));
 
   SendPacket(kSamplePacketAdditionalOnly, sizeof(kSamplePacket1));
 
   RunUntilIdle();
 
   ExpectPtrRecord("_privet._tcp.local", "hello._privet._tcp.local",
                   record_privet);
 }
 
 TEST_F(MDnsTest, TransactionTimeout) {
   ExpectPacket(kQueryPacketPrivet, sizeof(kQueryPacketPrivet));
 
   scoped_ptr<MDnsTransaction> transaction_privet =
       test_client_->CreateTransaction(
           dns_protocol::kTypePTR, "_privet._tcp.local",
+          kMDnsTransactionQueryNetwork |
+          kMDnsTransactionQueryCache |
+          kMDnsTransactionSingleResult,
           base::Bind(&MDnsTest::MockableRecordCallback,
                      base::Unretained(this)));
 
   transaction_privet->Start();
 
   EXPECT_TRUE(test_client_->IsListeningForTests());
 
-  EXPECT_CALL(*this, MockableRecordCallback(kMDnsTransactionTimeout, NULL))
+  EXPECT_CALL(*this, MockableRecordCallback(kMDnsTransactionNoResults, NULL))
       .Times(Exactly(1))
       .WillOnce(InvokeWithoutArgs(this, &MDnsTest::Stop));
 
   RunFor(base::TimeDelta::FromSeconds(4));
 
   EXPECT_FALSE(test_client_->IsListeningForTests());
 }
 
+TEST_F(MDnsTest, TransactionMultipleRecords) {
+  ExpectPacket(kQueryPacketPrivet, sizeof(kQueryPacketPrivet));
+
+  scoped_ptr<MDnsTransaction> transaction_privet =
+      test_client_->CreateTransaction(
+          dns_protocol::kTypePTR, "_privet._tcp.local",
+          kMDnsTransactionQueryNetwork |
+          kMDnsTransactionQueryCache ,
+          base::Bind(&MDnsTest::MockableRecordCallback,
+                     base::Unretained(this)));
+
+  transaction_privet->Start();
+
+  EXPECT_TRUE(test_client_->IsListeningForTests());
+
+  RecordParsedCopyContainer record_privet;
+  RecordParsedCopyContainer record_privet2;
+
+  EXPECT_CALL(*this, MockableRecordCallback(kMDnsTransactionRecord, _))
+      .Times(Exactly(2))
+      .WillOnce(Invoke(&record_privet,
+                       &RecordParsedCopyContainer::SaveWithDummyArg))
+      .WillOnce(Invoke(&record_privet2,
+                       &RecordParsedCopyContainer::SaveWithDummyArg));
+
+  SendPacket(kSamplePacket1, sizeof(kSamplePacket1));
+  SendPacket(kSamplePacket2, sizeof(kSamplePacket2));
+
+  RunUntilIdle();
+
+  ExpectPtrRecord("_privet._tcp.local", "hello._privet._tcp.local",
+                  record_privet);
+
+  ExpectPtrRecord("_privet._tcp.local", "zzzzz._privet._tcp.local",
+                  record_privet2);
+
+  EXPECT_CALL(*this, MockableRecordCallback(kMDnsTransactionDone, NULL))
+      .WillOnce(InvokeWithoutArgs(this, &MDnsTest::Stop));
+
+  RunFor(base::TimeDelta::FromSeconds(4));
+
+  EXPECT_FALSE(test_client_->IsListeningForTests());
+}
+
+
 TEST_F(MDnsTest, TransactionReentrantDelete) {
   ExpectPacket(kQueryPacketPrivet, sizeof(kQueryPacketPrivet));
 
   transaction_ = test_client_->CreateTransaction(
       dns_protocol::kTypePTR, "_privet._tcp.local",
+      kMDnsTransactionQueryNetwork |
+      kMDnsTransactionQueryCache |
+      kMDnsTransactionSingleResult,
       base::Bind(&MDnsTest::MockableRecordCallback,
                  base::Unretained(this)));
 
   transaction_->Start();
 
   EXPECT_TRUE(test_client_->IsListeningForTests());
 
-  EXPECT_CALL(*this, MockableRecordCallback(kMDnsTransactionTimeout, NULL))
+  EXPECT_CALL(*this, MockableRecordCallback(kMDnsTransactionNoResults, NULL))
       .Times(Exactly(1))
       .WillOnce(DoAll(InvokeWithoutArgs(this, &MDnsTest::DeleteTransaction),
                       InvokeWithoutArgs(this, &MDnsTest::Stop)));
 
 
   RunFor(base::TimeDelta::FromSeconds(4));
 
   EXPECT_EQ(NULL, transaction_.get());
 
   EXPECT_FALSE(test_client_->IsListeningForTests());
 }
 
+
+TEST_F(MDnsTest, TransactionReentrantDeleteFromCache) {
+  StrictMock<MockListenerDelegate> delegate_irrelevant;
+  scoped_ptr<MDnsListener> listener_irrelevant = test_client_->CreateListener(
+      dns_protocol::kTypeA, "codereview.chromium.local",
+      &delegate_irrelevant);
+  listener_irrelevant->Start();
+
+  ASSERT_TRUE(test_client_->IsListeningForTests());
+
+  SendPacket(kSamplePacket1, sizeof(kSamplePacket1));
+
+  transaction_ = test_client_->CreateTransaction(
+      dns_protocol::kTypePTR, "_privet._tcp.local",
+      kMDnsTransactionQueryNetwork |
+      kMDnsTransactionQueryCache,
+      base::Bind(&MDnsTest::MockableRecordCallback,
+                 base::Unretained(this)));
+
+
+  EXPECT_CALL(*this, MockableRecordCallback(kMDnsTransactionRecord, _))
+      .Times(Exactly(1))
+      .WillOnce(InvokeWithoutArgs(this, &MDnsTest::DeleteTransaction));
+
+  transaction_->Start();
+
+
+  RunUntilIdle();
+
+  EXPECT_EQ(NULL, transaction_.get());
+}
+
+
 // In order to reliably test reentrant listener deletes, we create two listeners
 // and have each of them delete both, so we're guaranteed to try and deliver a
 // callback to at least one deleted listener.
 
 TEST_F(MDnsTest, ListenerReentrantDelete) {
   StrictMock<MockListenerDelegate> delegate_privet;
 
   listener1_ = test_client_->CreateListener(
-      dns_protocol::kTypePTR, "_privet._tcp.local", false, false,
+      dns_protocol::kTypePTR, "_privet._tcp.local",
       &delegate_privet);
 
   listener2_ = test_client_->CreateListener(
-      dns_protocol::kTypePTR, "_privet._tcp.local", false, false,
+      dns_protocol::kTypePTR, "_privet._tcp.local",
       &delegate_privet);
 
   listener1_->Start();
 
   listener2_->Start();
 
   EXPECT_CALL(delegate_privet, OnRecordUpdate(kMDnsRecordAdded, _))
       .Times(Exactly(1))
       .WillOnce(InvokeWithoutArgs(this, &MDnsTest::DeleteBothListeners));
 
   EXPECT_TRUE(test_client_->IsListeningForTests());
 
   SendPacket(kSamplePacket1, sizeof(kSamplePacket1));
 
   RunUntilIdle();
 
   EXPECT_EQ(NULL, listener1_.get());
   EXPECT_EQ(NULL, listener2_.get());
 
   EXPECT_FALSE(test_client_->IsListeningForTests());
 }
 
-// Connection tests. These tests are disabled because, due to their dependence
-// on multicast networking, they are falky.
+class MockDatagramServerSocket : public DatagramServerSocket {
+ public:
+  int Listen(const IPEndPoint& address) {
+    return ListenInternal(address.ToString());
+  }
 
-// TODO (noamsml): Migrate these tests to use a MockDatagramServerSocket.
+  MOCK_METHOD1(ListenInternal, int(const std::string& address));
+
+  void SetResponsePacket(std::string response_packet) {
+    response_packet_ = response_packet;
+  }
+
+  int RespondImmediately(IOBuffer* buffer, int size, IPEndPoint* address,
+                         const CompletionCallback& callback) {
+    int to_copy = std::min(response_packet_.size(), (unsigned)size);
+    memcpy(buffer->data(), response_packet_.data(), to_copy);
+    return to_copy;
+  }
+
+  int RespondDelayed(IOBuffer* buffer, int size, IPEndPoint* address,
+                     const CompletionCallback& callback) {
+    int rv = RespondImmediately(buffer, size, address, callback);
+    MessageLoop::current()->PostTask(FROM_HERE, base::Bind(callback, rv));
+    return ERR_IO_PENDING;
+  }
+
+  MOCK_METHOD4(RecvFrom, int(IOBuffer* buffer, int size,
+                             IPEndPoint* address,
+                             const CompletionCallback& callback));
+
+  int SendTo(IOBuffer* buf, int buf_len, const IPEndPoint& address,
+             const CompletionCallback& callback) {
+    return SendToInternal(std::string(buf->data(), buf_len), address.ToString(),
+                          callback);
+  }
+
+  MOCK_METHOD3(SendToInternal, int(const std::string&, const std::string,
+                                   const CompletionCallback&));
+
+  MOCK_METHOD1(SetReceiveBufferSize, bool(int32 size));
+  MOCK_METHOD1(SetSendBufferSize, bool(int32 size));
+
+  MOCK_METHOD0(Close, void());
+
+  MOCK_CONST_METHOD1(GetPeerAddress, int(IPEndPoint* address));
+  MOCK_CONST_METHOD1(GetLocalAddress, int(IPEndPoint* address));
+  MOCK_CONST_METHOD0(NetLog, const BoundNetLog&());
+
+  MOCK_METHOD0(AllowAddressReuse, void());
+  MOCK_METHOD0(AllowBroadcast, void());
+
+  int JoinGroup(const IPAddressNumber& group_address) const {
+    return JoinGroupInternal(IPAddressToString(group_address));
+  }
+
+  MOCK_CONST_METHOD1(JoinGroupInternal, int(const std::string& group));
+
+  int LeaveGroup(const IPAddressNumber& group_address) const {
+    return JoinGroupInternal(IPAddressToString(group_address));
+  }
+
+  MOCK_CONST_METHOD1(LeaveGroupInternal, int(const std::string& group));
+
+  MOCK_METHOD1(SetMulticastTimeToLive, int(int ttl));
+
+  MOCK_METHOD1(SetMulticastLoopbackMode, int(bool loopback));
+
+ private:
+  std::string response_packet_;
+};
 
 class MDnsConnectionTest : public ::testing::Test {
  public:
-  MDnsConnectionTest() : connection_(&delegate_) {
-    EXPECT_TRUE(ParseIPLiteralToNumber("224.0.0.251", &multicast_address4_));
-    // TODO(noamsml): Figure out why client socket fails to connect to ipv6
-    // multicast address.
-    EXPECT_TRUE(ParseIPLiteralToNumber("::1", &multicast_address6_));
+  MDnsConnectionTest() : connection_(&socket_ipv4_, &socket_ipv6_, &delegate_) {
+  }
+
+ protected:
+  // Follow successful connection initialization
+  void InitConnection() {
+    EXPECT_CALL(socket_ipv4_, AllowAddressReuse());
+    EXPECT_CALL(socket_ipv6_, AllowAddressReuse());
+
+    EXPECT_CALL(socket_ipv4_, SetMulticastLoopbackMode(false));
+    EXPECT_CALL(socket_ipv6_, SetMulticastLoopbackMode(false));
+
+    EXPECT_CALL(socket_ipv4_, ListenInternal("0.0.0.0:5353"))
+        .WillOnce(Return(OK));
+    EXPECT_CALL(socket_ipv6_, ListenInternal("[::]:5353"))
+        .WillOnce(Return(OK));
+
+    EXPECT_CALL(socket_ipv4_, JoinGroupInternal("224.0.0.251"))
+        .WillOnce(Return(OK));
+    EXPECT_CALL(socket_ipv6_, JoinGroupInternal("ff02::fb"))
+        .WillOnce(Return(OK));
+
     connection_.Init();
   }
- protected:
-  IPAddressNumber multicast_address4_;
-  IPAddressNumber multicast_address6_;
+
   StrictMock<MockMDnsConnectionDelegate> delegate_;
 
+  MockDatagramServerSocket socket_ipv4_;
+  MockDatagramServerSocket socket_ipv6_;
   MDnsConnectionImpl connection_;
   TestCompletionCallback callback_;
 };
 
-TEST_F(MDnsConnectionTest, DISABLED_Recieve4) {
-  scoped_refptr<IOBuffer> buf(new IOBuffer(sizeof(kSamplePacket1)));
-  memcpy(buf->data(), kSamplePacket1, sizeof(kSamplePacket1));
-  UDPSocket socket(DatagramSocket::DEFAULT_BIND,
-                   RandIntCallback(),
-                   NULL, NetLog::Source());
-  EXPECT_CALL(delegate_, OnHandlePacket(std::string(kSamplePacket1,
-                                                    sizeof(kSamplePacket1))));
+TEST_F(MDnsConnectionTest, RecieveSynchronous) {
+  std::string sample_packet =
+      std::string(kSamplePacket1, sizeof(kSamplePacket1));
+  socket_ipv6_.SetResponsePacket(sample_packet);
+  EXPECT_CALL(socket_ipv4_, RecvFrom(_, _, _, _))
+      .WillOnce(Return(ERR_IO_PENDING));
+  EXPECT_CALL(socket_ipv6_, RecvFrom(_, _, _, _))
+      .WillOnce(
+          Invoke(&socket_ipv6_, &MockDatagramServerSocket::RespondImmediately))
+      .WillOnce(Return(ERR_IO_PENDING));
 
-  EXPECT_EQ(OK, socket.Connect(IPEndPoint(multicast_address4_, 5353)));
-  int rv = socket.Write(buf, sizeof(kSamplePacket1), callback_.callback());
-  if (rv == ERR_IO_PENDING) {
-    rv = callback_.GetResult(rv);
-  }
-  EXPECT_GT(rv, OK);
-  base::MessageLoop::current()->RunUntilIdle();  // Socket uses message loop.
+  EXPECT_CALL(delegate_, HandlePacketInternal(sample_packet));
+
+  InitConnection();
 }
 
-TEST_F(MDnsConnectionTest, DISABLED_Recieve6) {
-  scoped_refptr<IOBuffer> buf(new IOBuffer(sizeof(kSamplePacket2)));
-  memcpy(buf->data(), kSamplePacket2, sizeof(kSamplePacket2));
-  UDPSocket socket(DatagramSocket::DEFAULT_BIND,
-                   RandIntCallback(),
-                   NULL, NetLog::Source());
-  EXPECT_CALL(delegate_, OnHandlePacket(std::string(kSamplePacket2,
-                                                    sizeof(kSamplePacket2))));
+TEST_F(MDnsConnectionTest, RecieveAsynchronous) {
+  std::string sample_packet =
+      std::string(kSamplePacket1, sizeof(kSamplePacket1));
+  socket_ipv6_.SetResponsePacket(sample_packet);
+  EXPECT_CALL(socket_ipv4_, RecvFrom(_, _, _, _))
+      .WillOnce(Return(ERR_IO_PENDING));
+  EXPECT_CALL(socket_ipv6_, RecvFrom(_, _, _, _))
+      .WillOnce(
+          Invoke(&socket_ipv6_, &MockDatagramServerSocket::RespondDelayed))
+      .WillOnce(Return(ERR_IO_PENDING));
 
-  EXPECT_EQ(OK, socket.Connect(IPEndPoint(multicast_address6_, 5353)));
-  int rv = socket.Write(buf, sizeof(kSamplePacket2), callback_.callback());
-  if (rv == ERR_IO_PENDING) {
-    rv = callback_.GetResult(rv);
-  }
-  EXPECT_GT(rv, OK);
-  base::MessageLoop::current()->RunUntilIdle();  // Socket uses message loop.
+  InitConnection();
+
+  EXPECT_CALL(delegate_, HandlePacketInternal(sample_packet));
+
+  base::MessageLoop::current()->RunUntilIdle();
+}
+
+TEST_F(MDnsConnectionTest, Send) {
+  std::string sample_packet =
+      std::string(kSamplePacket1, sizeof(kSamplePacket1));
+
+  scoped_refptr<IOBufferWithSize> buf(
+      new IOBufferWithSize(sizeof kSamplePacket1));
+  memcpy(buf->data(), kSamplePacket1, sizeof(kSamplePacket1));
+
+  EXPECT_CALL(socket_ipv4_, RecvFrom(_, _, _, _))
+      .WillOnce(Return(ERR_IO_PENDING));
+  EXPECT_CALL(socket_ipv6_, RecvFrom(_, _, _, _))
+      .WillOnce(Return(ERR_IO_PENDING));
+
+  InitConnection();
+
+  EXPECT_CALL(socket_ipv4_,
+              SendToInternal(sample_packet, "224.0.0.251:5353", _));
+  EXPECT_CALL(socket_ipv6_,
+              SendToInternal(sample_packet, "[ff02::fb]:5353", _));
+
+  connection_.Send(buf, buf->size());
+}
+
+TEST_F(MDnsConnectionTest, Error) {
+  CompletionCallback callback;
+
+  EXPECT_CALL(socket_ipv4_, RecvFrom(_, _, _, _))
+      .WillOnce(Return(ERR_IO_PENDING));
+  EXPECT_CALL(socket_ipv6_, RecvFrom(_, _, _, _))
+      .WillOnce(DoAll(SaveArg<3>(&callback),  Return(ERR_IO_PENDING)));
+
+  InitConnection();
+
+  EXPECT_CALL(delegate_, OnConnectionError(ERR_SOCKET_NOT_CONNECTED));
+  callback.Run(ERR_SOCKET_NOT_CONNECTED);
 }
 
 }  // namespace
 
 }  // namespace net