[PATCH 1/1] buffer: add and use ssh_buffer_ensure_allocated

[Pino Toscano <ptoscano@xxxxxxxxxx>]

Add a small helper for ssh_buffer to ensure that the buffer has a
certain amount of space already preallocated. This can be useful in case
it is known how much data is going to be added to a buffer, to avoid
multiple reallocations.

Make use of it in few places in the library.

Signed-off-by: Pino Toscano <ptoscano@xxxxxxxxxx>
---
 include/libssh/buffer.h |   1 +
 src/agent.c             |  14 +++++
 src/buffer.c            |  25 +++++++++
 src/dh.c                |  12 ++++
 src/pcap.c              |  12 ++++
 src/sftp.c              | 119 ++++++++++++++++++++++++++++++++++++++++
 6 files changed, 183 insertions(+)

diff --git a/include/libssh/buffer.h b/include/libssh/buffer.h
index 14ce5e67..b3e38d7a 100644
--- a/include/libssh/buffer.h
+++ b/include/libssh/buffer.h
@@ -47,6 +47,7 @@ int ssh_buffer_add_u8(ssh_buffer buffer, uint8_t data);
 int ssh_buffer_add_u16(ssh_buffer buffer, uint16_t data);
 int ssh_buffer_add_u32(ssh_buffer buffer, uint32_t data);
 int ssh_buffer_add_u64(ssh_buffer buffer, uint64_t data);
+int ssh_buffer_ensure_allocated(ssh_buffer buffer, uint32_t len);
 
 int ssh_buffer_validate_length(struct ssh_buffer_struct *buffer, size_t len);
 
diff --git a/src/agent.c b/src/agent.c
index 0b145ff3..c19e3b31 100644
--- a/src/agent.c
+++ b/src/agent.c
@@ -536,6 +536,20 @@ ssh_string ssh_agent_sign_data(ssh_session session,
         return NULL;
     }
 
+    /*
+     * make sure it already can contain all the expected content:
+     * - 1 x uint8_t
+     * - 2 x uint32_t
+     * - 1 x ssh_string (uint8_t + data)
+     */
+    if (ssh_buffer_ensure_allocated(request,
+                                    sizeof(uint8_t) * 2 +
+                                    sizeof(uint32_t) * 2 +
+                                    ssh_string_len(key_blob)) < 0) {
+        ssh_buffer_free(request);
+        return NULL;
+    }
+
     /* adds len + blob */
     rc = ssh_buffer_add_ssh_string(request, key_blob);
     ssh_string_free(key_blob);
diff --git a/src/buffer.c b/src/buffer.c
index 5859e926..f5607aac 100644
--- a/src/buffer.c
+++ b/src/buffer.c
@@ -248,6 +248,31 @@ int ssh_buffer_add_data(struct ssh_buffer_struct *buffer, const void *data, uint
   return 0;
 }
 
