Merge pull request #28 from CopernicaMarketingSoftware/buffer-sizes

Prevent one buffer copy

Author: EmielBruijntjes

include/dnscpp/message.h

@@ -168,7 +168,7 @@ class Message
      *  Is the message truncated? In that case it is better to use tcp
      *  @return bool
      */
-    bool truncated() const { return flag(ns_f_tc); }
+    bool truncated() const { /*if (rand() % 2 == 0) return true; else*/ return flag(ns_f_tc); }
 
     /**
      *  Is recursion desired (false for responses)

include/dnscpp/query.h

@@ -24,6 +24,12 @@
  */
 namespace DNS {
 
+// we advertise that we support 1200 bytes for our response buffer size,
+// this is the same buffer size as libresolv seems to use. Their ratio
+// is that this limits the risk that dgram message get fragmented,
+// which makes the system vulnerable for injection
+constexpr size_t EDNSPacketSize = 1200;
+
 /**
  *  Forward declarations
  */

include/dnscpp/socket.h

@@ -56,7 +56,7 @@ class Socket : public Inbound, public Watchable
      *  All the buffered responses that came in
      *  @var std::list
      */
-    std::list<std::pair<Ip,std::basic_string<unsigned char>>> _responses;
+    std::list<std::pair<Ip,std::vector<unsigned char>>> _responses;
 
 protected:
     /**
@@ -74,10 +74,9 @@ class Socket : public Inbound, public Watchable
      *  Add a message for delayed processing
      *  @param  addr    the address from which the message was received
      *  @param  buffer  the response buffer
-     *  @param  size    size of the buffer
      */
-    void add(const sockaddr *addr, const unsigned char *buffer, size_t size);
-    void add(const Ip &addr, const unsigned char *buffer, size_t size);
+    void add(const sockaddr *addr, std::vector<unsigned char> &&buffer);
+    void add(const Ip &addr, std::vector<unsigned char> &&buffer);
 
 public:
     /**

include/dnscpp/tcp.h

@@ -71,19 +71,25 @@ class Tcp : public Socket, private Monitor
      *  @var int
      */
     int _fd;
-    
+
+    /**
+     *  Size of the buffer to fill
+     *  @var uint16_t
+     */
+    uint16_t _size;
+
     /**
-     *  The buffer that is being filled right now (first two bytes contain the size)
-     *  @var unsigned char *
+     *  The response that is being filled right now
+     *  @var vector
      */
     std::vector<unsigned char> _buffer;
 
     /**
-     *  How far is the _buffer now filled?
+     *  How many bytes of the frame were transferred?
      *  @var size_t
      */
-    size_t _filled = 0;
-    
+    size_t _transferred = 0;
+
     /**
      *  Identifier user for monitoring the filedescriptor in the event loop
      *  @var void *
@@ -140,23 +146,18 @@ class Tcp : public Socket, private Monitor
      */
     virtual void notify() override;
 
-    /**
-     *  Reallocate the buffer if it turns out that our buffer is too small for the expected response
-     *  @return bool
-     */
-    bool reallocate();
-
-    /**
-     *  Size of the response -- this method only works if we have already received the frist two bytes
-     *  @return uint16_t
-     */
-    uint16_t responsesize() const;
-    
     /**
      *  Number of bytes that we expect in the next read operation
      *  @return size_t
      */
     size_t expected() const;
+
+    /**
+     *  Check return value of a recv syscall
+     *  @param  bytes  The bytes transferred
+     *  @return true if we should leap out (an error occurred or we'd block), false if not
+     */
+    bool updatetransferred(ssize_t bytes);
 
     /**
      *  Method that is called when there are no more subscribers, and that 

src/query.cpp

@@ -194,12 +194,9 @@ bool Query::edns(bool dnssec)
 
     // the type of the pseudo-record is "opt"
     put16(TYPE_OPT);
-    
-    // we advertise that we support 1200 bytes for our response buffer size, 
-    // this is the same buffer size as libresolv seems to use. Their ratio
-    // is that this limits the risk that dgram message get fragmented,
-    // which makes the system vulnerable for injection
-    put16(1200);
+
+    // tell the server that we support larger UDP packets
+    put16(EDNSPacketSize);
 
     // extended rcode (0 because the normal rcode is good enough) and the 
     // edns version (also 0 because that is the mose up-to-date version)

src/socket.cpp

@@ -24,18 +24,16 @@ namespace DNS {
  *  Add a message for delayed processing
  *  @param  addr    the address from which the message was received
  *  @param  buffer  the response buffer
- *  @param  size    size of the buffer
  */
-void Socket::add(const sockaddr *addr, const unsigned char *buffer, size_t size)
+void Socket::add(const sockaddr *addr, std::vector<unsigned char> &&buffer)
 {
     // avoid exceptions (in case the ip cannot be parsed)
     try
     {
         // @todo inbound messages that do not come from port 53 can be ignored
 
         // remember the response for now
-        // @todo make this more efficient (without all the string-copying)
-        _responses.emplace_back(addr, std::basic_string<unsigned char>(buffer, size));
+        _responses.emplace_back(addr, move(buffer));
     }
     catch (...)
     {
@@ -50,12 +48,11 @@ void Socket::add(const sockaddr *addr, const unsigned char *buffer, size_t size)
  *  Add a message for delayed processing
  *  @param  addr    the address from which the message was received
  *  @param  buffer  the response buffer
- *  @param  size    size of the buffer
  */
-void Socket::add(const Ip &addr, const unsigned char *buffer, size_t size)
+void Socket::add(const Ip &addr, std::vector<unsigned char> &&buffer)
 {
     // add to the list
-    _responses.emplace_back(addr, std::basic_string<unsigned char>(buffer, size));
+    _responses.emplace_back(addr, move(buffer));
 
     // reschedule the processing of messages
     _handler->onBuffered(this);

src/tcp.cpp

@@ -174,52 +174,20 @@ Inbound *Tcp::send(const Query &query)
     return result2 >= (ssize_t)query.size() ? this : nullptr;
 }
 
-/**
- *  Size of the response -- this method only works if we have already received the frist two bytes
- *  @return uint16_t
- */
-uint16_t Tcp::responsesize() const
-{
-    // result variable
-    uint16_t result;
-
-    // get the first two bytes from the buffer
-    memcpy(&result, _buffer.data(), 2);
-    
-    // put the bytes in the right order
-    return ntohs(result);
-}
-
 /**
  *  Number of bytes that we expect in the next read operation
  *  @return size_t
  */
 size_t Tcp::expected() const
 {
     // if we have not yet received the first two bytes, we expect those first
-    switch (_filled) {
-    case 0:     return 2;
-    case 1:     return 1;
-    default:    return responsesize() - (_filled - 2);
+    switch (_transferred) {
+    case 0:     return sizeof(uint16_t);
+    case 1:     return sizeof(uint16_t) - 1;
+    default:    return _size - (_transferred - sizeof(uint16_t));
     }
 }
 
-/**
- *  Reallocate the buffer if it turns out that our buffer is too small for the expected response
- *  @return bool
- */
-bool Tcp::reallocate()
-{
-    // preferred buffer size
-    size_t preferred = responsesize() + 2;
-    
-    // reallocate the buffer (but do not shrink)
-    _buffer.resize(std::max(_buffer.size(), preferred));
-    
-    // report result
-    return true;
-}
-
 /**
  *  Upgrate the socket from a _connecting_ socket to a _connected_ socket
  */
@@ -230,10 +198,7 @@ void Tcp::upgrade()
 
     // if the connection failed
     if (!_connected) return fail();
-    
-    // already allocate enough data for the first two bytes (holding the size of a response)
-    _buffer.resize(2);
-    
+
     // we no longer monitor for writability, but for readability instead
     _loop->update(_identifier, _fd, 1, this);
 
@@ -268,6 +233,26 @@ void Tcp::fail()
     _handler->onBuffered(this);
 }
 
+/**
+ *  Check return value of a recv syscall
+ *  @param  bytes  The bytes transferred
+ *  @return true if we should leap out (an error occurred), false if not
+ */
+bool Tcp::updatetransferred(ssize_t result)
+{
+    // the operation would block, but don't leap out
+    if (result < 0 && (errno == EAGAIN || errno == EWOULDBLOCK)) return false;
+
+    // if there is a failure we leap out as well
+    if (result <= 0) return fail(), true;
+
+    // update the number of transferred bytes
+    _transferred += result;
+
+    // don't leap out
+    return false;
+}
+
 /**
  *  Method that is called when the socket becomes active (readable in our case)
  */
@@ -279,30 +264,45 @@ void Tcp::notify()
 
     // if the socket is not yet connected, it might be connected right now
     if (!_connected) return upgrade();
-    
-    // receive data from the socket
-    auto result = ::recv(_fd, _buffer.data() + _filled, expected(), MSG_DONTWAIT);
-    
-    // do nothing if the operation is blocking
-    if (result < 0 && (errno == EAGAIN || errno == EWOULDBLOCK)) return;
-    
-    // if there is a failure we leap out
-    if (result <= 0) return fail();
-    
-    // update the number of bytes received
-    _filled += result;
-    
-    // after we've received the first two bytes, we can reallocate the buffer so that it is of sufficient size
-    if (_filled == 2 && !reallocate()) return fail();
-    
+
+    // We can be in two receive states: the first state is that we're waiting for the
+    // size of the buffer. The second state is that we are waiting for the response content itself.
+    // To determine in what state we're in, we can check how many bytes have been transferred.
+    // If that's less than 2 then we're still waiting for the response size.
+    if (_transferred < sizeof(uint16_t))
+    {
+        const auto result = ::recv(_fd, (uint8_t*)&_size + _transferred, sizeof(uint16_t) - _transferred, MSG_DONTWAIT);
+
+        // if there is a failure we leap out
+        if (updatetransferred(result)) return;
+
+        // if we still haven't received the two bytes we should leap out here
+        else if (_transferred < sizeof(uint16_t)) return;
+
+        // OK: the size of the rest of the frame was received, we know how much to allocate
+        // update the size
+        _size = ntohs(_size);
+
+        // size the buffer accordingly
+        _buffer.resize(_size);
+    }
+
+    // calculate offset into the buffer
+    size_t offset = _transferred - sizeof(uint16_t);
+
+    // This is the second state of the Tcp state machine. At this point we know we have
+    // received at least two bytes of the frame, and so we know we have resized the
+    // buffer accordingly. All that's left to do is to await the full response content
+    if (updatetransferred(::recv(_fd, _buffer.data() + offset, _buffer.size() - offset, MSG_DONTWAIT))) return;
+
     // continue waiting if we have not yet received everything there is
     if (expected() > 0) return;
-    
-    // all data has been received, buffer the response for now
-    add(_ip, _buffer.data() + 2, _filled - 2);
+
+    // all data has been received, we can move the response content into a deferred list to be processed later
+    add(_ip, move(_buffer));
 
     // for the next response we empty the buffer again
-    _filled = 0;
+    _transferred = 0;
 }
 
 /**

src/udp.cpp

@@ -5,6 +5,7 @@
 #include "../include/dnscpp/processor.h"
 #include "../include/dnscpp/query.h"
 #include <unistd.h>
+#include <cassert>
 
 /**
  *  Begin namespace
@@ -168,25 +169,29 @@ void Udp::notify()
     // do nothing if there is no socket (how is that possible!?)
     if (!valid()) return;
 
-    // the buffer to receive the response in
-    // @todo use a macro
-    unsigned char buffer[65536];
-
     // structure will hold the source address (we use an ipv6 struct because that is also big enough for ipv4)
     struct sockaddr_in6 from; socklen_t fromlen = sizeof(from);
 
-    // we want to get as much messages at onces as possible, but not run forever
-    // @todo use scatter-gather io to optimize this further
-    for (size_t messages = 0; messages < 1024; ++messages)
+    // the buffer to receive the response in
+    std::vector<unsigned char> buffer;
+
+    // read all messages until depleted
+    while (true)
     {
-        // reveive the message (the DONTWAIT option is needed because this is a blocking socket, but we dont want to block now)
-        auto bytes = recvfrom(_fd, buffer, sizeof(buffer), MSG_DONTWAIT, (struct sockaddr *)&from, &fromlen);
+        // make the buffer large enough
+        buffer.resize(EDNSPacketSize);
+
+        // receive the message (the DONTWAIT option is needed because this is a blocking socket, but we dont want to block now)
+        auto bytes = recvfrom(_fd, buffer.data(), buffer.size(), MSG_DONTWAIT, (struct sockaddr *)&from, &fromlen);
 
         // if there were no bytes, leap out
         if (bytes <= 0) break;
-        
+
+        // shrink to fit
+        buffer.resize(bytes);
+
         // add to the buffer
-        add((const sockaddr *)&from, buffer, bytes);
+        add((const sockaddr *)&from, move(buffer));
     }
 }