Add helper to receive sftp response messages

For the sake of reducing code repetition, this commit
adds a helper function to receive sftp response
messages. The function can operate in both blocking
and non-blocking modes.

Signed-off-by: Eshan Kelkar <eshankelkar@galorithm.com>
Reviewed-by: Jakub Jelen <jjelen@redhat.com>

include/libssh/sftp_priv.h

@@ -21,6 +21,8 @@
 #ifndef SFTP_PRIV_H
 #define SFTP_PRIV_H
 
+#include <stdbool.h>
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -62,6 +64,39 @@ int sftp_read_and_dispatch(sftp_session sftp);
 
 sftp_message sftp_dequeue(sftp_session sftp, uint32_t id);
 
+/**
+ * @brief Receive the response of an sftp request
+ *
+ * In blocking mode, if the response hasn't arrived at the time of call, this
+ * function waits for the response to arrive.
+ *
+ * @param sftp           The sftp session via which the request was sent.
+ *
+ * @param id             The request identifier of the request whose
+ *                       corresponding response is required.
+ *
+ * @param blocking       Flag to indicate the operating mode. true indicates
+ *                       blocking mode and false indicates non-blocking mode
+ *
+ * @param msg_ptr        Pointer to the location to store the response message.
+ *                       In case of success, the message is allocated
+ *                       dynamically and must be freed (using
+ *                       sftp_message_free()) by the caller after usage. In case
+ *                       of failure, this is left untouched.
+ *
+ * @returns SSH_OK on success
+ * @returns SSH_ERROR on failure with the sftp and ssh errors set
+ * @returns SSH_AGAIN in case of non-blocking mode if the response hasn't
+ *          arrived yet.
+ *
+ * @warning In blocking mode, this may block indefinitely for an invalid request
+ *          identifier.
+ */
+int sftp_recv_response_msg(sftp_session sftp,
+                           uint32_t id,
+                           bool blocking,
+                           sftp_message *msg_ptr);
+
 /*
  * Assigns a new SFTP ID for new requests and assures there is no collision
  * between them.

src/sftp.c

@@ -701,13 +701,10 @@ sftp_dir sftp_opendir(sftp_session sftp, const char *path)
         return NULL;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        if (sftp_read_and_dispatch(sftp) < 0) {
+    if (rc != SSH_OK) {
-            /* something nasty has happened */
         return NULL;
     }