+/**
+ * @brief Ensure the buffer has at least a certain preallocated size.
+ *
+ * @param[in]  buffer   The buffer to enlarge.
+ *
+ * @param[in]  len      The length to ensure as allocated.
+ *
+ * @return              0 on success, < 0 on error.
+ */
+int ssh_buffer_ensure_allocated(struct ssh_buffer_struct *buffer, uint32_t len)
+{
+  buffer_verify(buffer);
+
+  if (buffer->allocated < len) {
+    if(buffer->pos > 0)
+      buffer_shift(buffer);
+    if (realloc_buffer(buffer, len) < 0) {
+      return -1;
+    }
+  }
+
+  buffer_verify(buffer);
+  return 0;
+}
+
 /**
  * @internal
  *
diff --git a/src/dh.c b/src/dh.c
index d2ddfabd..8581f756 100644
--- a/src/dh.c
+++ b/src/dh.c
@@ -791,6 +791,12 @@ int ssh_hashbufout_add_cookie(ssh_session session) {
     return -1;
   }
 
+  if (ssh_buffer_ensure_allocated(session->out_hashbuf,
+                                  sizeof(uint8_t) + 16) < 0) {
+    ssh_buffer_reinit(session->out_hashbuf);
+    return -1;
+  }
+
   if (ssh_buffer_add_u8(session->out_hashbuf, 20) < 0) {
     ssh_buffer_reinit(session->out_hashbuf);
     return -1;
@@ -819,6 +825,12 @@ int ssh_hashbufin_add_cookie(ssh_session session, unsigned char *cookie) {
     return -1;
   }
 
+  if (ssh_buffer_ensure_allocated(session->in_hashbuf,
+                                  sizeof(uint8_t) + 20 + 16) < 0) {
+    ssh_buffer_reinit(session->in_hashbuf);
+    return -1;
+  }
+
   if (ssh_buffer_add_u8(session->in_hashbuf, 20) < 0) {
     ssh_buffer_reinit(session->in_hashbuf);
     return -1;
diff --git a/src/pcap.c b/src/pcap.c
index bfb237ba..81bf0fd8 100644
--- a/src/pcap.c
+++ b/src/pcap.c
@@ -165,6 +165,12 @@ int ssh_pcap_file_write_packet(ssh_pcap_file pcap, ssh_buffer packet, uint32_t o
 	if(header == NULL)
 		return SSH_ERROR;
 	gettimeofday(&now,NULL);
+    err = ssh_buffer_ensure_allocated(header,
+                                      sizeof(uint32_t) * 4 +
+                                      ssh_buffer_get_len(packet));
+    if (err < 0) {
+        goto error;
+    }
     err = ssh_buffer_add_u32(header,htonl(now.tv_sec));
     if (err < 0) {
         goto error;
@@ -209,6 +215,12 @@ int ssh_pcap_file_open(ssh_pcap_file pcap, const char *filename){
 	header=ssh_buffer_new();
 	if(header==NULL)
 		return SSH_ERROR;
+    err = ssh_buffer_ensure_allocated(header,
+                                      sizeof(uint32_t) * 5 +
+                                      sizeof(uint16_t) * 2);
+    if (err < 0) {
+        goto error;
+    }
     err = ssh_buffer_add_u32(header,htonl(PCAP_MAGIC));
     if (err < 0) {
         goto error;
diff --git a/src/sftp.c b/src/sftp.c
index 80870d4e..490001c1 100644
--- a/src/sftp.c
+++ b/src/sftp.c
@@ -318,6 +318,10 @@ sftp_packet sftp_packet_read(sftp_session sftp) {
     return NULL;
   }
 
+  if (ssh_buffer_ensure_allocated(packet->payload, 4) < 0) {
+    ssh_set_error_oom(sftp->session);
+    goto error;
+  }
   r=0;
   do {
     // read from channel until 4 bytes have been read or an error occurs
@@ -355,6 +359,10 @@ sftp_packet sftp_packet_read(sftp_session sftp) {
   }
   size--;
 
+  if (ssh_buffer_ensure_allocated(packet->payload, size) < 0) {
+    ssh_set_error_oom(sftp->session);
+    goto error;
+  }
   while (size > 0 && size < UINT_MAX) {
     r=ssh_channel_read(sftp->channel,buffer,
         sizeof(buffer)>size ? size:sizeof(buffer),0);
@@ -864,6 +872,14 @@ sftp_dir sftp_opendir(sftp_session sftp, const char *path){
   }
 
   id = sftp_get_new_id(sftp);
+  if (ssh_buffer_ensure_allocated(payload,
+                                  sizeof(uint32_t) * 2 +
+                                  ssh_string_len(path_s)) < 0) {
+    ssh_set_error_oom(sftp->session);
+    ssh_buffer_free(payload);
+    ssh_string_free(path_s);
+    return NULL;
+  }
   if (ssh_buffer_add_u32(payload, htonl(id)) < 0 ||
       ssh_buffer_add_ssh_string(payload, path_s) < 0) {
     ssh_set_error_oom(sftp->session);
@@ -1392,6 +1408,13 @@ sftp_attributes sftp_readdir(sftp_session sftp, sftp_dir dir) {
       return NULL;
     }
 
+    if (ssh_buffer_ensure_allocated(payload,
+                                    sizeof(uint32_t) * 2 +
+                                    ssh_string_len(dir->handle)) < 0) {
+      ssh_set_error_oom(sftp->session);
+      ssh_buffer_free(payload);
+      return NULL;
+    }
     id = sftp_get_new_id(sftp);
     if (ssh_buffer_add_u32(payload, htonl(id)) < 0 ||
         ssh_buffer_add_ssh_string(payload, dir->handle) < 0) {
@@ -1516,6 +1539,14 @@ static int sftp_handle_close(sftp_session sftp, ssh_string handle) {
     return -1;
   }
 
+  if (ssh_buffer_ensure_allocated(buffer,
+                                  sizeof(uint32_t) * 2 +
+                                  ssh_string_len(handle)) < 0) {
+    ssh_set_error_oom(sftp->session);
+    ssh_buffer_free(buffer);
+    return -1;
+  }
+
   id = sftp_get_new_id(sftp);
   if (ssh_buffer_add_u32(buffer, htonl(id)) < 0 ||
       ssh_buffer_add_ssh_string(buffer, handle) < 0) {
@@ -1644,6 +1675,14 @@ sftp_file sftp_open(sftp_session sftp, const char *file, int flags,
       sftp_flags |= SSH_FXF_APPEND;
   }
   SSH_LOG(SSH_LOG_PACKET,"Opening file %s with sftp flags %x",file,sftp_flags);
+  if (ssh_buffer_ensure_allocated(buffer,
+                                  sizeof(uint32_t) * 4 +
+                                  ssh_string_len(filename)) < 0) {
+    ssh_set_error_oom(sftp->session);
+    ssh_buffer_free(buffer);
+    ssh_string_free(filename);
+    return NULL;
+  }
   id = sftp_get_new_id(sftp);
   if (ssh_buffer_add_u32(buffer, htonl(id)) < 0 ||
       ssh_buffer_add_ssh_string(buffer, filename) < 0) {
@@ -2233,6 +2272,15 @@ int sftp_mkdir(sftp_session sftp, const char *directory, mode_t mode) {
   attr.permissions = mode;
   attr.flags = SSH_FILEXFER_ATTR_PERMISSIONS;
 
+  if (ssh_buffer_ensure_allocated(buffer,
+                                  sizeof(uint32_t) * 2 +
+                                  ssh_string_len(path)) < 0) {
+    ssh_set_error_oom(sftp->session);
+    ssh_buffer_free(buffer);
+    ssh_string_free(path);
+    return -1;
+  }
+
   id = sftp_get_new_id(sftp);
   if (ssh_buffer_add_u32(buffer, htonl(id)) < 0 ||
       ssh_buffer_add_ssh_string(buffer, path) < 0 ||
@@ -2399,6 +2447,15 @@ int sftp_setstat(sftp_session sftp, const char *file, sftp_attributes attr) {
     return -1;
   }
 
+  if (ssh_buffer_ensure_allocated(buffer,
+                                  sizeof(uint32_t) * 2 +
+                                  ssh_string_len(path)) < 0) {
+    ssh_set_error_oom(sftp->session);
+    ssh_buffer_free(buffer);
+    ssh_string_free(path);
+    return -1;
+  }
+
   id = sftp_get_new_id(sftp);
   if (ssh_buffer_add_u32(buffer, htonl(id)) < 0 ||
       ssh_buffer_add_ssh_string(buffer, path) < 0 ||
@@ -2616,6 +2673,15 @@ char *sftp_readlink(sftp_session sftp, const char *path) {
     return NULL;
   }
 
+  if (ssh_buffer_ensure_allocated(buffer,
+                                  sizeof(uint32_t) * 2 +
+                                  ssh_string_len(path_s)) < 0) {
+    ssh_set_error_oom(sftp->session);
+    ssh_buffer_free(buffer);
+    ssh_string_free(path_s);
+    return NULL;
+  }
+
   id = sftp_get_new_id(sftp);
   if (ssh_buffer_add_u32(buffer, htonl(id)) < 0 ||
       ssh_buffer_add_ssh_string(buffer, path_s) < 0) {
@@ -2743,6 +2809,16 @@ sftp_statvfs_t sftp_statvfs(sftp_session sftp, const char *path) {
     return NULL;
   }
 
+  if (ssh_buffer_ensure_allocated(buffer,
+                                  sizeof(uint32_t) * 3 + ssh_string_len(ext) +
+                                  ssh_string_len(pathstr)) < 0) {
+    ssh_set_error_oom(sftp->session);
+    ssh_buffer_free(buffer);
+    ssh_string_free(ext);
+    ssh_string_free(pathstr);
+    return NULL;
+  }
+
   id = sftp_get_new_id(sftp);
   if (ssh_buffer_add_u32(buffer, htonl(id)) < 0 ||
       ssh_buffer_add_ssh_string(buffer, ext) < 0 ||
@@ -2823,6 +2899,14 @@ int sftp_fsync(sftp_file file)
         goto done;
     }
 
+    rc = ssh_buffer_ensure_allocated(buffer,
+                                  sizeof(uint32_t) * 3 + ssh_string_len(ext) +
+                                  ssh_string_len(file->handle));
+    if (rc < 0) {
+        ssh_set_error_oom(sftp->session);
+        goto done;
+    }
+
     id = sftp_get_new_id(sftp);
     rc = ssh_buffer_add_u32(buffer, htonl(id));
     if (rc < 0) {
@@ -2934,6 +3018,15 @@ sftp_statvfs_t sftp_fstatvfs(sftp_file file) {
     return NULL;
   }
 
+  if (ssh_buffer_ensure_allocated(buffer,
+                                  sizeof(uint32_t) * 3 + ssh_string_len(ext) +
+                                  ssh_string_len(file->handle)) < 0) {
+    ssh_set_error_oom(sftp->session);
+    ssh_buffer_free(buffer);
+    ssh_string_free(ext);
+    return NULL;
+  }
+
   id = sftp_get_new_id(sftp);
   if (ssh_buffer_add_u32(buffer, htonl(id)) < 0 ||
       ssh_buffer_add_ssh_string(buffer, ext) < 0 ||
@@ -3023,6 +3116,15 @@ char *sftp_canonicalize_path(sftp_session sftp, const char *path) {
     return NULL;
   }
 
+  if (ssh_buffer_ensure_allocated(buffer,
+                                  sizeof(uint32_t) * 2 +
+                                  ssh_string_len(pathstr)) < 0) {
+    ssh_set_error_oom(sftp->session);
+    ssh_buffer_free(buffer);
+    ssh_string_free(pathstr);
+    return NULL;
+  }
+
   id = sftp_get_new_id(sftp);
   if (ssh_buffer_add_u32(buffer, htonl(id)) < 0 ||
       ssh_buffer_add_ssh_string(buffer, pathstr) < 0) {
@@ -3106,6 +3208,15 @@ static sftp_attributes sftp_xstat(sftp_session sftp, const char *path,
     return NULL;
   }
 
+  if (ssh_buffer_ensure_allocated(buffer,
+                                  sizeof(uint32_t) * 2 +
+                                  ssh_string_len(pathstr)) < 0) {
+    ssh_set_error_oom(sftp->session);
+    ssh_buffer_free(buffer);
+    ssh_string_free(pathstr);
+    return NULL;
+  }
+
   id = sftp_get_new_id(sftp);
   if (ssh_buffer_add_u32(buffer, htonl(id)) < 0 ||
       ssh_buffer_add_ssh_string(buffer, pathstr) < 0) {
@@ -3173,6 +3284,14 @@ sftp_attributes sftp_fstat(sftp_file file) {
     return NULL;
   }
 
+  if (ssh_buffer_ensure_allocated(buffer,
+                                  sizeof(uint32_t) * 2 +
+                                  ssh_string_len(file->handle)) < 0) {
+    ssh_set_error_oom(file->sftp->session);
+    ssh_buffer_free(buffer);
+    return NULL;
+  }
+
   id = sftp_get_new_id(file->sftp);
   if (ssh_buffer_add_u32(buffer, htonl(id)) < 0 ||
       ssh_buffer_add_ssh_string(buffer, file->handle) < 0) {
-- 
2.17.1