Merge branch 'buf_for_stringbuffer_squashed'

7 files changed, 197 insertions, 79 deletions

diff --git a/changes/ticket22342 b/changes/ticket22342
new file mode 100644
index 0000000000..53505509d2
--- /dev/null
+++ b/changes/ticket22342
@@ -0,0 +1,3 @@
+  o Code simplification and refactoring:
+    - Small changes to Tor's buf_t API to make it suitable for use as
+      a general-purpose safe string constructor. Closes ticket 22342.
diff --git a/src/common/buffers.c b/src/common/buffers.c
index c45e13d551..a21cf2ab5b 100644
--- a/src/common/buffers.c
+++ b/src/common/buffers.c
@@ -714,6 +714,53 @@ buf_add(buf_t *buf, const char *string, size_t string_len)
   return (int)buf->datalen;
 }
 
+/** Add a nul-terminated <b>string</b> to <b>buf</b>, not including the
+ * terminating NUL. */
+void
+buf_add_string(buf_t *buf, const char *string)
+{
+  buf_add(buf, string, strlen(string));
+}
+
+/** As tor_snprintf, but write the results into a buf_t */
+void
+buf_add_printf(buf_t *buf, const char *format, ...)
+{
+  va_list ap;
+  va_start(ap,format);
+  buf_add_vprintf(buf, format, ap);
+  va_end(ap);
+}
+
+/** As tor_vsnprintf, but write the results into a buf_t. */
+void
+buf_add_vprintf(buf_t *buf, const char *format, va_list args)
+{
+  /* XXXX Faster implementations are easy enough, but let's optimize later */
+  char *tmp;
+  tor_vasprintf(&tmp, format, args);
+  buf_add(buf, tmp, strlen(tmp));
+  tor_free(tmp);
+}
+
+/** Return a heap-allocated string containing the contents of <b>buf</b>, plus
+ * a NUL byte. If <b>sz_out</b> is provided, set *<b>sz_out</b> to the length
+ * of the returned string, not including the terminating NUL. */
+char *
+buf_extract(buf_t *buf, size_t *sz_out)
+{
+  tor_assert(buf);
+
+  size_t sz = buf_datalen(buf);
+  char *result;
+  result = tor_malloc(sz+1);
+  buf_peek(buf, result, sz);
+  result[sz] = 0;
+  if (sz_out)
+    *sz_out = sz;
+  return result;
+}
+
 /** Helper: copy the first <b>string_len</b> bytes from <b>buf</b>
  * onto <b>string</b>.
  */
@@ -795,6 +842,28 @@ buf_move_to_buf(buf_t *buf_out, buf_t *buf_in, size_t *buf_flushlen)
   return (int)cp;
 }
 
+/** Moves all data from <b>buf_in</b> to <b>buf_out</b>, without copying.
+ */
+void
+buf_move_all(buf_t *buf_out, buf_t *buf_in)
+{
+  tor_assert(buf_out);
+  if (!buf_in)
+    return;
+
+  if (buf_out->head == NULL) {
+    buf_out->head = buf_in->head;
+    buf_out->tail = buf_in->tail;
+  } else {
+    buf_out->tail->next = buf_in->head;
+    buf_out->tail = buf_in->tail;
+  }
+
+  buf_out->datalen += buf_in->datalen;
+  buf_in->head = buf_in->tail = NULL;
+  buf_in->datalen = 0;
+}
+
 /** Internal structure: represents a position in a buffer. */
 typedef struct buf_pos_t {
   const chunk_t *chunk; /**< Which chunk are we pointing to? */
diff --git a/src/common/buffers.h b/src/common/buffers.h
index 1eaa5f2d04..dd96ddcb6c 100644
--- a/src/common/buffers.h
+++ b/src/common/buffers.h
@@ -43,9 +43,15 @@ int buf_flush_to_socket(buf_t *buf, tor_socket_t s, size_t sz,
                         size_t *buf_flushlen);
 
 int buf_add(buf_t *buf, const char *string, size_t string_len);
+void buf_add_string(buf_t *buf, const char *string);
+void buf_add_printf(buf_t *buf, const char *format, ...)
+  CHECK_PRINTF(2, 3);
+void buf_add_vprintf(buf_t *buf, const char *format, va_list args)
+  CHECK_PRINTF(2, 0);
 int buf_add_compress(buf_t *buf, struct tor_compress_state_t *state,
                           const char *data, size_t data_len, int done);
 int buf_move_to_buf(buf_t *buf_out, buf_t *buf_in, size_t *buf_flushlen);
+void buf_move_all(buf_t *buf_out, buf_t *buf_in);
 void buf_peek(const buf_t *buf, char *string, size_t string_len);
 void buf_drain(buf_t *buf, size_t n);
 int buf_get_bytes(buf_t *buf, char *string, size_t string_len);
@@ -62,6 +68,7 @@ void buf_assert_ok(buf_t *buf);
 int buf_find_string_offset(const buf_t *buf, const char *s, size_t n);
 void buf_pullup(buf_t *buf, size_t bytes,
                 const char **head_out, size_t *len_out);
+char *buf_extract(buf_t *buf, size_t *sz_out);
 
 #ifdef BUFFERS_PRIVATE
 #ifdef TOR_UNIT_TESTS
diff --git a/src/or/connection.c b/src/or/connection.c
index 431ab3a3a0..e6aae71627 100644
--- a/src/or/connection.c
+++ b/src/or/connection.c
@@ -4047,6 +4047,68 @@ connection_flush(connection_t *conn)
   return connection_handle_write(conn, 1);
 }
 
