Message validity is checked in the runtime

Signed-off-by: Martin Sustrik <sustrik@250bpm.com>

4 files changed, 48 insertions, 10 deletions

diff --git a/src/msg.cpp b/src/msg.cpp
index e9d1da7..e800bd6 100644
--- a/src/msg.cpp
+++ b/src/msg.cpp
@@ -26,12 +26,13 @@
 #include <new>
 
 #include "stdint.hpp"
+#include "likely.hpp"
 #include "err.hpp"
 
 int zmq_msg_init (zmq_msg_t *msg_)
 {
     msg_->content = (zmq::msg_content_t*) ZMQ_VSM;
-    msg_->flags = 0;
+    msg_->flags = (unsigned char) ~ZMQ_MSG_MASK;
     msg_->vsm_size = 0;
     return 0;
 }
@@ -40,7 +41,7 @@ int zmq_msg_init_size (zmq_msg_t *msg_, size_t size_)
 {
     if (size_ <= ZMQ_MAX_VSM_SIZE) {
         msg_->content = (zmq::msg_content_t*) ZMQ_VSM;
-        msg_->flags = 0;
+        msg_->flags = (unsigned char) ~ZMQ_MSG_MASK;
         msg_->vsm_size = (uint8_t) size_;
     }
     else {
@@ -50,7 +51,7 @@ int zmq_msg_init_size (zmq_msg_t *msg_, size_t size_)
             errno = ENOMEM;
             return -1;
         }
-        msg_->flags = 0;
+        msg_->flags = (unsigned char) ~ZMQ_MSG_MASK;
         
         zmq::msg_content_t *content = (zmq::msg_content_t*) msg_->content;
         content->data = (void*) (content + 1);
@@ -67,7 +68,7 @@ int zmq_msg_init_data (zmq_msg_t *msg_, void *data_, size_t size_,
 {
     msg_->content = (zmq::msg_content_t*) malloc (sizeof (zmq::msg_content_t));
     alloc_assert (msg_->content);
-    msg_->flags = 0;
+    msg_->flags = (unsigned char) ~ZMQ_MSG_MASK;
     zmq::msg_content_t *content = (zmq::msg_content_t*) msg_->content;
     content->data = data_;
     content->size = size_;
@@ -79,6 +80,12 @@ int zmq_msg_init_data (zmq_msg_t *msg_, void *data_, size_t size_,
 
 int zmq_msg_close (zmq_msg_t *msg_)
 {
+    //  Check the validity tag.
+    if (unlikely (msg_->flags | ZMQ_MSG_MASK) != 0xff) {
+        errno = EFAULT;
+        return -1;
+    }
+
     //  For VSMs and delimiters there are no resources to free.
     if (msg_->content != (zmq::msg_content_t*) ZMQ_DELIMITER &&
           msg_->content != (zmq::msg_content_t*) ZMQ_VSM) {
@@ -98,17 +105,21 @@ int zmq_msg_close (zmq_msg_t *msg_)
         }
     }
 
-    //  As a safety measure, let's make the deallocated message look like
-    //  an empty message.
-    msg_->content = (zmq::msg_content_t*) ZMQ_VSM;
+    //  Remove the validity tag from the message.
     msg_->flags = 0;
-    msg_->vsm_size = 0;
 
     return 0;
 }
 
 int zmq_msg_move (zmq_msg_t *dest_, zmq_msg_t *src_)
 {
+    //  Check the validity tags.
+    if (unlikely ((dest_->flags | ZMQ_MSG_MASK) != 0xff ||
+          (src_->flags | ZMQ_MSG_MASK) != 0xff)) {
+        errno = EFAULT;
+        return -1;
+    }
+
     zmq_msg_close (dest_);
     *dest_ = *src_;
     zmq_msg_init (src_);
@@ -117,6 +128,13 @@ int zmq_msg_move (zmq_msg_t *dest_, zmq_msg_t *src_)
 
 int zmq_msg_copy (zmq_msg_t *dest_, zmq_msg_t *src_)
 {
+    //  Check the validity tags.
+    if (unlikely ((dest_->flags | ZMQ_MSG_MASK) != 0xff ||
+          (src_->flags | ZMQ_MSG_MASK) != 0xff)) {
+        errno = EFAULT;
+        return -1;
+    }
+
     zmq_msg_close (dest_);
 
     //  VSMs and delimiters require no special handling.
@@ -140,6 +158,9 @@ int zmq_msg_copy (zmq_msg_t *dest_, zmq_msg_t *src_)
 
 void *zmq_msg_data (zmq_msg_t *msg_)
 {
+    //  Check the validity tag.
+    zmq_assert ((msg_->flags | ZMQ_MSG_MASK) == 0xff);
+
     if (msg_->content == (zmq::msg_content_t*) ZMQ_VSM)
         return msg_->vsm_data;
     if (msg_->content == (zmq::msg_content_t*) ZMQ_DELIMITER)
@@ -150,6 +171,9 @@ void *zmq_msg_data (zmq_msg_t *msg_)
 
 size_t zmq_msg_size (zmq_msg_t *msg_)
 {
+    //  Check the validity tag.
+    zmq_assert ((msg_->flags | ZMQ_MSG_MASK) == 0xff);
+
     if (msg_->content == (zmq::msg_content_t*) ZMQ_VSM)
         return msg_->vsm_size;
     if (msg_->content == (zmq::msg_content_t*) ZMQ_DELIMITER)
diff --git a/src/req.cpp b/src/req.cpp
index f495492..503f221 100644
--- a/src/req.cpp
+++ b/src/req.cpp
@@ -49,7 +49,7 @@ int zmq::req_t::xsend (zmq_msg_t *msg_, int flags_)
         zmq_msg_t prefix;
         int rc = zmq_msg_init (&prefix);
         zmq_assert (rc == 0);
-        prefix.flags = ZMQ_MSG_MORE;
+        prefix.flags |= ZMQ_MSG_MORE;
         rc = xreq_t::xsend (&prefix, flags_);
         if (rc != 0)
             return rc;
diff --git a/src/socket_base.cpp b/src/socket_base.cpp
index c9b5c31..9f3b1f6 100644
--- a/src/socket_base.cpp
+++ b/src/socket_base.cpp
@@ -466,11 +466,18 @@ int zmq::socket_base_t::connect (const char *addr_)
 
 int zmq::socket_base_t::send (::zmq_msg_t *msg_, int flags_)
 {
+    //  Check whether the library haven't been shut down yet.
     if (unlikely (ctx_terminated)) {
         errno = ETERM;
         return -1;
     }
 
+    //  Check whether message passed to the function is valid.
+    if (unlikely ((msg_->flags | ZMQ_MSG_MASK) != 0xff)) {
+        errno = EFAULT;
+        return -1;
+    }
+
     //  Process pending commands, if any.
     int rc = process_commands (false, true);
     if (unlikely (rc != 0))
@@ -504,11 +511,18 @@ int zmq::socket_base_t::send (::zmq_msg_t *msg_, int flags_)
 
 int zmq::socket_base_t::recv (::zmq_msg_t *msg_, int flags_)
 {
+    //  Check whether the library haven't been shut down yet.
     if (unlikely (ctx_terminated)) {
         errno = ETERM;
         return -1;
     }
 
+    //  Check whether message passed to the function is valid.
+    if (unlikely ((msg_->flags | ZMQ_MSG_MASK) != 0xff)) {
+        errno = EFAULT;
+        return -1;
+    }
+
     //  Get the message.
     int rc = xrecv (msg_, flags_);
     int err = errno;
diff --git a/src/xrep.cpp b/src/xrep.cpp
index 7f0da4d..75dc30e 100644
--- a/src/xrep.cpp
+++ b/src/xrep.cpp
@@ -269,7 +269,7 @@ int zmq::xrep_t::xrecv (zmq_msg_t *msg_, int flags_)
             zmq_assert (rc == 0);
             memcpy (zmq_msg_data (msg_), inpipes [current_in].identity.data (),
                 zmq_msg_size (msg_));
-            msg_->flags = ZMQ_MSG_MORE;
+            msg_->flags |= ZMQ_MSG_MORE;
             return 0;
         }