Browse Source

[trac499] Checkpoint comment

Stephen Morris 14 years ago
parent
commit
9ae710aa95

+ 1 - 1
src/lib/asiolink/io_fetch.h

@@ -106,7 +106,7 @@ public:
 
         /// \brief Callback method
         ///
-        /// This is the method called when the fecth completes.
+        /// This is the method called when the fetch completes.
         ///
         /// \param result Result of the fetch
         virtual void operator()(Result result) = 0;

+ 59 - 9
src/lib/asiolink/recursive_query.cc

@@ -55,10 +55,22 @@ RecursiveQuery::RecursiveQuery(DNSService& dns_service,
     unsigned retries) :
     dns_service_(dns_service), upstream_(new AddressVector(upstream)),
     upstream_root_(new AddressVector(upstream_root)),
+    test_server_("", 0),
     query_timeout_(query_timeout), client_timeout_(client_timeout),
     lookup_timeout_(lookup_timeout), retries_(retries)
 {}
 
+// Set the test server - only used for unit testing.
+
+void
+RecursiveQuery::setTestServer(const std::string& address, uint16_t port) {
+    dlog("Setting test server to " + address + "(" +
+            boost::lexical_cast<std::string>(port) + ")");
+    test_server_.first = address;
+    test_server_.second = port;
+}
+
+
 namespace {
 
 typedef std::pair<std::string, uint16_t> addr_t;
@@ -88,6 +100,10 @@ private:
     // root servers...just copied over to the zone_servers_
     boost::shared_ptr<AddressVector> upstream_root_;
 
+    // Test server - only used for testing.  This takes precedence over all
+    // other servers if the port is non-zero.
+    std::pair<std::string, uint16_t> test_server_;
+
     // Buffer to store the result.
     OutputBufferPtr buffer_;
 
@@ -95,6 +111,12 @@ private:
     //shared_ptr<DNSServer> server_;
     isc::resolve::ResolverInterface::CallbackPtr resolvercallback_;
 
+    // Protocol used for the last query.  This is set to IOFetch::UDP when a
+    // new upstream query is initiated, and changed to IOFetch::TCP if a
+    // packet is returned with the TC bit set.  It is stored here to detect the
+    // case of a TCP packet being returned with the TC bit set.
+    IOFetch::Protocol protocol_;
+
     // To prevent both unreasonably long cname chains and cname loops,
     // we simply keep a counter of the number of CNAMEs we have
     // followed so far (and error if it exceeds RESOLVER_MAX_CNAME_CHAIN
@@ -155,15 +177,27 @@ private:
     }
 
     // (re)send the query to the server.