+/** Helper for connection_write_to_buf_impl and connection_write_buf_to_buf:
+ *
+ * Return true iff it is okay to queue bytes on <b>conn</b>'s outbuf for
+ * writing.
+ */
+static int
+connection_may_write_to_buf(connection_t *conn)
+{
+  /* if it's marked for close, only allow write if we mean to flush it */
+  if (conn->marked_for_close && !conn->hold_open_until_flushed)
+    return 0;
+
+  return 1;
+}
+
+/** Helper for connection_write_to_buf_impl and connection_write_buf_to_buf:
+ *
+ * Called when an attempt to add bytes on <b>conn</b>'s outbuf has failed;
+ * mark the connection and warn as appropriate.
+ */
+static void
+connection_write_to_buf_failed(connection_t *conn)
+{
+  if (CONN_IS_EDGE(conn)) {
+    /* if it failed, it means we have our package/delivery windows set
+       wrong compared to our max outbuf size. close the whole circuit. */
+    log_warn(LD_NET,
+             "write_to_buf failed. Closing circuit (fd %d).", (int)conn->s);
+    circuit_mark_for_close(circuit_get_by_edge_conn(TO_EDGE_CONN(conn)),
+                           END_CIRC_REASON_INTERNAL);
+  } else if (conn->type == CONN_TYPE_OR) {
+    or_connection_t *orconn = TO_OR_CONN(conn);
+    log_warn(LD_NET,
+             "write_to_buf failed on an orconn; notifying of error "
+             "(fd %d)", (int)(conn->s));
+    connection_or_close_for_error(orconn, 0);
+  } else {
+    log_warn(LD_NET,
+             "write_to_buf failed. Closing connection (fd %d).",
+             (int)conn->s);
+    connection_mark_for_close(conn);
+  }
+}
+
+/** Helper for connection_write_to_buf_impl and connection_write_buf_to_buf:
+ *
+ * Called when an attempt to add bytes on <b>conn</b>'s outbuf has succeeded:
+ * record the number of bytes added.
+ */
+static void
+connection_write_to_buf_commit(connection_t *conn, size_t len)
+{
+  /* If we receive optimistic data in the EXIT_CONN_STATE_RESOLVING
+   * state, we don't want to try to write it right away, since
+   * conn->write_event won't be set yet.  Otherwise, write data from
+   * this conn as the socket is available. */
+  if (conn->write_event) {
+    connection_start_writing(conn);
+  }
+  conn->outbuf_flushlen += len;
+}
+
 /** Append <b>len</b> bytes of <b>string</b> onto <b>conn</b>'s
  * outbuf, and ask it to start writing.
  *
@@ -4061,58 +4123,52 @@ connection_write_to_buf_impl_,(const char *string, size_t len,
 {
   /* XXXX This function really needs to return -1 on failure. */
   int r;
-  size_t old_datalen;
   if (!len && !(zlib<0))
     return;
-  /* if it's marked for close, only allow write if we mean to flush it */
-  if (conn->marked_for_close && !conn->hold_open_until_flushed)
+
+  if (!connection_may_write_to_buf(conn))
     return;
 
