# HG changeset patch
# User jbe
# Date 1472246156 -7200
# Node ID 28bffa2af1ec0b2576f56a9f1596567c9415d029
# Parent  53962483bf1c0a94b51df069950da14eb680d99a
Perform TLS handshake in :read and :write methods and store TLS_WANT_POLLIN, TLS_WANT_POLLOUT in handle structure

diff -r 53962483bf1c -r 28bffa2af1ec moonbridge_io.c
--- a/moonbridge_io.c	Fri Aug 26 22:04:08 2016 +0200
+++ b/moonbridge_io.c	Fri Aug 26 23:15:56 2016 +0200
@@ -83,7 +83,9 @@
   char readbuf[MOONBR_IO_READBUFLEN];
   char writebuf[MOONBR_IO_WRITEBUFLEN];
 #ifdef MOONBR_IO_USE_TLS
-  struct tls *tls, *servertls;
+  struct tls *tls;
+  struct tls *servertls;
+  int tlshandshake;
 #endif
 } moonbr_io_handle_t;
 
@@ -277,15 +279,34 @@
     handle->readbufout = 0;
 #ifdef MOONBR_IO_USE_TLS
     if (handle->tls) {
-      do bytesread = tls_read(handle->tls, handle->readbuf, MOONBR_IO_READBUFLEN);
-      while (!nonblocking && (bytesread == TLS_WANT_POLLIN || bytesread == TLS_WANT_POLLOUT));
-      if (bytesread == TLS_WANT_POLLIN || bytesread == TLS_WANT_POLLOUT) {
-        errno = EAGAIN;
-      } else if (bytesread < 0) {
-        lua_pushnil(L);
-        lua_pushstring(L, tls_error(handle->tls));
-        return 2;
-      }
+      do {
+        if (!handle->tlshandshake) {
+          do bytesread = tls_handshake(handle->tls);
+          while (!nonblocking && (bytesread == TLS_WANT_POLLIN || bytesread == TLS_WANT_POLLOUT));
+          if (bytesread == TLS_WANT_POLLIN || bytesread == TLS_WANT_POLLOUT) {
+            handle->tlshandshake = bytesread;
+            errno = EAGAIN;
+            break;
+          }
+          if (bytesread < 0) {
+            lua_pushnil(L);
+            lua_pushstring(L, tls_error(handle->tls));
+            return 2;
+          }
+          handle->tlshandshake = 1;
+        }
+        do bytesread = tls_read(handle->tls, handle->readbuf, MOONBR_IO_READBUFLEN);
+        while (!nonblocking && (bytesread == TLS_WANT_POLLIN || bytesread == TLS_WANT_POLLOUT));
+        if (bytesread == TLS_WANT_POLLIN || bytesread == TLS_WANT_POLLOUT) {
+          errno = EAGAIN;
+          break;
+        }
+        if (bytesread < 0) {
+          lua_pushnil(L);
+          lua_pushstring(L, tls_error(handle->tls));
+          return 2;
+        }
+      } while (0);
     }
     else
 #endif
@@ -452,6 +473,41 @@
 
 moonbr_io_yield_wrapper(moonbr_io_drain_yield, moonbr_io_drain_call);
 
