1 files changed, 32 insertions, 31 deletions

diff --git a/netio.c b/netio.c
index c955430..74eea7b 100644
--- a/netio.c
+++ b/netio.c
@@ -70,13 +70,13 @@ static void connect_try_next(struct dropbear_progress_connection *c) {
 	struct addrinfo *r;
 	int res = 0;
 	int fastopen = 0;
-#ifdef DROPBEAR_TCP_FAST_OPEN
+#ifdef DROPBEAR_CLIENT_TCP_FAST_OPEN
 	struct msghdr message;
 #endif
 
 	for (r = c->res_iter; r; r = r->ai_next)
 	{
-		assert(c->sock == -1);
+		dropbear_assert(c->sock == -1);
 
 		c->sock = socket(c->res_iter->ai_family, c->res_iter->ai_socktype, c->res_iter->ai_protocol);
 		if (c->sock < 0) {
@@ -91,31 +91,36 @@ static void connect_try_next(struct dropbear_progress_connection *c) {
 		set_piggyback_ack(c->sock);
 #endif
 
-#ifdef DROPBEAR_TCP_FAST_OPEN
+#ifdef DROPBEAR_CLIENT_TCP_FAST_OPEN
 		fastopen = (c->writequeue != NULL);
 
-		memset(&message, 0x0, sizeof(message));
-		message.msg_name = r->ai_addr;
-		message.msg_namelen = r->ai_addrlen;
-
-		if (c->writequeue) {
-			int iovlen; /* Linux msg_iovlen is a size_t */
-			message.msg_iov = packet_queue_to_iovec(c->writequeue, &iovlen);
+		if (fastopen) {
+			memset(&message, 0x0, sizeof(message));
+			message.msg_name = r->ai_addr;
+			message.msg_namelen = r->ai_addrlen;
+			/* 6 is arbitrary, enough to hold initial packets */
+			unsigned int iovlen = 6; /* Linux msg_iovlen is a size_t */
+			struct iovec iov[6];
+			packet_queue_to_iovec(c->writequeue, iov, &iovlen);
+			message.msg_iov = iov;
 			message.msg_iovlen = iovlen;
 			res = sendmsg(c->sock, &message, MSG_FASTOPEN);
-			if (res < 0 && errno != EINPROGRESS) {
-				m_free(c->errstring);
-				c->errstring = m_strdup(strerror(errno));
-				/* Not entirely sure which kind of errors are normal - 2.6.32 seems to 
-				return EPIPE for any (nonblocking?) sendmsg(). just fall back */
-				TRACE(("sendmsg tcp_fastopen failed, falling back. %s", strerror(errno)));
-				/* No kernel MSG_FASTOPEN support. Fall back below */
-				fastopen = 0;
-				/* Set to NULL to avoid trying again */
-				c->writequeue = NULL;
+			/* Returns EINPROGRESS if FASTOPEN wasn't available */
+			if (res < 0) {
+				if (errno != EINPROGRESS) {
+					m_free(c->errstring);
+					c->errstring = m_strdup(strerror(errno));
+					/* Not entirely sure which kind of errors are normal - 2.6.32 seems to 
+					return EPIPE for any (nonblocking?) sendmsg(). just fall back */
+					TRACE(("sendmsg tcp_fastopen failed, falling back. %s", strerror(errno)));
+					/* No kernel MSG_FASTOPEN support. Fall back below */
+					fastopen = 0;
+					/* Set to NULL to avoid trying again */
+					c->writequeue = NULL;
+				}
+			} else {
+				packet_queue_consume(c->writequeue, res);
 			}
-			m_free(message.msg_iov);
-			packet_queue_consume(c->writequeue, res);
 		}
 #endif
 
@@ -163,7 +168,7 @@ struct dropbear_progress_connection *connect_remote(const char* remotehost, cons
 
 	memset(&hints, 0, sizeof(hints));
 	hints.ai_socktype = SOCK_STREAM;
-	hints.ai_family = PF_UNSPEC;
+	hints.ai_family = AF_UNSPEC;
 
 	err = getaddrinfo(remotehost, remoteport, &hints, &c->res);
 	if (err) {
@@ -258,8 +263,7 @@ void connect_set_writequeue(struct dropbear_progress_connection *c, struct Queue
 	c->writequeue = writequeue;
 }
 
-struct iovec * packet_queue_to_iovec(struct Queue *queue, int *ret_iov_count) {
-	struct iovec *iov = NULL;
+void packet_queue_to_iovec(struct Queue *queue, struct iovec *iov, unsigned int *iov_count) {
 	struct Link *l;
 	unsigned int i;
 	int len;
@@ -269,10 +273,9 @@ struct iovec * packet_queue_to_iovec(struct Queue *queue, int *ret_iov_count) {
 	#define IOV_MAX UIO_MAXIOV
 	#endif
 
-	*ret_iov_count = MIN(queue->count, IOV_MAX);
+	*iov_count = MIN(MIN(queue->count, IOV_MAX), *iov_count);
 
-	iov = m_malloc(sizeof(*iov) * *ret_iov_count);
-	for (l = queue->head, i = 0; l; l = l->link, i++)
+	for (l = queue->head, i = 0; i < *iov_count; l = l->link, i++)
 	{
 		writebuf = (buffer*)l->item;
 		len = writebuf->len - 1 - writebuf->pos;
@@ -282,8 +285,6 @@ struct iovec * packet_queue_to_iovec(struct Queue *queue, int *ret_iov_count) {
 		iov[i].iov_base = buf_getptr(writebuf, len);
 		iov[i].iov_len = len;
 	}
-
-	return iov;
 }
 
 void packet_queue_consume(struct Queue *queue, ssize_t written) {
@@ -312,7 +313,7 @@ void set_sock_nodelay(int sock) {
 	setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, (void*)&val, sizeof(val));
 }
 
-#ifdef DROPBEAR_TCP_FAST_OPEN
+#ifdef DROPBEAR_SERVER_TCP_FAST_OPEN
 void set_listen_fast_open(int sock) {
 	int qlen = MAX(MAX_UNAUTH_PER_IP, 5);
 	if (setsockopt(sock, SOL_TCP, TCP_FASTOPEN, &qlen, sizeof(qlen)) != 0) {