common: tls, fix for pending reads

1 files changed, 33 insertions, 10 deletions

diff --git a/common/trans.c b/common/trans.c
index 5503ea61..5f435bcd 100644
--- a/common/trans.c
+++ b/common/trans.c
@@ -47,6 +47,17 @@ trans_tls_send(struct trans *self, const void *data, int len)
 
 /*****************************************************************************/
 int APP_CC
+trans_tls_can_recv(struct trans *self, int sck, int millis)
+{
+    if (self->tls == NULL)
+    {
+        return 1;
+    }
+    return xrdp_tls_can_recv(self->tls, sck, millis);
+}
+
+/*****************************************************************************/
+int APP_CC
 trans_tcp_recv(struct trans *self, void *ptr, int len)
 {
     return g_tcp_recv(self->sck, ptr, len, 0);
@@ -60,6 +71,13 @@ trans_tcp_send(struct trans *self, const void *data, int len)
 }
 
 /*****************************************************************************/
+int APP_CC
+trans_tcp_can_recv(struct trans *self, int sck, int millis)
+{
+    return g_tcp_can_recv(sck, millis);
+}
+
+/*****************************************************************************/
 struct trans *
 APP_CC
 trans_create(int mode, int in_size, int out_size)
@@ -79,6 +97,7 @@ trans_create(int mode, int in_size, int out_size)
         /* assign tcp calls by default */
         self->trans_recv = trans_tcp_recv;
         self->trans_send = trans_tcp_send;
+        self->trans_can_recv = trans_tcp_can_recv;
     }
 
     return self;
@@ -133,6 +152,16 @@ trans_get_wait_objs(struct trans *self, tbus *objs, int *count)
 
     objs[*count] = self->sck;
     (*count)++;
+
+    if (self->tls != 0)
+    {
+        if (self->tls->rwo != 0)
+        {
+            objs[*count] = self->tls->rwo;
+            (*count)++;
+        }
+    }
+
     return 0;
 }
 
@@ -141,19 +170,11 @@ int APP_CC
 trans_get_wait_objs_rw(struct trans *self, tbus *robjs, int *rcount,
                        tbus *wobjs, int *wcount)
 {
-    if (self == 0)
+    if (trans_get_wait_objs(self, robjs, rcount) != 0)
     {
         return 1;
     }
 
-    if (self->status != TRANS_STATUS_UP)
-    {
-        return 1;
-    }
-
-    robjs[*rcount] = self->sck;
-    (*rcount)++;
-
     if (self->wait_s != 0)
     {
         wobjs[*wcount] = self->sck;
@@ -288,7 +309,7 @@ trans_check_wait_objs(struct trans *self)
     }
     else /* connected server or client (2 or 3) */
     {
-        if (g_tcp_can_recv(self->sck, 0))
+        if (self->trans_can_recv(self, self->sck, 0))
         {
             read_so_far = (int) (self->in_s->end - self->in_s->data);
             to_read = self->header_size - read_so_far;
@@ -716,6 +737,7 @@ trans_set_tls_mode(struct trans *self, const char *key, const char *cert)
     /* assign tls functions */
     self->trans_recv = trans_tls_recv;
     self->trans_send = trans_tls_send;
+    self->trans_can_recv = trans_tls_can_recv;
 
     return 0;
 }
@@ -732,6 +754,7 @@ trans_shutdown_tls_mode(struct trans *self)
     /* assign callback back to tcp cal */
     self->trans_recv = trans_tcp_recv;
     self->trans_send = trans_tcp_send;
+    self->trans_can_recv = trans_tcp_can_recv;
 
     return 0;
 }