-        msg = sftp_dequeue(sftp, id);
-    }
 
     switch (msg->packet_type) {
         case SSH_FXP_STATUS:
@@ -798,13 +795,10 @@ sftp_attributes sftp_readdir(sftp_session sftp, sftp_dir dir)
         SSH_LOG(SSH_LOG_PACKET,
                 "Sent a ssh_fxp_readdir with id %" PRIu32, id);
 
-        while (msg == NULL) {
+        rc = sftp_recv_response_msg(sftp, id, true, &msg);
-            if (sftp_read_and_dispatch(sftp) < 0) {
+        if (rc != SSH_OK) {
-                /* something nasty has happened */
             return NULL;
         }
-            msg = sftp_dequeue(sftp, id);
-        }
 
         switch (msg->packet_type){
             case SSH_FXP_STATUS:
@@ -928,13 +922,10 @@ static int sftp_handle_close(sftp_session sftp, ssh_string handle)
         return -1;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        if (sftp_read_and_dispatch(sftp) < 0) {
+    if (rc != SSH_OK) {
-            /* something nasty has happened */
         return -1;
     }
-        msg = sftp_dequeue(sftp,id);
-    }
 
     switch (msg->packet_type) {
         case SSH_FXP_STATUS:
@@ -1073,13 +1064,10 @@ sftp_file sftp_open(sftp_session sftp,
         return NULL;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        if (sftp_read_and_dispatch(sftp) < 0) {
+    if (rc != SSH_OK) {
-            /* something nasty has happened */
         return NULL;
     }
-        msg = sftp_dequeue(sftp, id);
-    }
 
     switch (msg->packet_type) {
     case SSH_FXP_STATUS:
@@ -1197,18 +1185,17 @@ ssize_t sftp_read(sftp_file handle, void *buf, size_t count) {
   }
   SSH_BUFFER_FREE(buffer);
 
-  while (msg == NULL) {
+  rc = sftp_recv_response_msg(handle->sftp, id, !handle->nonblocking, &msg);
-    if (handle->nonblocking) {
+  if (rc == SSH_ERROR) {
-      if (ssh_channel_poll(handle->sftp->channel, 0) == 0) {
-        /* we cannot block */
-        return 0;
-      }
-    }
-    if (sftp_read_and_dispatch(handle->sftp) < 0) {
-      /* something nasty has happened */
       return -1;
   }
-    msg = sftp_dequeue(handle->sftp, id);
+
+  if (rc == SSH_AGAIN) {
+      /*
+       * file opened in non blocking mode and the response has not arrived yet.
+       * Since we cannot block, return 0 as the number of bytes read.
+       */
+      return 0;
   }
 
   switch (msg->packet_type) {
@@ -1309,7 +1296,7 @@ int sftp_async_read(sftp_file file, void *data, uint32_t size, uint32_t id){
   sftp_message msg = NULL;
   sftp_status_message status;
   ssh_string datastring;
-  int err = SSH_OK;
+  int rc, err = SSH_OK;
   uint32_t len;
 
   if (file == NULL) {
@@ -1322,20 +1309,9 @@ int sftp_async_read(sftp_file file, void *data, uint32_t size, uint32_t id){
   }
 
   /* handle an existing request */
-  while (msg == NULL) {
+  rc = sftp_recv_response_msg(sftp, id, !file->nonblocking, &msg);
-    if (file->nonblocking){
+  if (rc == SSH_ERROR || rc == SSH_AGAIN) {
-      if (ssh_channel_poll(sftp->channel, 0) == 0) {
+      return rc;
-        /* we cannot block */
-        return SSH_AGAIN;
-      }
-    }
-
-    if (sftp_read_and_dispatch(sftp) < 0) {
-      /* something nasty has happened */
-      return SSH_ERROR;
-    }
-
-    msg = sftp_dequeue(sftp,id);
   }
 
   switch (msg->packet_type) {
@@ -1446,13 +1422,11 @@ ssize_t sftp_write(sftp_file file, const void *buf, size_t count) {
         "Could not write as much data as expected");
   }
 
-  while (msg == NULL) {
+  /* Wait for the response in blocking mode */
-    if (sftp_read_and_dispatch(file->sftp) < 0) {
+  rc = sftp_recv_response_msg(sftp, id, true, &msg);
-      /* something nasty has happened */
+  if (rc != SSH_OK) {
       return -1;
   }
-    msg = sftp_dequeue(file->sftp, id);
-  }
 
   switch (msg->packet_type) {
     case SSH_FXP_STATUS:
@@ -1558,12 +1532,10 @@ int sftp_unlink(sftp_session sftp, const char *file) {
   }
   SSH_BUFFER_FREE(buffer);
 
-  while (msg == NULL) {
+  rc = sftp_recv_response_msg(sftp, id, true, &msg);
-    if (sftp_read_and_dispatch(sftp)) {
+  if (rc != SSH_OK) {
       return -1;
   }
-    msg = sftp_dequeue(sftp, id);
-  }
 
   if (msg->packet_type == SSH_FXP_STATUS) {
     /* by specification, this command's only supposed to return SSH_FXP_STATUS */
@@ -1632,12 +1604,10 @@ int sftp_rmdir(sftp_session sftp, const char *directory) {
   }
   SSH_BUFFER_FREE(buffer);
 
-  while (msg == NULL) {
+  rc = sftp_recv_response_msg(sftp, id, true, &msg);
-    if (sftp_read_and_dispatch(sftp) < 0) {
+  if (rc != SSH_OK) {
       return -1;
   }
-    msg = sftp_dequeue(sftp, id);
-  }
 
   /* By specification, this command returns SSH_FXP_STATUS */
   if (msg->packet_type == SSH_FXP_STATUS) {
@@ -1718,12 +1688,10 @@ int sftp_mkdir(sftp_session sftp, const char *directory, mode_t mode)
         return -1;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        if (sftp_read_and_dispatch(sftp) < 0) {
+    if (rc != SSH_OK) {
         return -1;
     }
-        msg = sftp_dequeue(sftp, id);
-    }
 
     /* By specification, this command only returns SSH_FXP_STATUS */
     if (msg->packet_type == SSH_FXP_STATUS) {
@@ -1842,12 +1810,10 @@ int sftp_rename(sftp_session sftp, const char *original, const char *newname)
         return -1;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        if (sftp_read_and_dispatch(sftp) < 0) {
+    if (rc != SSH_OK) {
         return -1;
     }
-        msg = sftp_dequeue(sftp, id);
-    }
 
     /* By specification, this command only returns SSH_FXP_STATUS */
     if (msg->packet_type == SSH_FXP_STATUS) {
@@ -1931,12 +1897,10 @@ int sftp_setstat(sftp_session sftp, const char *file, sftp_attributes attr)
         return -1;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        if (sftp_read_and_dispatch(sftp) < 0) {
+    if (rc != SSH_OK) {
         return -1;
     }
-        msg = sftp_dequeue(sftp, id);
-    }
 
     /* By specification, this command only returns SSH_FXP_STATUS */
     if (msg->packet_type == SSH_FXP_STATUS) {
@@ -2016,12 +1980,10 @@ sftp_lsetstat(sftp_session sftp, const char *file, sftp_attributes attr)
         return -1;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        if (sftp_read_and_dispatch(sftp) < 0) {
+    if (rc != SSH_OK) {
         return -1;
     }
-        msg = sftp_dequeue(sftp, id);
-    }
 
     /* By specification, this command only returns SSH_FXP_STATUS */
     if (msg->packet_type == SSH_FXP_STATUS) {
@@ -2156,12 +2118,10 @@ int sftp_symlink(sftp_session sftp, const char *target, const char *dest)
   }
   SSH_BUFFER_FREE(buffer);
 
-  while (msg == NULL) {
+  rc = sftp_recv_response_msg(sftp, id, true, &msg);
-    if (sftp_read_and_dispatch(sftp) < 0) {
+  if (rc != SSH_OK) {
       return -1;
   }
-    msg = sftp_dequeue(sftp, id);
-  }
 
   /* By specification, this command only returns SSH_FXP_STATUS */
   if (msg->packet_type == SSH_FXP_STATUS) {
@@ -2244,12 +2204,10 @@ char *sftp_readlink(sftp_session sftp, const char *path)
         return NULL;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        if (sftp_read_and_dispatch(sftp) < 0) {
+    if (rc != SSH_OK) {
         return NULL;
     }
-        msg = sftp_dequeue(sftp, id);
-    }
 
     if (msg->packet_type == SSH_FXP_NAME) {
         uint32_t ignored = 0;
@@ -2336,12 +2294,10 @@ int sftp_hardlink(sftp_session sftp, const char *oldpath, const char *newpath)
         return -1;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        if (sftp_read_and_dispatch(sftp) < 0) {
+    if (rc != SSH_OK) {
         return -1;
     }
-        msg = sftp_dequeue(sftp, id);
-    }
 
     /* By specification, this command only returns SSH_FXP_STATUS */
     if (msg->packet_type == SSH_FXP_STATUS) {
@@ -2459,12 +2415,10 @@ sftp_statvfs_t sftp_statvfs(sftp_session sftp, const char *path)
         return NULL;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        if (sftp_read_and_dispatch(sftp) < 0) {
+    if (rc != SSH_OK) {
         return NULL;
     }
-        msg = sftp_dequeue(sftp, id);
-    }
 
     if (msg->packet_type == SSH_FXP_EXTENDED_REPLY) {
         sftp_statvfs_t  buf = sftp_parse_statvfs(sftp, msg->payload);
@@ -2533,15 +2487,10 @@ int sftp_fsync(sftp_file file)
         goto done;
     }
 
-    do {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        rc = sftp_read_and_dispatch(sftp);
+    if (rc != SSH_OK) {
-        if (rc < 0) {
+        return -1;
-            ssh_set_error_oom(sftp->session);
-            rc = -1;
-            goto done;
     }
-        msg = sftp_dequeue(sftp, id);
-    } while (msg == NULL);
 
     /* By specification, this command only returns SSH_FXP_STATUS */
     if (msg->packet_type == SSH_FXP_STATUS) {
@@ -2634,12 +2583,10 @@ sftp_statvfs_t sftp_fstatvfs(sftp_file file)
         return NULL;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        if (sftp_read_and_dispatch(sftp) < 0) {
+    if (rc == -1) {
         return NULL;
     }
-        msg = sftp_dequeue(sftp, id);
-    }
 
     if (msg->packet_type == SSH_FXP_EXTENDED_REPLY) {
         sftp_statvfs_t buf = sftp_parse_statvfs(sftp, msg->payload);
@@ -2747,12 +2694,10 @@ static sftp_limits_t sftp_limits_use_extension(sftp_session sftp)
         return NULL;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        if (sftp_read_and_dispatch(sftp) < 0) {
+    if (rc != SSH_OK) {
         return NULL;
     }
-        msg = sftp_dequeue(sftp, id);
-    }
 
     if (msg->packet_type == SSH_FXP_EXTENDED_REPLY) {
         sftp_limits_t limits = sftp_parse_limits(sftp, msg->payload);
@@ -2897,12 +2842,10 @@ char *sftp_canonicalize_path(sftp_session sftp, const char *path)
         return NULL;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        if (sftp_read_and_dispatch(sftp) < 0) {
+    if (rc == -1) {
         return NULL;
     }
-        msg = sftp_dequeue(sftp, id);
-    }
 
     if (msg->packet_type == SSH_FXP_NAME) {
         uint32_t ignored = 0;
@@ -2988,12 +2931,10 @@ static sftp_attributes sftp_xstat(sftp_session sftp,
         return NULL;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        if (sftp_read_and_dispatch(sftp) < 0) {
+    if (rc != SSH_OK) {
         return NULL;
     }
-        msg = sftp_dequeue(sftp, id);
-    }
 
     if (msg->packet_type == SSH_FXP_ATTRS) {
         sftp_attributes attr = sftp_parse_attr(sftp, msg->payload, 0);
@@ -3066,12 +3007,10 @@ sftp_attributes sftp_fstat(sftp_file file)
         return NULL;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(file->sftp, id, true, &msg);
-        if (sftp_read_and_dispatch(file->sftp) < 0) {
+    if (rc != SSH_OK) {
         return NULL;
     }
-        msg = sftp_dequeue(file->sftp, id);
-    }
 
     if (msg->packet_type == SSH_FXP_ATTRS){
         sftp_attributes attr = sftp_parse_attr(file->sftp, msg->payload, 0);
@@ -3146,13 +3085,10 @@ char *sftp_expand_path(sftp_session sftp, const char *path)
         return NULL;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        rc = sftp_read_and_dispatch(sftp);
+    if (rc != SSH_OK) {
-        if (rc < 0) {
         return NULL;
     }
-        msg = sftp_dequeue(sftp, id);
-    }
 
     if (msg->packet_type == SSH_FXP_NAME) {
         uint32_t ignored = 0;
@@ -3233,13 +3169,10 @@ sftp_home_directory(sftp_session sftp, const char *username)
         return NULL;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, id, true, &msg);
-        rc = sftp_read_and_dispatch(sftp);
+    if (rc != SSH_OK) {
-        if (rc < 0) {
         return NULL;
     }
-        msg = sftp_dequeue(sftp, id);
-    }
 
     if (msg->packet_type == SSH_FXP_NAME) {
         uint32_t count = 0;

src/sftp_aio.c

@@ -202,21 +202,15 @@ ssize_t sftp_aio_wait_read(sftp_aio *aio,
     }
 
     /* handle an existing request */
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, (*aio)->id, !file->nonblocking, &msg);
-        if (file->nonblocking) {
+    if (rc == SSH_ERROR) {
-            if (ssh_channel_poll(sftp->channel, 0) == 0) {
-                /* we cannot block */
-                return SSH_AGAIN;
-            }
-        }
-
-        if (sftp_read_and_dispatch(sftp) < 0) {
-            /* something nasty has happened */
         SFTP_AIO_FREE(*aio);
         return SSH_ERROR;
     }
 
-        msg = sftp_dequeue(sftp, (*aio)->id);
+    if (rc == SSH_AGAIN) {
+        /* return without freeing the (*aio) */
+        return SSH_AGAIN;
     }
 
     /*
@@ -410,6 +404,7 @@ ssize_t sftp_aio_wait_write(sftp_aio *aio)
     sftp_session sftp = NULL;
     sftp_message msg = NULL;
     sftp_status_message status = NULL;
+    int rc;
 
     /*
      * This function releases the memory of the structure
@@ -446,21 +441,15 @@ ssize_t sftp_aio_wait_write(sftp_aio *aio)
         return SSH_ERROR;
     }
 
-    while (msg == NULL) {
+    rc = sftp_recv_response_msg(sftp, (*aio)->id, !file->nonblocking, &msg);
-        if (file->nonblocking) {
+    if (rc == SSH_ERROR) {
-            if (ssh_channel_poll(sftp->channel, 0) == 0) {
-                /* we cannot block */
-                return SSH_AGAIN;
-            }
-        }
-
-        if (sftp_read_and_dispatch(sftp) < 0) {
-            /* something nasty has happened */
         SFTP_AIO_FREE(*aio);
         return SSH_ERROR;
     }
 
-        msg = sftp_dequeue(sftp, (*aio)->id);
+    if (rc == SSH_AGAIN) {
+        /* Return without freeing the (*aio) */
+        return SSH_AGAIN;
     }
 
     /*

src/sftp_common.c

@@ -859,6 +859,56 @@ int sftp_read_and_dispatch(sftp_session sftp)
     return 0;
 }
 
+int sftp_recv_response_msg(sftp_session sftp,
+                           uint32_t id,
+                           bool blocking,
+                           sftp_message *msg_ptr)
+{
+    sftp_message msg = NULL;
+    int rc;
+
+    if (sftp == NULL) {
+        return SSH_ERROR;
+    }
+
+    if (msg_ptr == NULL) {
+        ssh_set_error_invalid(sftp->session);
+        sftp_set_error(sftp, SSH_FX_FAILURE);
+        return SSH_ERROR;
+    }
+
+    SSH_LOG(SSH_LOG_PACKET,
+            "Trying to receive response of request id %" PRIu32 " in %s mode",
+            id,
+            blocking ? "blocking" : "non-blocking");
+
+    do {
+        if (!blocking) {
+            rc = ssh_channel_poll(sftp->channel, 0);
+            if (rc == SSH_ERROR) {
+                sftp_set_error(sftp, SSH_FX_FAILURE);
+                return SSH_ERROR;
+            }
+
+            if (rc == 0) {
+                /* nothing available and we cannot block */
+                return SSH_AGAIN;
+            }
+        }
+
+        rc = sftp_read_and_dispatch(sftp);
+        if (rc == -1) {
+            /* something nasty has happened */
+            return SSH_ERROR;
+        }
+
+        msg = sftp_dequeue(sftp, id);
+    } while (msg == NULL);
+
+    *msg_ptr = msg;
+    return SSH_OK;
+}
+
 sftp_status_message parse_status_msg(sftp_message msg)
 {
     sftp_status_message status = NULL;