Merge pull request #208 from kurdybacha/poller-fix

Problem: poller_t invalid behavior on moved sockets

Author: sigiesec

tests/poller.cpp

@@ -228,4 +228,28 @@ TEST(poller, client_server)
     ASSERT_TRUE(got_pollout);
 }
 
+TEST(poller, poller_add_invalid_socket_throws)
+{
+    zmq::context_t context;
+    zmq::poller_t poller;
+    zmq::socket_t a {context, zmq::socket_type::router};
+    zmq::socket_t b {std::move (a)};
+    ASSERT_THROW (poller.add (a, ZMQ_POLLIN, zmq::poller_t::handler_t {}),
+                  zmq::error_t);
+    ASSERT_EQ (0u, poller.size ());
+}
+
+TEST(poller, poller_remove_invalid_socket_throws)
+{
+    zmq::context_t context;
+    zmq::socket_t socket {context, zmq::socket_type::router};
+    zmq::poller_t poller;
+    ASSERT_NO_THROW (poller.add (socket, ZMQ_POLLIN, zmq::poller_t::handler_t {}));
+    ASSERT_EQ (1u, poller.size ());
+    std::vector<zmq::socket_t> sockets;
+    sockets.emplace_back (std::move (socket));
+    ASSERT_THROW (poller.remove (socket), zmq::error_t);
+    ASSERT_EQ (1u, poller.size ());
+}
+
 #endif

zmq.hpp

@@ -536,6 +536,11 @@ namespace zmq
         {
             return ptr;
         }
+
+        inline operator bool() const ZMQ_NOTHROW
+        {
+            return ptr != NULL;
+        }
     private:
 
         void *ptr;
@@ -1053,26 +1058,24 @@ namespace zmq
 
         void add (zmq::socket_t &socket, short events, handler_t handler)
         {
-            handler_t *handler_ptr = nullptr;
-            /// \todo is it sensible to allow handler to be empty? doesn't this lead to an error when the event is signalled? (perhaps it should already lead to an error in zmq_poller_add then)
-            if (handler) {
-                auto emplace_res = handlers.emplace(&socket, std::move(handler));
-                handler_ptr = &emplace_res.first->second;
-            }
-            if (0 == zmq_poller_add (poller_ptr, socket.ptr, handler_ptr, events)) {
+            auto it = std::end (handlers);
+            auto inserted = false;
+            if (handler)
+                std::tie(it, inserted) = handlers.emplace (socket.ptr, std::move (handler));
+            if (0 == zmq_poller_add (poller_ptr, socket.ptr, inserted ? &(it->second) : nullptr, events)) {
                 poller_events.emplace_back (zmq_poller_event_t ());
                 return;
             }
+            // rollback
+            if (inserted)
+                handlers.erase (socket.ptr);
             throw error_t ();
         }
 
         void remove (zmq::socket_t &socket)
         {
             if (0 == zmq_poller_remove (poller_ptr, socket.ptr)) {
-                auto it = handlers.find (&socket);
-                if (it != handlers.end ()) { /// \todo this may only be false if handler was empty on add
-                    handlers.erase (it);
-                }
+                handlers.erase (socket.ptr);
                 poller_events.pop_back ();
                 return;
             }
@@ -1108,7 +1111,7 @@ namespace zmq
     private:
         void *poller_ptr;
         std::vector<zmq_poller_event_t> poller_events;
-        std::unordered_map<socket_t*, handler_t> handlers;
+        std::unordered_map<void*, handler_t> handlers;
     };  // class poller_t
 #endif //  defined(ZMQ_BUILD_DRAFT_API) && defined(ZMQ_CPP11) && defined(ZMQ_HAVE_POLLER)