Buffer: add ssh_buffer_(un)pack()

That function permits chaining of buffer values to minimize buffer handling
in packet sending code.

Reviewed-by: Andreas Schneider <asn@cryptomilk.org>

2 files changed, 301 insertions, 0 deletions

diff --git a/include/libssh/buffer.h b/include/libssh/buffer.h
index 3224af28..942089f0 100644
--- a/include/libssh/buffer.h
+++ b/include/libssh/buffer.h
@@ -21,6 +21,8 @@
 #ifndef BUFFER_H_
 #define BUFFER_H_
 
+#include <stdarg.h>
+
 #include "libssh/libssh.h"
 /*
  * Describes a buffer state
@@ -46,6 +48,10 @@ int buffer_add_u16(ssh_buffer buffer, uint16_t data);
 int buffer_add_u32(ssh_buffer buffer, uint32_t data);
 int buffer_add_u64(ssh_buffer buffer, uint64_t data);
 int ssh_buffer_add_data(ssh_buffer buffer, const void *data, uint32_t len);
+int ssh_buffer_pack_va(struct ssh_buffer_struct *buffer, const char *format, va_list ap);
+int ssh_buffer_pack(struct ssh_buffer_struct *buffer, const char *format, ...);
+int ssh_buffer_unpack_va(struct ssh_buffer_struct *buffer, const char *format, va_list ap);
+int ssh_buffer_unpack(struct ssh_buffer_struct *buffer, const char *format, ...);
 int buffer_prepend_data(ssh_buffer buffer, const void *data, uint32_t len);
 int buffer_add_buffer(ssh_buffer buffer, ssh_buffer source);
 int ssh_buffer_reinit(ssh_buffer buffer);
diff --git a/src/buffer.c b/src/buffer.c
index fdee6a09..c2879b4b 100644
--- a/src/buffer.c
+++ b/src/buffer.c
@@ -24,6 +24,7 @@
 #include <limits.h>
 #include <stdlib.h>
 #include <string.h>
+#include <stdarg.h>
 
 #ifndef _WIN32
 #include <netinet/in.h>
@@ -32,6 +33,8 @@
 
 #include "libssh/priv.h"
 #include "libssh/buffer.h"
+#include "libssh/misc.h"
+#include "libssh/dh.h"
 
 /**
  * @defgroup libssh_buffer The SSH buffer functions.
@@ -645,6 +648,298 @@ struct ssh_string_struct *buffer_get_mpint(struct ssh_buffer_struct *buffer) {
   return str;
 }
 
+/** @internal
+ * @brief Add multiple values in a buffer on a single function call
+ * @param[in] buffer    The buffer to add to
+ * @param[in] format    A format string of arguments.
+ * @param[in] ap        A va_list of arguments.
+ * @returns             SSH_OK on success
+ *                      SSH_ERROR on error
+ * @see ssh_buffer_add_format() for format list values.
+ */
+int ssh_buffer_pack_va(struct ssh_buffer_struct *buffer, const char *format, va_list ap){
+    int rc;
+    const char *p;
+    union {
+        uint8_t byte;
+        uint16_t word;
+        uint32_t dword;
+        uint64_t qword;
+        ssh_string string;
+        void *data;
+    } o;
+    char *cstring;
+    bignum b;
+    size_t len;
+
+    for (p = format; *p != '\0'; p++) {
+        switch(*p) {
+        case 'b':
+            o.byte = (uint8_t)va_arg(ap, unsigned int);
+            rc = buffer_add_u8(buffer, o.byte);
+            break;
+        case 'w':
+            o.word = (uint16_t)va_arg(ap, unsigned int);
+            o.word = htons(o.word);
+            rc = buffer_add_u16(buffer, o.word);
+            break;
+        case 'd':
+            o.dword = va_arg(ap, uint32_t);
+            o.dword = htonl(o.dword);
+            rc = buffer_add_u32(buffer, o.dword);
+            break;
+        case 'q':
+            o.qword = va_arg(ap, uint64_t);
+            o.qword = htonll(o.qword);
+            rc = buffer_add_u64(buffer, o.qword);
+            break;
+        case 'S':
+            o.string = va_arg(ap, ssh_string);
+            rc = buffer_add_ssh_string(buffer, o.string);
+            o.string = NULL;
+            break;
+        case 's':
+            cstring = va_arg(ap, char *);
+            len = strlen(cstring);
+            rc = buffer_add_u32(buffer, htonl(len));
+            if (rc == SSH_OK){
+                rc = ssh_buffer_add_data(buffer, cstring, len);
+            }
+            cstring = NULL;
+            break;
+        case 'P':
+            len = va_arg(ap, size_t);
+            o.data = va_arg(ap, void *);
+            rc = ssh_buffer_add_data(buffer, o.data, len);
+            o.data = NULL;
+            break;
+        case 'B':
+            b = va_arg(ap, bignum);
+            o.string = make_bignum_string(b);
+            if(o.string == NULL){
+                rc = SSH_ERROR;
+                break;
+            }
+            rc = buffer_add_ssh_string(buffer, o.string);
+            SAFE_FREE(o.string);
+            break;
+        default:
+            SSH_LOG(SSH_LOG_WARN, "Invalid buffer format %c", *p);
+            rc = SSH_ERROR;
+        }
+        if (rc != SSH_OK){
+            break;
+        }
+    }
+
+    return rc;
+}
+
+/** @internal
+ * @brief Add multiple values in a buffer on a single function call
+ * @param[in] buffer    The buffer to add to
+ * @param[in] format    A format string of arguments. This string contains single
+ *                      letters describing the order and type of arguments:
+ *                         'b': uint8_t  (pushed in network byte order)
+ *                         'w': uint16_t (pushed in network byte order)
+ *                         'd': uint32_t (pushed in network byte order)
+ *                         'q': uint64_t (pushed in network byte order)
+ *                         'S': ssh_string
+ *                         's': char * (C string, pushed as SSH string)
+ *                         'P': size_t, void * (len of data, pointer to data)
+ *                              only pushes data.
+ *                         'S': bignum (pushed as SSH string)
+ * @returns             SSH_OK on success
+ *                      SSH_ERROR on error
+ * @warning             when using 'P' with a constant size (e.g. 8), do not
+ *                      forget to cast to (size_t).
+ */
+int ssh_buffer_pack(struct ssh_buffer_struct *buffer, const char *format, ...){
+    va_list ap;
+    int rc;
+
+    va_start(ap, format);
+    rc = ssh_buffer_pack_va(buffer, format, ap);
+    va_end(ap);
+    return rc;
+}
+
+/** @internal
+ * @brief Get multiple values from a buffer on a single function call
+ * @param[in] buffer    The buffer to get from
+ * @param[in] format    A format string of arguments.
+ * @param[in] ap        A va_list of arguments.
+ * @returns             SSH_OK on success
+ *                      SSH_ERROR on error
+ * @see ssh_buffer_get_format() for format list values.
+ */
+int ssh_buffer_unpack_va(struct ssh_buffer_struct *buffer, const char *format, va_list ap){
+    int rc;
+    const char *p, *last;
+    union {
+        uint8_t *byte;
+        uint16_t *word;
+        uint32_t *dword;
+        uint64_t *qword;
+        ssh_string *string;
+        char **cstring;
+        void **data;
+    } o;
+    size_t len, rlen;
+    uint32_t u32len;
+    va_list ap_copy;
+
+    /* copy the argument list in case a rollback is needed */
+    va_copy(ap_copy, ap);
+
+    for (p = format; *p != '\0'; p++) {
+        switch (*p) {
+        case 'b':
+            o.byte = va_arg(ap, uint8_t *);
+            rlen = buffer_get_u8(buffer, o.byte);
+            rc = rlen==1 ? SSH_OK : SSH_ERROR;
+            break;
+        case 'w':
+            o.word = va_arg(ap,  uint16_t *);
+            rlen = buffer_get_data(buffer, o.word, sizeof(uint16_t));
+            *o.word = ntohs(*o.word);
+            rc = rlen==2 ? SSH_OK : SSH_ERROR;
+            break;
+        case 'd':
+            o.dword = va_arg(ap, uint32_t *);
+            rlen = buffer_get_u32(buffer, o.dword);
+            *o.dword = ntohl(*o.dword);
+            rc = rlen==4 ? SSH_OK : SSH_ERROR;
+            break;
+        case 'q':
+            o.qword = va_arg(ap, uint64_t*);
+            rlen = buffer_get_u64(buffer, o.qword);
+            *o.qword = ntohll(*o.qword);
+            rc = rlen==8 ? SSH_OK : SSH_ERROR;
+            break;
+        case 'S':
+            o.string = va_arg(ap, ssh_string *);
+            *o.string = buffer_get_ssh_string(buffer);
+            rc = *o.string != NULL ? SSH_OK : SSH_ERROR;
+            o.string = NULL;
+            break;
+        case 's':
+            o.cstring = va_arg(ap, char **);
+            *o.cstring = NULL;
+            rc = buffer_get_u32(buffer, &u32len);
+            if (rc != 4){
+                rc = SSH_ERROR;
+                break;
+            }
+            len = ntohl(u32len);
+            if (len > UINT_MAX - 1){
+                rc = SSH_ERROR;
+                break;
+            }
+            *o.cstring = malloc(len + 1);
+            if (*o.cstring == NULL){
+                rc = SSH_ERROR;
+                break;
+            }
+            rlen = buffer_get_data(buffer, *o.cstring, len);
+            if (rlen != len){
+                SAFE_FREE(*o.cstring);
+                rc = SSH_ERROR;
+                break;
+            }
+            (*o.cstring)[len] = '\0';
+            o.cstring = NULL;
+            rc = SSH_OK;
+            break;
+        case 'P':
+            len = va_arg(ap, size_t);
+            o.data = va_arg(ap, void **);
+            *o.data = malloc(len);
+            if(*o.data == NULL){
+                rc = SSH_ERROR;
+                break;
+            }
+            rlen = buffer_get_data(buffer, *o.data, len);
+            if (rlen != len){
+                SAFE_FREE(*o.data);
+                rc = SSH_ERROR;
+                break;
+            }
+            o.data = NULL;
+            rc = SSH_OK;
+            break;
+        default:
+            SSH_LOG(SSH_LOG_WARN, "Invalid buffer format %c", *p);
+            rc = SSH_ERROR;
+        }
+        if (rc != SSH_OK) {
+            break;
+        }
+    }
+
+    if (rc != SSH_OK){
+        /* Reset the format string and erase everything that was allocated */
+        last = p;
+        for(p=format;p<last;++p){
+            switch(*p){
+            case 'b':
+            case 'w':
+            case 'd':
+            case 'q':
+                va_arg(ap_copy, void *);
+                break;
+            case 'S':
+                o.string=va_arg(ap_copy, ssh_string *);
+                SAFE_FREE(*o.string);
+                break;
+            case 's':
+                o.cstring=va_arg(ap_copy, char **);
+                SAFE_FREE(*o.cstring);
+                break;
+            case 'P':
+                va_arg(ap_copy, size_t);
+                o.data = va_arg(ap_copy, void **);
+                SAFE_FREE(*o.data);
+                break;
+            default:
+                va_arg(ap_copy, void *);
+                break;
+            }
+        }
+    }
+    va_end(ap_copy);
+
+    return rc;
+}
+
+/** @internal
+ * @brief Get multiple values from a buffer on a single function call
+ * @param[in] buffer    The buffer to get from
+ * @param[in] format    A format string of arguments. This string contains single
+ *                      letters describing the order and type of arguments:
+ *                         'b': uint8_t *  (pulled in network byte order)
+ *                         'w': uint16_t * (pulled in network byte order)
+ *                         'd': uint32_t * (pulled in network byte order)
+ *                         'q': uint64_t * (pulled in network byte order)
+ *                         'S': ssh_string *
+ *                         's': char ** (C string, pulled as SSH string)
+ *                         'P': size_t, void ** (len of data, pointer to data)
+ *                              only pulls data.
+ * @returns             SSH_OK on success
+ *                      SSH_ERROR on error
+ * @warning             when using 'P' with a constant size (e.g. 8), do not
+ *                      forget to cast to (size_t).
+ */
+int ssh_buffer_unpack(struct ssh_buffer_struct *buffer, const char *format, ...){
+    va_list ap;
+    int rc;
+
+    va_start(ap, format);
+    rc = ssh_buffer_unpack_va(buffer, format, ap);
+    va_end(ap);
+    return rc;
+}
+
 /** @} */
 
 /* vim: set ts=4 sw=4 et cindent: */