-    void send() {
+    //
+    // \param protocol Protocol to use for the fetch (default is UDP)
+    void send(IOFetch::Protocol protocol = IOFetch::UDP) {
         const int uc = upstream_->size();
         const int zs = zone_servers_.size();
+        protocol_ = protocol;   // Store protocol being used for this
         buffer_->clear();
-        if (uc > 0) {
+        if (test_server_.second != 0) {
+            dlog("Sending upstream query (" + question_.toText() +
+                 ") to test server at " + test_server_.first);
+            IOFetch query(protocol, io_, question_,
+                test_server_.first,
+                test_server_.second, buffer_, this,
+                query_timeout_);
+            ++queries_out_;
+            io_.get_io_service().post(query);
+        } else if (uc > 0) {
             int serverIndex = rand() % uc;
             dlog("Sending upstream query (" + question_.toText() +
                 ") to " + upstream_->at(serverIndex).first);
-            IOFetch query(IOFetch::UDP, io_, question_,
+            IOFetch query(protocol, io_, question_,
                 upstream_->at(serverIndex).first,
                 upstream_->at(serverIndex).second, buffer_, this,
                 query_timeout_);
@@ -173,7 +207,7 @@ private:
             int serverIndex = rand() % zs;
             dlog("Sending query to zone server (" + question_.toText() +
                 ") to " + zone_servers_.at(serverIndex).first);
-            IOFetch query(IOFetch::UDP, io_, question_,
+            IOFetch query(protocol, io_, question_,
                 zone_servers_.at(serverIndex).first,
                 zone_servers_.at(serverIndex).second, buffer_, this,
                 query_timeout_);
@@ -291,6 +325,18 @@ private:
                 return true;
             }
             break;
+        case isc::resolve::ResponseClassifier::TRUNCATED:
+            // Truncated packet.  If the protocol we used for the last one is
+            // UDP, re-query using TCP.  Otherwise regard it as an error.
+            if (protocol_ == IOFetch::UDP) {
+                dlog("Response truncated, re-querying over TCP");
+                send(IOFetch::TCP);
+                break;
+            }
+            // Was a TCP query so we have received a packet over TCP with the TC
+            // bit set: drop through to common error processing.
+            // TODO: Can we use what we have received instead of discarding it?
+
         case isc::resolve::ResponseClassifier::EMPTY:
         case isc::resolve::ResponseClassifier::EXTRADATA:
         case isc::resolve::ResponseClassifier::INVNAMCLASS:
@@ -302,7 +348,7 @@ private:
         case isc::resolve::ResponseClassifier::NOTSINGLE:
         case isc::resolve::ResponseClassifier::OPCODE:
         case isc::resolve::ResponseClassifier::RCODE:
-        case isc::resolve::ResponseClassifier::TRUNCATED:
+
             // Should we try a different server rather than SERVFAIL?
             isc::resolve::makeErrorMessage(answer_message_,
                                            Rcode::SERVFAIL());
@@ -320,6 +366,7 @@ public:
         MessagePtr answer_message,
         boost::shared_ptr<AddressVector> upstream,
         boost::shared_ptr<AddressVector> upstream_root,
+        std::pair<std::string, uint16_t>& test_server,
         OutputBufferPtr buffer,
         isc::resolve::ResolverInterface::CallbackPtr cb,
         int query_timeout, int client_timeout, int lookup_timeout,
@@ -330,8 +377,10 @@ public:
         answer_message_(answer_message),
         upstream_(upstream),
         upstream_root_(upstream_root),
+        test_server_(test_server),
         buffer_(buffer),
         resolvercallback_(cb),
+        protocol_(IOFetch::UDP),
         cname_count_(0),
         query_timeout_(query_timeout),
         retries_(retries),
@@ -441,7 +490,6 @@ public:
 
     // This function is used as callback from DNSQuery.
     virtual void operator()(IOFetch::Result result) {
-        // XXX is this the place for TCP retry?
         --queries_out_;
         if (!done_ && result != IOFetch::TIME_OUT) {
             // we got an answer
@@ -496,7 +544,8 @@ RecursiveQuery::resolve(const QuestionPtr& question,
         dlog("Message not found in cache, starting recursive query");
         // It will delete itself when it is done
         new RunningQuery(io, *question, answer_message, upstream_,
-                         upstream_root_, buffer, callback, query_timeout_,
+                         upstream_root_, test_server_,
+                         buffer, callback, query_timeout_,
                          client_timeout_, lookup_timeout_, retries_,
                          cache_);
     }
@@ -533,8 +582,9 @@ RecursiveQuery::resolve(const Question& question,
         dlog("Message not found in cache, starting recursive query");
         // It will delete itself when it is done
         new RunningQuery(io, question, answer_message, upstream_, upstream_root_,
-                             buffer, crs, query_timeout_, client_timeout_,
-                             lookup_timeout_, retries_, cache_);
+                         test_server_,
+                         buffer, crs, query_timeout_, client_timeout_,
+                         lookup_timeout_, retries_, cache_);
     }
 }
 

+ 14 - 0
src/lib/asiolink/recursive_query.h

@@ -98,12 +98,26 @@ public:
                  isc::dns::MessagePtr answer_message,
                  isc::dns::OutputBufferPtr buffer,
                  DNSServer* server);
+
+    /// \brief Set Test Server
+    ///
+    /// This method is *only* for unit testing the class.  If set, it enables
+    /// recursive behaviour but, regardless of responses received, sends every
+    /// query to the test server.
+    ///
+    /// The test server is enabled by setting a non-zero port number.
+    ///
+    /// \param address IP address of the test server.
+    /// \param port Port number of the test server
+    void setTestServer(const std::string& address, uint16_t port);
+    
 private:
     DNSService& dns_service_;
     boost::shared_ptr<std::vector<std::pair<std::string, uint16_t> > >
         upstream_;
     boost::shared_ptr<std::vector<std::pair<std::string, uint16_t> > >
         upstream_root_;
+    std::pair<std::string, uint16_t> test_server_;
     int query_timeout_;
     int client_timeout_;
     int lookup_timeout_;

+ 0 - 28
src/lib/asiolink/tests/io_fetch_unittest.cc

@@ -95,34 +95,6 @@ public:
         msg.toWire(renderer);
     }
 
-    /// \brief Read uint16_t from network-byte-order buffer
-    ///
-    /// Adapted from isc::dns::InputBuffer::readUint16().
-    ///
-    /// \param data Pointer to at least two bytes of data which are in network
-    ///        byte order.
-    ///
-    /// \return uint16_t value in host byte order.
-    uint16_t readUint16(const void* data) {
-        const uint8_t* cp = static_cast<const uint8_t*>(data);
-
-        uint16_t value = ((unsigned int)(cp[0])) << 8;
-        value |= ((unsigned int)(cp[1]));
-
-        return (value);
-    }
-
-    /// \brief Write uint16_t to network-byte-order buffer
-    ///
-    /// Adapted from isc::dns::OutputBuffer::writeUint16().
-    ///
-    /// \param value The 16-bit integer to be written into the buffer.
-    /// \param data Pointer to buffer at least two bytes long
-    void writeUint16(uint16_t value, uint8_t* data) {
-        data[0] = static_cast<uint8_t>((value & 0xff00U) >> 8);
-        data[1] = static_cast<uint8_t>(value & 0x00ffU);
-    }
-
     /// \brief UDP Response handler (the "remote UDP DNS server")
     ///
     /// When IOFetch is sending data, this response handler emulates the remote

+ 509 - 0
src/lib/asiolink/tests/recursive_query_unittest_2.cc

@@ -0,0 +1,509 @@
+// Copyright (C) 2011  Internet Systems Consortium, Inc. ("ISC")
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND ISC DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS.  IN NO EVENT SHALL ISC BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE
+// OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+#include <algorithm>
+#include <cstdlib>
+#include <string>
+#include <iostream>
+
+#include <gtest/gtest.h>
+#include <boost/bind.hpp>
+
+
+#include <asio.hpp>
+
+#include <dns/buffer.h>
+#include <dns/question.h>
+#include <dns/message.h>
+#include <dns/messagerenderer.h>
+#include <dns/opcode.h>
+#include <dns/name.h>
+#include <dns/rcode.h>
+
+#include <asiolink/asiolink_utilities.h>
+#include <asiolink/io_address.h>
+#include <asiolink/io_endpoint.h>
+#include <asiolink/io_fetch.h>
+#include <asiolink/io_service.h>
+
+using namespace asio;
+using namespace isc::dns;
+using namespace asio::ip;
+using namespace std;
+
+/// RecursiveQuery Test - 2
+///
+/// The second part of the RecursiveQuery unit tests, this attempts to get the
+/// RecursiveQuery object to follow a set of referrals for "www.example.org" to
+/// and to invoke TCP fallback on one of the queries.  In particular, we
+/// expect that the test will do the following in an attempt to resolve
+/// www.example.org:
+///
+/// - Send a question over UDP to the root servers - get referral to "org".
+/// - Send question over UDP to "org" - get referral to "example.org" with TC bit set.
+/// - Send question over TCP to "org" - get referral to "example.org".
+/// - Send question over UDP to "example.org" - get response for www.example.org.
+///
+/// The order of queries is partly to test that after there is a fallover to TCP,
+/// queries revert to UDP.
+///
+/// By using the "test_server_" element of RecursiveQuery, all queries are
+/// directed to one or other of the "servers" in the RecursiveQueryTest2 class.
+
+namespace asiolink {
+
+const asio::ip::address TEST_HOST(asio::ip::address::from_string("127.0.0.1"));
+const uint16_t TEST_PORT(5301);
+
+
+/// \brief Test fixture for the asiolink::IOFetch.
+class RecursiveQueryTest2 : public virtual ::testing::Test, public virtual IOFetch::Callback
+{
+public:
+
+    /// \brief Status of query
+    ///
+    /// Set before the query and then by each "server" when responding.
+    enum QueryStatus {
+        NONE = 0,                   ///< Default
+        UDP_ROOT = 1,               ///< Query root server over UDP
+        UDP_ORG = 2,                ///< Query ORG server over UDP
+        TCP_ORG = 3,                ///< Query ORG server over TCP
+        UDP_EXAMPLE_ORG = 4,        ///< Query EXAMPLE.ORG server over UDP
+        COMPLETE = 5                ///< Query is complete
+    };
+
+    IOService       service_;                   ///< Service to run everything
+    Question        question_;                  ///< What to ask
+    QueryStatus     last_;                      ///< Last state
+    QueryStatus     expected_;                  ///< Expected next state
+    OutputBufferPtr question_buffer_;           ///< Question we expect to receive
+
+    size_t          tcp_cumulative_;            ///< Cumulative TCP data received
+    tcp::endpoint   tcp_endpoint_;              ///< Endpoint for TCP receives
+    size_t          tcp_length_;                ///< Expected length value
+    uint8_t         tcp_receive_buffer_[512];   ///< Receive buffer for TCP I/O
+    OutputBufferPtr tcp_send_buffer_;           ///< Send buffer for TCP I/O
+    tcp::socket     tcp_socket_;                ///< Socket used by TCP server
+
+    /// Data for UDP
+    udp::endpoint   udp_endpoint_;              ///< Endpoint for UDP receives
+    size_t          udp_length_;                ///< Expected length value
+    uint8_t         udp_receive_buffer_[512];   ///< Receive buffer for UDP I/O
+    OutputBufferPtr udp_send_buffer_;           ///< Send buffer for UDP I/O
+    udp::socket     udp_socket_;                ///< Socket used by UDP server
+
+
+    /// \brief Constructor
+    RecursiveQueryTest2() :
+        service_(),
+        question_(Name("www.example.org"), RRClass::IN(), RRType::A()),
+        last_(NONE),
+        expected_(NONE),
+        question_buffer_(),
+        tcp_cumulative_(0),
+        tcp_endpoint_(TEST_HOST, TEST_PORT),
+        tcp_length_(0),
+        tcp_receive_buffer_(),
+        tcp_send_buffer_(),
+        tcp_socket_(service_.get_io_service()),
+        udp_endpoint_(),
+        udp_length_(0),
+        udp_receive_buffer_(),
+        udp_send_buffer_(),
+        udp_socket_(service_.get_io_service(), udp::v4())
+    {
+
+    }
+
+
+    /// \brief Set Common Message Bits
+    ///
+    /// Sets up the common bits of a response message returned by the handlers.
+    ///
+    /// \param msg Message buffer in RENDER mode.
+    /// \param qid QIT to set the message to
+    void setCommonMessage(isc::dns::Message& msg, uint16_t qid = 0) {
+        msg.setQid(qid);
+        msg.setHeaderFlag(Message::HEADERFLAG_QR);
+        msg.setOpcode(Opcode::QUERY());
+        msg.setHeaderFlag(Message::HEADERFLAG_AA);
+        msg.setRcode(Rcode::NOERROR());
+        msg.addQuestion(question_);
+    }
+
+    /// \brief Set Referral to "org"
+    ///
+    /// Sets up the passed-in message (expected to be in "RENDER" mode to
+    /// indicate a referral to fictitious .org nameservers.
+    ///
+    /// \param msg Message to update with referral information.
+    void setReferralOrg(isc::dns::Message& msg) {
+
+        // Do a referral to org.  We'll define all NS records as "in-zone"
+        // nameservers (and so supply glue) to avoid the possibility of
+        // the resolver doing another lookup.
+        RRSetPtr org_ns(new RRSet(Name("org."), RRClass::IN(), RRType::NS(), RRTTL(300)));
+        org_ns->addRdata(NS("ns1.org."));
+        org_ns->addRdata(NS("ns2.org."));
+        msg.addRRset(Message::SECTION_AUTHORITY, org_ns);
+
+        RRsetPtr org_ns1(new RRSet(Name("ns1.org."), RRClass::IN(), RRType::A(), RRTTL(300)));
+        org_ns1->addRdata(A("192.0.2.1"));
+        msg.addRRset(Message::SECTION_ADDITIONAL, org_ns1);
+
+        RRsetPtr org_ns1(new RRSet(Name("ns2.org."), RRClass::IN(), RRType::A(), RRTTL(300)));
+        org_ns2->addRdata(A("192.0.2.2"));
+        msg.addRRset(Message::SECTION_ADDITIONAL, org_ns2);
+    }
+
+    /// \brief Set Referral to "example.org"
+    ///
+    /// Sets up the passed-in message (expected to be in "RENDER" mode to
+    /// indicate a referral to fictitious example.org nameservers.
+    ///
+    /// \param msg Message to update with referral information.
+    void setReferralExampleOrg(isc::dns::Message& msg) {
+
+        // Do a referral to example.org.  As before, we'll define all NS
+        // records as "in-zone" nameservers (and so supply glue) to avoid
+        // the possibility of the resolver doing another lookup.
+        RRSetPtr example_org_ns(new RRSet(Name("example.org."), RRClass::IN(), RRType::NS(), RRTTL(300)));
+        example_org_ns->addRdata(NS("ns1.example.org."));
+        example_org_ns->addRdata(NS("ns2.example.org."));
+        msg.addRRset(Message::SECTION_AUTHORITY, example_org_ns);
+
+        RRsetPtr example_org_ns1(new RRSet(Name("ns1.example.org."), RRClass::IN(), RRType::A(), RRTTL(300)));
+        example_org_ns1->addRdata(A("192.0.2.11"));
+        msg.addRRset(Message::SECTION_ADDITIONAL, example_org_ns1);
+
+        RRsetPtr org_ns1(new RRSet(Name("ns2.example.org."), RRClass::IN(), RRType::A(), RRTTL(300)));
+        example_org_ns2->addRdata(A("192.0.2.12"));
+        msg.addRRset(Message::SECTION_ADDITIONAL, example_org_ns2);
+    }
+
+    /// \brief Set Answer to "www.example.org"
+    ///
+    /// Sets up the passed-in message (expected to be in "RENDER" mode to
+    /// indicate an authoritative answer to www.example.org.
+    ///
+    /// \param msg Message to update with referral information.
+    void setAnswerWwwExampleOrg(isc::dns::Message& msg) {
+
+        // Give a response for www.example.org.
+        RRsetPtr www_example_org_a(new RRSet(Name("www.example.org."), RRClass::IN(), RRType::NS(), RRTTL(300)));
+        www_example_org_a->addRdata(A("192.0.2.21"));
+        msg.addRRset(Message::SECTION_ANSWER, example_org_ns);
+
+        // ... and add the Authority and Additional sections. (These are the
+        // same as in the referral to example.org from the .org nameserver.)
+        setReferralExampleOrg(msg);
+    }
+
+    /// \brief UDP Receive Handler
+    ///
+    /// This is invoked when a message is received from the RecursiveQuery
+    /// Object.  It formats an answer and sends it, with the UdpSendHandler
+    /// method being specified as the completion handler.
+    ///
+    /// \param remote Endpoint to which to send the answer
+    /// \param socket Socket to use to send the answer
+    /// \param ec ASIO error code, completion code of asynchronous I/O issued
+    ///        by the "server" to receive data.
+    /// \param length Amount of data received.
+    void udpReceiveHandler(error_code ec = error_code(), size_t length = 0) {
+
+        // Expected state should be one greater than the last state.
+        EXPECT_EQ(static_cast<int>(expected_), static_cast<int>(last_) + 1);
+        last_ = expected_;
+
+        // The QID in the incoming data is random so set it to 0 for the
+        // data comparison check. (It is set to 0 in the buffer containing
+        // the expected data.)
+        uint16_t qid = readUint16(udp_receive_buffer_);
+        udp_receive_buffer_[0] = udp_receive_buffer_[1] = 0;
+
+        // Check that length of the received data and the expected data are
+        // identical, then check that the data is identical as well.
+        EXPECT_EQ(question_buff_->getLength(), length);
+        EXPECT_TRUE(equal(udp_receive_buffer_, (udp_receive_buffer_ + length - 1),
+                    static_cast<const uint8_t*>(question_buff_->getData())));
+
+        // The message returned depends on what state we are in.  Set up
+        // common stuff first: bits not mentioned are set to 0.
+        Message msg(Message::RENDER);
+        setCommonMessage(msg, qid);
+
+        // Set up state-dependent bits:
+        switch (expected_) {
+        case UDP_ROOT:
+            // Return a referral to org.  We then expect to query the "org"
+            // nameservers over UDP next.
+            setReferralOrg(msg);
+            expected_ = UDP_ORG;
+            break;
+
+         case UDP_ORG:
+            // Return a referral to example.org.  We explicitly set the TC bit to
+            // force a repeat query to the .org nameservers over TCP.
+            setReferralExampleOrg(msg);
+            msg.setHeaderFlag(Message::HEADERFLAG_TC);
+            expected_ = TCP_ORG;
+            break;
+
+         case UDP_EXAMPLE_ORG:
+            // Return the answer to the question.
+            setAnswerWwwExampleOrg(msg);
+            expected_ = COMPLETE;
+            break;
+
+         default:
+            FAIL() << "UdpReceiveHandler called with unknown state";
+        }
+
+        // Convert to wire format
+        MessageRenderer renderer(*udp_send_buffer_);
+        msg.toWire(renderer);
+
+        // Return a message back to the IOFetch object.
+        udp_socket_.send_to(asio::buffer(udp_send_buffer_.getData(),
+                                         udp_send_buffer_.getLength()),
+                            boost::bind(&RecursiveQueryTest2::UdpSendHandler,
+                                        this, _1, _2));
+
+        // Set the expected length for the send handler.
+        udp_length_ = udp_send_buffer_.getLength();
+    }
+
+    /// \brief UDP Send Handler
+    ///
+    /// Called when a send operation of the UDP server (i.e. a response
+    /// being sent to the RecursiveQuery) has completed, this re-issues
+    /// a read call.
+    void udpSendHandler(error_code ec = error_code(), size_t length = 0) {
+
+        // Check send was OK
+        EXPECT_EQ(0, ec.value());
+        EXPECT_EQ(udp_length_, length);
+
+        // Reissue the receive.
+        udp_socket_.async_receive_from(
+            asio::buffer(udp_receive_buffer_, sizeof(udp_receive_buffer_)),
+                udp_endpoint_
+                boost::bind(&recursiveQuery2::udpReceiveHandler, this, _1, _2));
+    }
+
+    /// \brief Completion Handler for Accepting TCP Data
+    ///
+    /// Called when the remote system connects to the "server".  It issues
+    /// an asynchronous read on the socket to read data.
+    ///
+    /// \param socket Socket on which data will be received
+    /// \param ec Boost error code, value should be zero.
+    void tcpAcceptHandler(error_code ec = error_code(), size_t length = 0)
+    {
+        // Expect that the accept completed without a problem.
+        EXPECT_EQ(0, ec.value());
+
+        // Initiate a read on the socket, indicating that nothing has yet been
+        // received.
+        tcp_cumulative_ = 0;
+        tcp_socket_.async_receive(
+            asio::buffer(tcp_receive_buffer_, sizeof(tcp_receive_buffer_)),
+            boost::bind(&recursiveQuery2::tcpReceiveHandler, this, _1, _2));
+    }
+
+    /// \brief Completion Handler for Receiving TCP Data
+    ///
+    /// Reads data from the RecursiveQuery object and loops, reissuing reads,
+    /// until all the message has been read.  It then sends
+    ///
+    /// \param socket Socket to use to send the answer
+    /// \param ec ASIO error code, completion code of asynchronous I/O issued
+    ///        by the "server" to receive data.
+    /// \param length Amount of data received.
+    void tcpReceiveHandler(error_code ec = error_code(), size_t length = 0)
+    {
+        // Expect that the receive completed without a problem.
+        EXPECT_EQ(0, ec.value());
+
+        // Have we received all the data?  We know this by checking if the two-
+        // byte length count in the message is equal to the data received.
+        tcp_cumulative_ += length;
+        bool complete = false;
+        if (tcp_cumulative_ > 2) {
+            uint16_t dns_length = readUint16(tcp_receive_buffer_);
+            complete = ((dns_length + 2) == tcp_cumulative_);
+        }
+
+        if (!complete) {
+
+            // Not complete yet, issue another read.
+            tcp_socket_.async_receive(
+                asio::buffer(tcp_receive_buffer_ + tcp_cumulative_,
+                             sizeof(tcp_receive_buffer_) - tcp_cumulative_),
+                boost::bind(&recursiveQuery2::tcpReceiveHandler, this, _1, _2));
+            return;
+        }
+
+        // Have received a TCP message.  Expected state should be one greater
+        // than the last state.
+        EXPECT_EQ(static_cast<int>(expected_), static_cast<int>(last_) + 1);
+        last_ = expected_;
+
+        // Check that length of the received data and the expected data are
+        // identical (taking into account the two-byte count), then check that
+        // the data is identical as well (after zeroing the QID).
+        EXPECT_EQ(question_buff_->getLength() + 2, tcp_cumulative_);
+        uint16_t qid = readUint16(&udp_receive_buffer_[2]);
+        tcp_receive_buffer_[2] = tcp_receive_buffer_[3] = 0;
+        EXPECT_TRUE(equal((tcp_receive_buffer_ + 2),
+                          (tcp_receive_buffer_ + tcp_cumulative_),
+                           static_cast<const uint8_t*>(question_buff_->getData())));
+
+
+        // Return a message back.  This is a referral to example.org, which
+        // should result in another query over UDP.
+        Message msg(Message::RENDER);
+        setCommonMessage(msg, qid);
+        setReferralExampleOrg(msg);
+
+        // Convert to wire format
+        MessageRenderer renderer(*tcp_send_buffer_);
+        msg.toWire(renderer);
+        
+        // Expected next state (when checked) is the UDP query to example.org.
+        expected_ = UDP_EXAMPLE_ORG;
+        
+        // We'll write the message in two parts, the count and the message
+        // itself.  When specifying the send handler, the expected size of the
+        // data written is passed as the first parameter.
+        uint8_t count[2];
+        writeUint16(tcp_send_buffer_->getLength(), count);
+        socket->async_send(asio::buffer(count, 2),
+                           boost::bind(&IOFetchTest::tcpSendHandler, this,
+                                       2, _1, _2));
+        socket->async_send(asio::buffer(tcp_send_buffer_->getData(),
+                                        tcp_send_buffer_->getLength()),
+                           boost::bind(&RecursiveQuery2::tcpSendHandler, this,
+                                       sizeof(TEST_DATA), _1, _2));
+    }
+
+    /// \brief Completion Handler for Sending TCP data
+    ///
+    /// Called when the asynchronous send of data back to the RecursiveQuery
+    /// by the TCP "server" in this class has completed.  (This send has to
+    /// be asynchronous because control needs to return to the caller in order
+    /// for the IOService "run()" method to be called to run the handlers.)
+    ///
+    /// \param expected Number of bytes that were expected to have been sent.
+    /// \param ec Boost error code, value should be zero.
+    /// \param length Number of bytes sent.
+    void tcpSendHandler(size_t expected = 0, error_code ec = error_code(),
+                        size_t length = 0)
+    {
+        EXPECT_EQ(0, ec.value());       // Expect no error
+        EXPECT_EQ(expected, length);    // And that amount sent is as expected
+    }
+
+    /// \brief Resolver Callback Completion
+    ///
+    /// This is the callback's operator() method which is called when the 
+    /// resolution of the query is complete.  It checks that the data received
+    /// is the wire format of the data sent back by the server.
+    ///
+    /// \param result Result indicated by the callback
+    void operator()(IOFetch::Result result) {
+
+        EXPECT_EQ(COMPLETE, expected_);
+        /*
+        EXPECT_EQ(expected_, result);   // Check correct result returned
+        EXPECT_FALSE(run_);             // Check it is run only once
+        run_ = true;                    // Note success
+
+        // If the expected result for SUCCESS, then this should have been called
+        // when one of the "servers" in this class has sent back the TEST_DATA.
+        // Check the data is as expected/
+        if (expected_ == IOFetch::SUCCESS) {
+            EXPECT_EQ(sizeof(TEST_DATA), result_buff_->getLength());
+
+            const uint8_t* start = static_cast<const uint8_t*>(result_buff_->getData());
+            EXPECT_TRUE(equal(TEST_DATA, (TEST_DATA + sizeof(TEST_DATA) - 1),
+                              start));
+        }
+
+        // ... and cause the run loop to exit.
+         * */
+        service_.stop();
+    }
+
+// Sets up the UDP and TCP "servers", then tries a resolution.
+
+TEST_F(RecursiveQueryTest2, Resolve) {
+
+    // Set the state of the 
+    // Set up the UDP server and issue the first read.
+    udp_socket_.set_option(socket_base::reuse_address(true));
+    udp_socket_.bind(udp::endpoint(TEST_HOST, TEST_PORT));
+    udp_socket.async_receive_from(asio::buffer(server_buff_, sizeof(server_buff_)),
+        remote,
+        boost::bind(&IOFetchTest::udpReceiveHandler, this, &remote, &socket,
+                    _1, _2));
+    service_.get_io_service().post(udp_fetch_);
+    service_.run();
+
+    socket.close();
+
+    EXPECT_TRUE(run_);;
+}
+
+// Do the same tests for TCP transport
+
+TEST_F(IOFetchTest, TcpStop) {
+    stopTest(IOFetch::TCP, tcp_fetch_);
+}
+
+TEST_F(IOFetchTest, TcpPrematureStop) {
+    prematureStopTest(IOFetch::TCP, tcp_fetch_);
+}
+
+TEST_F(IOFetchTest, TcpTimeout) {
+    timeoutTest(IOFetch::TCP, tcp_fetch_);
+}
+
+TEST_F(IOFetchTest, TcpSendReceive) {
+    protocol_ = IOFetch::TCP;
+    expected_ = IOFetch::SUCCESS;
+
+    // Socket into which the connection will be accepted
+    tcp::socket socket(service_.get_io_service());
+
+    // Acceptor object - called when the connection is made, the handler will
+    // initiate a read on the socket.
+    tcp::acceptor acceptor(service_.get_io_service(),
+                           tcp::endpoint(tcp::v4(), TEST_PORT));
+    acceptor.async_accept(socket,
+        boost::bind(&IOFetchTest::tcpAcceptHandler, this, &socket, _1));
+
+    // Post the TCP fetch object to send the query and receive the response.
+    service_.get_io_service().post(tcp_fetch_);
+
+    // ... and execute all the callbacks.  This exits when the fetch completes.
+    service_.run();
+    EXPECT_TRUE(run_);  // Make sure the callback did execute
+
+    socket.close();
+}
+
+} // namespace asiolink