ZMQII-34: ensure that poll won't return POLLIN event when the message will be filtered out anyway (POSIX)

4 files changed, 110 insertions, 46 deletions

diff --git a/src/socket_base.cpp b/src/socket_base.cpp
index 3e6488d..ad68fdb 100644
--- a/src/socket_base.cpp
+++ b/src/socket_base.cpp
@@ -331,9 +331,7 @@ int zmq::socket_base_t::recv (::zmq_msg_t *msg_, int flags_)
     }
     else  {
         while (rc != 0) {
-            if (errno == EINPROGRESS)
-                app_thread->process_commands (false, true);
-            else if (errno == EAGAIN)
+            if (errno == EAGAIN)
                 app_thread->process_commands (true, false);
             else
                 return -1;
diff --git a/src/sub.cpp b/src/sub.cpp
index bdc27da..b8e3990 100644
--- a/src/sub.cpp
+++ b/src/sub.cpp
@@ -26,14 +26,17 @@
 
 zmq::sub_t::sub_t (class app_thread_t *parent_) :
     socket_base_t (parent_),
-    all_count (0)
+    all_count (0),
+    has_message (false)
 {
     options.requires_in = true;
     options.requires_out = false;
+    zmq_msg_init (&message);
 }
 
 zmq::sub_t::~sub_t ()
 {
+    zmq_msg_close (&message);
 }
 
 void zmq::sub_t::xattach_pipes (class reader_t *inpipe_,
@@ -115,16 +118,72 @@ int zmq::sub_t::xflush ()
 
 int zmq::sub_t::xrecv (zmq_msg_t *msg_, int flags_)
 {
-    //  Get a message using fair queueing algorithm.
-    int rc = fq.recv (msg_, flags_);
+    //  If there's already a message prepared by a previous call to zmq_poll,
+    //  return it straight ahead.
+    if (has_message) {
+        zmq_msg_move (msg_, &message);
+        has_message = false;
+        return 0;
+    }
+
+    //  TODO: This can result in infinite loop in the case of continuous
+    //  stream of non-matching messages which breaks the non-blocking recv
+    //  semantics.
+    while (true) {
 
-    //  If there's no message available, return immediately.
-    if (rc != 0 && errno == EAGAIN)
-        return -1;
+        //  Get a message using fair queueing algorithm.
+        int rc = fq.recv (msg_, flags_);
 
+        //  If there's no message available, return immediately.
+        //  The same when error occurs.
+        if (rc != 0)
+            return -1;
+
+        //  Check whether the message matches at least one subscription.
+        if (match (msg_))
+            return 0;
+    }
+}
+
+bool zmq::sub_t::xhas_in ()
+{
+    //  If there's already a message prepared by a previous call to zmq_poll,
+    //  return straight ahead.
+    if (has_message)
+        return true;
+
+    //  TODO: This can result in infinite loop in the case of continuous
+    //  stream of non-matching messages.
+    while (true) {
+
+        //  Get a message using fair queueing algorithm.
+        int rc = fq.recv (&message, ZMQ_NOBLOCK);
+
+        //  If there's no message available, return immediately.
+        //  The same when error occurs.
+        if (rc != 0) {
+            zmq_assert (errno == EAGAIN);
+            return false;
+        }
+
+        //  Check whether the message matches at least one subscription.
+        if (match (&message)) {
+            has_message = true;
+            return true;
+        }
+    }
+}
+
+bool zmq::sub_t::xhas_out ()
+{
+    return false;
+}
+
+bool zmq::sub_t::match (zmq_msg_t *msg_)
+{
     //  If there is at least one * subscription, the message matches.
     if (all_count)
-        return 0;
+        return true;
 
     //  Check whether the message matches at least one prefix subscription.
     //  TODO: Make this efficient - O(log(n)) where n is number of characters in
@@ -135,25 +194,8 @@ int zmq::sub_t::xrecv (zmq_msg_t *msg_, int flags_)
         size_t sub_size = it->size ();
         if (sub_size <= msg_size &&
               memcmp (zmq_msg_data (msg_), it->data (), sub_size) == 0)
-            return 0;
+            return true;
     }
 
-    //  The message did not pass the filter. Trim it.
-    //  Note that we are returning a different error code so that the caller
-    //  knows there are more messages available. We cannot loop here as
-    //  a stream of non-matching messages would create a DoS situation.
-    zmq_msg_close (msg_);
-    zmq_msg_init (msg_);
-    errno = EINPROGRESS;
-    return -1;
-}
-
-bool zmq::sub_t::xhas_in ()
-{
-    return fq.has_in ();
-}
-
-bool zmq::sub_t::xhas_out ()
-{
     return false;
 }
diff --git a/src/sub.hpp b/src/sub.hpp
index df7d369..2044442 100644
--- a/src/sub.hpp
+++ b/src/sub.hpp
@@ -23,6 +23,8 @@
 #include <set>
 #include <string>
 
+#include "../bindings/c/zmq.h"
+
 #include "socket_base.hpp"
 #include "fq.hpp"
 
@@ -53,6 +55,9 @@ namespace zmq
 
     private:
 
+        //  Check whether the message matches at least one subscription.
+        bool match (zmq_msg_t *msg_);
+
         //  Fair queueing object for inbound pipes.
          fq_t fq;
 
@@ -62,6 +67,11 @@ namespace zmq
         typedef std::multiset <std::string> subscriptions_t;
         subscriptions_t subscriptions;
 
+        //  If true, 'message' contains a matching message to return on the
+        //  next recv call.
+        bool has_message;
+        zmq_msg_t message;
+
         sub_t (const sub_t&);
         void operator = (const sub_t&);
     };
diff --git a/src/zmq.cpp b/src/zmq.cpp
index 03064bc..1ed96df 100644
--- a/src/zmq.cpp
+++ b/src/zmq.cpp
@@ -324,27 +324,20 @@ int zmq_poll (zmq_pollitem_t *items_, int nitems_, long timeout_)
         npollfds++;
     }
 
