Merge pull request #3703 from somdoron/ws_mechanism

problem: wss transport return incorrect return code for error

src/stream_engine_base.cpp

@@ -354,7 +354,7 @@ void zmq::stream_engine_base_t::out_event ()
         //  If there is no data to send, stop polling for output.
         if (_outsize == 0) {
             _output_stopped = true;
-            reset_pollout (_handle);
+            reset_pollout ();
             return;
         }
     }
@@ -370,7 +370,7 @@ void zmq::stream_engine_base_t::out_event ()
     //  The engine is not terminated until we detect input error;
     //  this is necessary to prevent losing incoming messages.
     if (nbytes == -1) {
-        reset_pollout (_handle);
+        reset_pollout ();
         return;
     }
 
@@ -381,7 +381,7 @@ void zmq::stream_engine_base_t::out_event ()
     //  to send, stop polling for output.
     if (unlikely (_handshaking))
         if (_outsize == 0)
-            reset_pollout (_handle);
+            reset_pollout ();
 }
 
 void zmq::stream_engine_base_t::restart_output ()

src/stream_engine_base.hpp

@@ -102,6 +102,7 @@ class stream_engine_base_t : public io_object_t, public i_engine
     virtual int read (void *data, size_t size_);
     virtual int write (const void *data_, size_t size_);
 
+    void reset_pollout () { io_object_t::reset_pollout (_handle); }
     void set_pollout () { io_object_t::set_pollout (_handle); }
     void set_pollin () { io_object_t::set_pollin (_handle); }
     session_base_t *session () { return _session; }

src/wss_engine.cpp

@@ -128,21 +128,53 @@ void zmq::wss_engine_t::plug_internal ()
     in_event ();
 }
 
+void zmq::wss_engine_t::out_event ()
+{
+    if (_established)
+        return ws_engine_t::out_event ();
+
+    int rc = gnutls_handshake (_tls_session);
+
+    reset_pollout ();
+
+    if (rc == GNUTLS_E_SUCCESS) {
+        start_ws_handshake ();
+        _established = true;
+        return;
+    } else if (rc == GNUTLS_E_AGAIN) {
+        int direction = gnutls_record_get_direction (_tls_session);
+        if (direction == 1)
+            set_pollout ();
+
+        return;
+    } else if (rc == GNUTLS_E_INTERRUPTED
+               || rc == GNUTLS_E_WARNING_ALERT_RECEIVED) {
+        return;
+    } else {
+        error (zmq::i_engine::connection_error);
+        return;
+    }
+}
+
 bool zmq::wss_engine_t::handshake ()
 {
     if (!_established) {
         int rc = gnutls_handshake (_tls_session);
 
-        // TODO: when E_AGAIN is returned we might need to call gnutls_handshake for out_event as well, see gnutls_record_get_direction
-
         if (rc == GNUTLS_E_SUCCESS) {
             start_ws_handshake ();
             _established = true;
             return false;
-        } else if (rc == GNUTLS_E_AGAIN || rc == GNUTLS_E_INTERRUPTED
-                   || rc == GNUTLS_E_WARNING_ALERT_RECEIVED)
+        } else if (rc == GNUTLS_E_AGAIN) {
+            int direction = gnutls_record_get_direction (_tls_session);
+            if (direction == 1)
+                set_pollout ();
+
+            return false;
+        } else if (rc == GNUTLS_E_INTERRUPTED
+                   || rc == GNUTLS_E_WARNING_ALERT_RECEIVED) {
             return false;
-        else {
+        } else {
             error (zmq::i_engine::connection_error);
             return false;
         }
@@ -171,6 +203,11 @@ int zmq::wss_engine_t::read (void *data_, size_t size_)
         return -1;
     }
 
+    if (rc < 0) {
+        errno = EINVAL;
+        return -1;
+    }
+
     // TODO: change return type to ssize_t (signed)
     return rc;
 }
@@ -189,6 +226,11 @@ int zmq::wss_engine_t::write (const void *data_, size_t size_)
         return -1;
     }
 
+    if (rc < 0) {
+        errno = EINVAL;
+        return -1;
+    }
+
     // TODO: change return type to ssize_t (signed)
     return rc;
 }

src/wss_engine.hpp

@@ -49,6 +49,8 @@ class wss_engine_t : public ws_engine_t
                   const char *hostname_);
     ~wss_engine_t ();
 
+    void out_event ();
+
   protected:
     bool handshake ();
     void plug_internal ();