-  old_datalen = buf_datalen(conn->outbuf);
+  size_t written;
+
   if (zlib) {
+    size_t old_datalen = buf_datalen(conn->outbuf);
     dir_connection_t *dir_conn = TO_DIR_CONN(conn);
     int done = zlib < 0;
     CONN_LOG_PROTECT(conn, r = buf_add_compress(conn->outbuf,
-                                                     dir_conn->compress_state,
-                                                     string, len, done));
+                                                dir_conn->compress_state,
+                                                string, len, done));
+    written = buf_datalen(conn->outbuf) - old_datalen;
   } else {
     CONN_LOG_PROTECT(conn, r = buf_add(conn->outbuf, string, len));
+    written = len;
   }
   if (r < 0) {
-    if (CONN_IS_EDGE(conn)) {
-      /* if it failed, it means we have our package/delivery windows set
-         wrong compared to our max outbuf size. close the whole circuit. */
-      log_warn(LD_NET,
-               "write_to_buf failed. Closing circuit (fd %d).", (int)conn->s);
-      circuit_mark_for_close(circuit_get_by_edge_conn(TO_EDGE_CONN(conn)),
-                             END_CIRC_REASON_INTERNAL);
-    } else if (conn->type == CONN_TYPE_OR) {
-      or_connection_t *orconn = TO_OR_CONN(conn);
-      log_warn(LD_NET,
-               "write_to_buf failed on an orconn; notifying of error "
-               "(fd %d)", (int)(conn->s));
-      connection_or_close_for_error(orconn, 0);
-    } else {
-      log_warn(LD_NET,
-               "write_to_buf failed. Closing connection (fd %d).",
-               (int)conn->s);
-      connection_mark_for_close(conn);
-    }
+    connection_write_to_buf_failed(conn);
     return;
   }
+  connection_write_to_buf_commit(conn, written);
+}
 
-  /* If we receive optimistic data in the EXIT_CONN_STATE_RESOLVING
-   * state, we don't want to try to write it right away, since
-   * conn->write_event won't be set yet.  Otherwise, write data from
-   * this conn as the socket is available. */
-  if (conn->write_event) {
-    connection_start_writing(conn);
-  }
-  if (zlib) {
-    conn->outbuf_flushlen += buf_datalen(conn->outbuf) - old_datalen;
-  } else {
-    conn->outbuf_flushlen += len;
-  }
+/**
+ * Add all bytes from <b>buf</b> to <b>conn</b>'s outbuf, draining them
+ * from <b>buf</b>. (If the connection is marked and will soon be closed,
+ * nothing is drained.)
+ */
+void
+connection_buf_add_buf(connection_t *conn, buf_t *buf)
+{
+  tor_assert(conn);
+  tor_assert(buf);
+  size_t len = buf_datalen(buf);
+  if (len == 0)
+    return;
+
+  if (!connection_may_write_to_buf(conn))
+    return;
+
+  buf_move_all(conn->outbuf, buf);
+  connection_write_to_buf_commit(conn, len);
 }
 
 #define CONN_GET_ALL_TEMPLATE(var, test) \