+    //  First iteration just check for events, don't block. Waiting would
+    //  prevent exiting on any events that may already been signaled on
+    //  0MQ sockets.
+    int rc = poll (pollfds, npollfds, 0);
+    if (rc == -1 && errno == EINTR) {
+        free (pollfds);
+        return 0;
+    }
+    errno_assert (rc >= 0);
+
     int timeout = timeout_ > 0 ? timeout_ / 1000 : -1;
     int nevents = 0;
-    bool initial = true;
-    while (!nevents) {
-
-        //  Wait for activity. In the first iteration just check for events,
-        //  don't wait. Waiting would prevent exiting on any events that may
-        //  already be signaled on 0MQ sockets.
-        int rc = poll (pollfds, npollfds, initial ? 0 : timeout);
-        if (rc == -1 && errno == EINTR)
-            break;
-        errno_assert (rc >= 0);
-
-        //  If timeout was hit with no events signaled, return zero.
-        if (!initial && rc == 0) {
-            free (pollfds);
-            return 0;
-        }
 
-        //  From now on, perform blocking polling.
-        initial = false;
+    while (true) {
 
         //  Process 0MQ commands if needed.
         if (nsockets && pollfds [npollfds -1].revents & POLLIN)
@@ -376,6 +369,27 @@ int zmq_poll (zmq_pollitem_t *items_, int nitems_, long timeout_)
             if (items_ [i].revents)
                 nevents++;
         }
+
+        //  If there's at least one event, or if we are asked not to block,
+        //  return immediately.
+        if (nevents || !timeout)
+            break;
+
+        //  Wait for events.
+        rc = poll (pollfds, npollfds, timeout);
+        if (rc == -1 && errno == EINTR)
+            break;
+        errno_assert (rc >= 0);
+
+        //  If timeout was hit with no events signaled, return zero.
+        if (rc == 0)
+            break;
+
+        //  If timeout was already applied, we don't want to poll anymore.
+        //  Setting timeout to zero will cause termination of the function
+        //  once the events we've got are processed.
+        if (timeout > 0)
+            timeout = 0;
     }
 
     free (pollfds);