generalise write iovec handling

1 files changed, 42 insertions, 49 deletions

diff --git a/packet.c b/packet.c
index d02ec69..7b2b865 100644
--- a/packet.c
+++ b/packet.c
@@ -52,25 +52,52 @@ static buffer* buf_decompress(buffer* buf, unsigned int len);
 static void buf_compress(buffer * dest, buffer * src, unsigned int len);
 #endif
 
-#if 0
-struct iovec * dropbear_queue_to_iovec(struct Queue *queue) {
-
+struct iovec * packet_queue_to_iovec(struct Queue *queue, int *ret_iov_count) {
 	struct iovec *iov = NULL;
 	struct Link *l;
-	int iov_max_count;
+	unsigned int i, packet_type;
+	int len;
+	buffer *writebuf;
 
 	#ifndef IOV_MAX
 	#define IOV_MAX UIO_MAXIOV
 	#endif
 
-#error incomplete
+	*ret_iov_count = MIN(queue->count, IOV_MAX);
 
-}
+	iov = m_malloc(sizeof(*iov) * *ret_iov_count);
+	for (l = queue->head, i = 0; l; l = l->link, i++)
+	{
+		writebuf = (buffer*)l->item;
+		packet_type = writebuf->data[writebuf->len-1];
+		len = writebuf->len - 1 - writebuf->pos;
+		dropbear_assert(len > 0);
+		TRACE2(("write_packet writev #%d  type %d len %d/%d", i, packet_type,
+				len, writebuf->len-1))
+		iov[i].iov_base = buf_getptr(writebuf, len);
+		iov[i].iov_len = len;
+	}
 
-void dropbear_queue_consume(struct Queue *queue, ssize_t written) {
+	return iov;
+}
 
+void packet_queue_consume(struct Queue *queue, ssize_t written) {
+	buffer *writebuf;
+	int len;
+	while (written > 0) {
+		writebuf = (buffer*)examine(queue);
+		len = writebuf->len - 1 - writebuf->pos;
+		if (len > written) {
+			/* partial buffer write */
+			buf_incrpos(writebuf, written);
+			written = 0;
+		} else {
+			written -= len;
+			dequeue(queue);
+			buf_free(writebuf);
+		}
+	}
 }
-#endif
 
 /* non-blocking function writing out a current encrypted packet */
 void write_packet() {
@@ -83,7 +110,7 @@ void write_packet() {
 	struct iovec *iov = NULL;
 	int i;
 	struct Link *l;
-	int iov_max_count;
+	int iov_count;
 #endif
 	
 	TRACE2(("enter write_packet"))
@@ -91,62 +118,28 @@ void write_packet() {
 
 #if defined(HAVE_WRITEV) && (defined(IOV_MAX) || defined(UIO_MAXIOV))
 
-#ifndef IOV_MAX
-#define IOV_MAX UIO_MAXIOV
-#endif
-
-	/* Make sure the size of the iov is below the maximum allowed by the OS. */
-	iov_max_count = ses.writequeue.count;
-	if (iov_max_count > IOV_MAX)
-	{
-		iov_max_count = IOV_MAX;
-	}
-
-	iov = m_malloc(sizeof(*iov) * iov_max_count);
-	for (l = ses.writequeue.head, i = 0; l; l = l->link, i++)
-	{
-		writebuf = (buffer*)l->item;
-		packet_type = writebuf->data[writebuf->len-1];
-		len = writebuf->len - 1 - writebuf->pos;
-		dropbear_assert(len > 0);
-		TRACE2(("write_packet writev #%d  type %d len %d/%d", i, packet_type,
-				len, writebuf->len-1))
-		iov[i].iov_base = buf_getptr(writebuf, len);
-		iov[i].iov_len = len;
-	}
+	iov = packet_queue_to_iovec(&ses.writequeue, &iov_count);
 	/* This may return EAGAIN. The main loop sometimes
 	calls write_packet() without bothering to test with select() since
 	it's likely to be necessary */
-	written = writev(ses.sock_out, iov, iov_max_count);
+	written = writev(ses.sock_out, iov, iov_count);
 	if (written < 0) {
 		if (errno == EINTR || errno == EAGAIN) {
-			m_free(iov);
 			TRACE2(("leave write_packet: EINTR"))
+			m_free(iov);
 			return;
 		} else {
 			dropbear_exit("Error writing: %s", strerror(errno));
 		}
 	}
+	m_free(iov);
+
+	packet_queue_consume(&ses.writequeue, written);
 
 	if (written == 0) {
 		ses.remoteclosed();
 	}
 
-	while (written > 0) {
-		writebuf = (buffer*)examine(&ses.writequeue);
-		len = writebuf->len - 1 - writebuf->pos;
-		if (len > written) {
-			/* partial buffer write */
-			buf_incrpos(writebuf, written);
-			written = 0;
-		} else {
-			written -= len;
-			dequeue(&ses.writequeue);
-			buf_free(writebuf);
-		}
-	}
-
-	m_free(iov);
 #else /* No writev () */
 	/* Get the next buffer in the queue of encrypted packets to write*/
 	writebuf = (buffer*)examine(&ses.writequeue);