diff --git a/src/or/connection.h b/src/or/connection.h
index 7f488676ed..450229ce8f 100644
--- a/src/or/connection.h
+++ b/src/or/connection.h
@@ -156,6 +156,7 @@ connection_buf_add_compress(const char *string, size_t len,
 {
   connection_write_to_buf_impl_(string, len, TO_CONN(conn), done ? -1 : 1);
 }
+void connection_buf_add_buf(connection_t *conn, buf_t *buf);
 
 /* DOCDOC connection_get_inbuf_len */
 static size_t connection_get_inbuf_len(connection_t *conn);
diff --git a/src/or/directory.c b/src/or/directory.c
index 777972d576..4dd6da592a 100644
--- a/src/or/directory.c
+++ b/src/or/directory.c
@@ -3515,63 +3515,47 @@ write_http_response_header_impl(dir_connection_t *conn, ssize_t length,
                            long cache_lifetime)
 {
   char date[RFC1123_TIME_LEN+1];
-  char tmp[1024];
-  char *cp;
   time_t now = time(NULL);
+  buf_t *buf = buf_new_with_capacity(1024);
 
   tor_assert(conn);
 
   format_rfc1123_time(date, now);
-  cp = tmp;
-  tor_snprintf(cp, sizeof(tmp),
-               "HTTP/1.0 200 OK\r\nDate: %s\r\n",
-               date);
-  cp += strlen(tmp);
+
+  buf_add_printf(buf, "HTTP/1.0 200 OK\r\nDate: %s\r\n", date);
   if (type) {
-    tor_snprintf(cp, sizeof(tmp)-(cp-tmp), "Content-Type: %s\r\n", type);
-    cp += strlen(cp);
+    buf_add_printf(buf, "Content-Type: %s\r\n", type);
   }
   if (!is_local_addr(&conn->base_.addr)) {
     /* Don't report the source address for a nearby/private connection.
      * Otherwise we tend to mis-report in cases where incoming ports are
      * being forwarded to a Tor server running behind the firewall. */
-    tor_snprintf(cp, sizeof(tmp)-(cp-tmp),
-                 X_ADDRESS_HEADER "%s\r\n", conn->base_.address);
-    cp += strlen(cp);
+    buf_add_printf(buf, X_ADDRESS_HEADER "%s\r\n", conn->base_.address);
   }
   if (encoding) {
-    tor_snprintf(cp, sizeof(tmp)-(cp-tmp),
-                 "Content-Encoding: %s\r\n", encoding);
-    cp += strlen(cp);
+    buf_add_printf(buf, "Content-Encoding: %s\r\n", encoding);
   }
   if (length >= 0) {
-    tor_snprintf(cp, sizeof(tmp)-(cp-tmp),
-                 "Content-Length: %ld\r\n", (long)length);
-    cp += strlen(cp);
+    buf_add_printf(buf, "Content-Length: %ld\r\n", (long)length);
   }
   if (cache_lifetime > 0) {
     char expbuf[RFC1123_TIME_LEN+1];
     format_rfc1123_time(expbuf, (time_t)(now + cache_lifetime));
     /* We could say 'Cache-control: max-age=%d' here if we start doing
      * http/1.1 */
-    tor_snprintf(cp, sizeof(tmp)-(cp-tmp),
-                 "Expires: %s\r\n", expbuf);
-    cp += strlen(cp);
+    buf_add_printf(buf, "Expires: %s\r\n", expbuf);
   } else if (cache_lifetime == 0) {
     /* We could say 'Cache-control: no-cache' here if we start doing
      * http/1.1 */
-    strlcpy(cp, "Pragma: no-cache\r\n", sizeof(tmp)-(cp-tmp));
-    cp += strlen(cp);
+    buf_add_string(buf, "Pragma: no-cache\r\n");
   }
   if (extra_headers) {
-    strlcpy(cp, extra_headers, sizeof(tmp)-(cp-tmp));
-    cp += strlen(cp);
+    buf_add_string(buf, extra_headers);
   }
-  if (sizeof(tmp)-(cp-tmp) > 3)
-    memcpy(cp, "\r\n", 3);
-  else
-    tor_assert(0);
-  connection_buf_add(tmp, strlen(tmp), TO_CONN(conn));
+  buf_add_string(buf, "\r\n");
+
+  connection_buf_add_buf(TO_CONN(conn), buf);
+  buf_free(buf);
 }
 
 /** As write_http_response_header_impl, but sets encoding and content-typed
diff --git a/src/or/geoip.c b/src/or/geoip.c
index 3944b2cf69..c976b8d276 100644
--- a/src/or/geoip.c
+++ b/src/or/geoip.c
@@ -30,6 +30,7 @@
 #define GEOIP_PRIVATE
 #include "or.h"
 #include "ht.h"
+#include "buffers.h"
 #include "config.h"
 #include "control.h"
 #include "dnsserv.h"
@@ -930,9 +931,9 @@ static char *
 geoip_get_dirreq_history(dirreq_type_t type)
 {
   char *result = NULL;
+  buf_t *buf = NULL;
   smartlist_t *dirreq_completed = NULL;
   uint32_t complete = 0, timeouts = 0, running = 0;
-  int bufsize = 1024, written;
   dirreq_map_entry_t **ptr, **next;
   struct timeval now;
 
@@ -965,13 +966,9 @@ geoip_get_dirreq_history(dirreq_type_t type)
                                               DIR_REQ_GRANULARITY);
   running = round_uint32_to_next_multiple_of(running,
                                              DIR_REQ_GRANULARITY);
-  result = tor_malloc_zero(bufsize);
-  written = tor_snprintf(result, bufsize, "complete=%u,timeout=%u,"
-                         "running=%u", complete, timeouts, running);
-  if (written < 0) {
-    tor_free(result);
-    goto done;
-  }
+  buf = buf_new_with_capacity(1024);
+  buf_add_printf(buf, "complete=%u,timeout=%u,"
+                 "running=%u", complete, timeouts, running);
 
 #define MIN_DIR_REQ_RESPONSES 16
   if (complete >= MIN_DIR_REQ_RESPONSES) {
@@ -992,7 +989,7 @@ geoip_get_dirreq_history(dirreq_type_t type)
       dltimes[ent_sl_idx] = bytes_per_second;
     } SMARTLIST_FOREACH_END(ent);
     median_uint32(dltimes, complete); /* sorts as a side effect. */
-    written = tor_snprintf(result + written, bufsize - written,
+    buf_add_printf(buf,
                            ",min=%u,d1=%u,d2=%u,q1=%u,d3=%u,d4=%u,md=%u,"
                            "d6=%u,d7=%u,q3=%u,d8=%u,d9=%u,max=%u",
                            dltimes[0],
@@ -1008,14 +1005,15 @@ geoip_get_dirreq_history(dirreq_type_t type)
                            dltimes[8*complete/10-1],
                            dltimes[9*complete/10-1],
                            dltimes[complete-1]);
-    if (written<0)
-      tor_free(result);
     tor_free(dltimes);
   }
- done:
+
+  result = buf_extract(buf, NULL);
+
   SMARTLIST_FOREACH(dirreq_completed, dirreq_map_entry_t *, ent,
                     tor_free(ent));
   smartlist_free(dirreq_completed);
+  buf_free(buf);
   return result;
 }