+#ifdef MOONBR_IO_USE_TLS
+#define moonbr_io_write_tls(buf, buflen) \
+  if (handle->tls) { \
+    do { \
+      if (!handle->tlshandshake) { \
+        do written = tls_handshake(handle->tls); \
+        while (!nonblocking && (written == TLS_WANT_POLLIN || written == TLS_WANT_POLLOUT)); \
+        if (written == TLS_WANT_POLLIN || written == TLS_WANT_POLLOUT) { \
+          handle->tlshandshake = written; \
+          errno = EAGAIN; \
+          break; \
+        } \
+        if (written < 0) { \
+          lua_pushnil(L); \
+          lua_pushstring(L, tls_error(handle->tls)); \
+          return 2; \
+        } \
+        handle->tlshandshake = 1; \
+      } \
+      do written = tls_write(handle->tls, (buf), (buflen)); \
+      while (!nonblocking && (written == TLS_WANT_POLLIN || written == TLS_WANT_POLLOUT)); \
+      if (written == TLS_WANT_POLLIN || written == TLS_WANT_POLLOUT) { \
+        errno = EAGAIN; \
+        break; \
+      } \
+      if (written < 0) { \
+        lua_pushnil(L); \
+        lua_pushstring(L, tls_error(handle->tls)); \
+        return 2; \
+      } \
+    } while (0); \
+  } \
+  else
+#endif
+
 static int moonbr_io_write_impl(lua_State *L, int nonblocking, int flush) {
   moonbr_io_handle_t *handle;
   int i, top;
@@ -504,18 +560,10 @@
         while (handle->writebufout < MOONBR_IO_WRITEBUFLEN) {
           moonbr_io_handle_set_nopush(L, handle, 1);
 #ifdef MOONBR_IO_USE_TLS
-          if (handle->tls) {
-            do written = tls_write(handle->tls, handle->writebuf + handle->writebufout, MOONBR_IO_WRITEBUFLEN - handle->writebufout);
-            while (!nonblocking && (written == TLS_WANT_POLLIN || written == TLS_WANT_POLLOUT));
-            if (written == TLS_WANT_POLLIN || written == TLS_WANT_POLLOUT) {
-              errno = EAGAIN;
-            } else if (written < 0) {
-              lua_pushnil(L);
-              lua_pushstring(L, tls_error(handle->tls));
-              return 2;
-            }
-          }
-          else
+          moonbr_io_write_tls(
+            handle->writebuf + handle->writebufout,
+            MOONBR_IO_WRITEBUFLEN - handle->writebufout
+          )
 #endif
           written = write(
             handle->fd,
@@ -556,18 +604,10 @@
   while (handle->flushedleft) {
     moonbr_io_handle_set_nopush(L, handle, 1);
 #ifdef MOONBR_IO_USE_TLS
-    if (handle->tls) {
-      do written = tls_write(handle->tls, handle->writebuf + handle->writebufout, handle->writebufin - handle->writebufout);
-      while (!nonblocking && (written == TLS_WANT_POLLIN || written == TLS_WANT_POLLOUT));
-      if (written == TLS_WANT_POLLIN || written == TLS_WANT_POLLOUT) {
-        errno = EAGAIN;
-      } else if (written < 0) {
-        lua_pushnil(L);
-        lua_pushstring(L, tls_error(handle->tls));
-        return 2;
-      }
-    }
-    else
+    moonbr_io_write_tls(
+      handle->writebuf + handle->writebufout,
+      handle->writebufin - handle->writebufout
+    )
 #endif
     written = write(
       handle->fd,
@@ -844,6 +884,7 @@
 #ifdef MOONBR_IO_USE_TLS
   handle->tls = NULL;
   handle->servertls = NULL;
+  handle->tlshandshake = 0;
 #endif
   handle->fd = *fd;  /* required for set_linger call */
   moonbr_io_handle_set_linger(L, handle, 0);
@@ -1583,10 +1624,34 @@
 
 moonbr_io_yield_wrapper(moonbr_io_wait_yield, moonbr_io_wait_call);
 
+#ifdef MOONBR_IO_USE_TLS
+#define moonbr_io_poll_tls() \
+  if (!handle->tlshandshake) { \
+    lua_pushboolean(L, 1); \
+    return 1; \
+  } \
+  if (handle->tlshandshake == TLS_WANT_POLLIN) { \
+    if (fd < 0) { \
+      lua_pushboolean(L, 1); \
+      return 1; \
+    } \
+    FD_SET(fd, &readfds); \
+    if (fd+1 > nfds) nfds = fd+1; \
+    continue; \
+  } \
+  if (handle->tlshandshake == TLS_WANT_POLLOUT) { \
+    if (fd < 0) { \
+      lua_pushboolean(L, 1); \
+      return 1; \
+    } \
+    FD_SET(fd, &writefds); \
+    if (fd+1 > nfds) nfds = fd+1; \
+    continue; \
+  } \
+  while (0)
+#endif
+
 static int moonbr_io_poll(lua_State *L) {
-#if MOONBR_IO_USE_TLS
-  // TODO: tls_handshake must probably not be called when handshake has been completed
-#endif
   moonbr_io_handle_t *handle;
   moonbr_io_listener_t *listener;
   int fd, isnum;
@@ -1606,16 +1671,7 @@
           if (handle->closed) luaL_error(L, "Attempt to poll a closed connection");
           fd = handle->fd;
 #if MOONBR_IO_USE_TLS
-          if (handle->tls && tls_handshake(handle->tls) == TLS_WANT_POLLOUT) {
-            if (fd < 0) {
-              // TODO?
-              lua_pushboolean(L, 1);
-              return 1;
-            }
-            FD_SET(fd, &writefds);
-            if (fd+1 > nfds) nfds = fd+1;
-            continue;
-          }
+          moonbr_io_poll_tls();
 #endif
           if (
             fd < 0 ||  /* fake EOF to simulate shutdown if fd < 0 */
@@ -1650,11 +1706,7 @@
           if (handle->finished) luaL_error(L, "Attempt to write-poll a finished connection");
           fd = handle->fd;
 #if MOONBR_IO_USE_TLS
-          if (handle->tls && tls_handshake(handle->tls) == TLS_WANT_POLLIN) {
-            FD_SET(fd, &readfds);
-            if (fd+1 > nfds) nfds = fd+1;
-            continue;
-          }
+          moonbr_io_poll_tls();
 #endif
         } else {
           listener = luaL_testudata(L, -2, MOONBR_IO_LISTENER_MT_REGKEY);