diff --git a/src/buffer.c b/src/buffer.c index 9e7d267e..8cc01c74 100644 --- a/src/buffer.c +++ b/src/buffer.c @@ -786,6 +786,122 @@ struct ssh_string_struct *ssh_buffer_get_ssh_string(struct ssh_buffer_struct *bu return str; } +/** + * @brief Pre-calculate the size we need for packing the buffer. + * + * This makes sure that enough memory is allocated for packing the buffer and + * we only have to do one memory allocation. + * + * @param[in] buffer The buffer to allocate + * + * @param[in] format A format string of arguments. + * + * @param[in] argc The number of arguments. + * + * @param[in] ap The va_list of arguments. + * + * @return SSH_OK on success, SSH_ERROR on error. + */ +static int ssh_buffer_pack_allocate_va(struct ssh_buffer_struct *buffer, + const char *format, + int argc, + va_list ap) +{ + const char *p = NULL; + ssh_string string = NULL; + char *cstring = NULL; + size_t needed_size = 0; + size_t count; + size_t len; + int rc = SSH_OK; + + for (p = format, count = 0; *p != '\0'; p++, count++) { + /* Invalid number of arguments passed */ + if (argc != -1 && count > argc) { + return SSH_ERROR; + } + + switch(*p) { + case 'b': + va_arg(ap, unsigned int); + needed_size += sizeof(uint8_t); + break; + case 'w': + va_arg(ap, unsigned int); + needed_size += sizeof(uint16_t); + break; + case 'd': + va_arg(ap, uint32_t); + needed_size += sizeof(uint32_t); + break; + case 'q': + va_arg(ap, uint64_t); + needed_size += sizeof(uint64_t); + break; + case 'S': + string = va_arg(ap, ssh_string); + needed_size += 4 + ssh_string_len(string); + string = NULL; + break; + case 's': + cstring = va_arg(ap, char *); + needed_size += sizeof(uint32_t) + strlen(cstring); + cstring = NULL; + break; + case 'P': + len = va_arg(ap, size_t); + needed_size += len; + va_arg(ap, void *); + count++; /* increase argument count */ + break; + case 'B': + va_arg(ap, bignum); + /* + * Use a fixed size for a bignum + * (they should normaly be around 32) + */ + needed_size += 64; + break; + case 't': + cstring = va_arg(ap, char *); + needed_size += strlen(cstring); + cstring = NULL; + break; + default: + SSH_LOG(SSH_LOG_WARN, "Invalid buffer format %c", *p); + rc = SSH_ERROR; + } + if (rc != SSH_OK){ + break; + } + } + + if (argc != -1 && argc != count) { + return SSH_ERROR; + } + + if (rc != SSH_ERROR){ + /* + * Check if our canary is intact, if not, something really bad happened. + */ + uint32_t canary = va_arg(ap, uint32_t); + if (canary != SSH_BUFFER_PACK_END) { + if (argc == -1){ + return SSH_ERROR; + } else { + abort(); + } + } + } + + rc = ssh_buffer_allocate_size(buffer, needed_size); + if (rc != 0) { + return SSH_ERROR; + } + + return SSH_OK; +} + /** @internal * @brief Add multiple values in a buffer on a single function call * @param[in] buffer The buffer to add to @@ -935,9 +1051,18 @@ int _ssh_buffer_pack(struct ssh_buffer_struct *buffer, va_list ap; int rc; + va_start(ap, argc); + rc = ssh_buffer_pack_allocate_va(buffer, format, argc, ap); + va_end(ap); + + if (rc != SSH_OK) { + return rc; + } + va_start(ap, argc); rc = ssh_buffer_pack_va(buffer, format, argc, ap); va_end(ap); + return rc; }