From a9bb526403cd17eff13108db804237e76ad105c5 Mon Sep 17 00:00:00 2001
From: somdoron <somdoron@gmail.com>
Date: Fri, 4 Oct 2019 20:45:54 +0300
Subject: [PATCH] problem: wss transport return incorrect return code for error

Solution: return -1 for errors
---
 src/stream_engine_base.cpp |  6 ++---
 src/stream_engine_base.hpp |  1 +
 src/wss_engine.cpp         | 52 ++++++++++++++++++++++++++++++++++----
 src/wss_engine.hpp         |  2 ++
 4 files changed, 53 insertions(+), 8 deletions(-)

diff --git a/src/stream_engine_base.cpp b/src/stream_engine_base.cpp
index 66c989bf..46a2fac4 100644
--- a/src/stream_engine_base.cpp
+++ b/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 ()
diff --git a/src/stream_engine_base.hpp b/src/stream_engine_base.hpp
index 1639ecf5..a932d92b 100644
--- a/src/stream_engine_base.hpp
+++ b/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; }
diff --git a/src/wss_engine.cpp b/src/wss_engine.cpp
index 62018e53..34750225 100644
--- a/src/wss_engine.cpp
+++ b/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 {
+        } else if (rc == GNUTLS_E_INTERRUPTED
+                   || rc == GNUTLS_E_WARNING_ALERT_RECEIVED) {
+            return false;
+        } 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;
 }
diff --git a/src/wss_engine.hpp b/src/wss_engine.hpp
index 76c3cdb3..1f744e83 100644
--- a/src/wss_engine.hpp
+++ b/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 ();