Make main socket nonblocking. Limit writequeue size.

1 files changed, 16 insertions, 6 deletions

diff --git a/common-session.c b/common-session.c
index 164dc85..44e7030 100644
--- a/common-session.c
+++ b/common-session.c
@@ -64,6 +64,13 @@ void common_session_init(int sock_in, int sock_out) {
 	ses.sock_out = sock_out;
 	ses.maxfd = MAX(sock_in, sock_out);
 
+	if (sock_in >= 0) {
+		setnonblocking(sock_in);
+	}
+	if (sock_out >= 0) {
+		setnonblocking(sock_out);
+	}
+
 	ses.socket_prio = DROPBEAR_PRIO_DEFAULT;
 	/* Sets it to lowdelay */
 	update_channel_prio();
@@ -145,6 +152,7 @@ void session_loop(void(*loophandler)()) {
 
 	/* main loop, select()s for all sockets in use */
 	for(;;) {
+		const int writequeue_has_space = (ses.writequeue_len <= 2*TRANS_MAX_PAYLOAD_LEN);
 
 		timeout.tv_sec = select_timeout();
 		timeout.tv_usec = 0;
@@ -155,8 +163,12 @@ void session_loop(void(*loophandler)()) {
 		/* We delay reading from the input socket during initial setup until
 		after we have written out our initial KEXINIT packet (empty writequeue). 
 		This means our initial packet can be in-flight while we're doing a blocking
-		read for the remote ident */
-		if (ses.sock_in != -1 && (ses.remoteident || isempty(&ses.writequeue))) {
+		read for the remote ident.
+		We also avoid reading from the socket if the writequeue is full, that avoids
+		replies backing up */
+		if (ses.sock_in != -1 
+			&& (ses.remoteident || isempty(&ses.writequeue)) 
+			&& writequeue_has_space) {
 			FD_SET(ses.sock_in, &readfd);
 		}
 		if (ses.sock_out != -1 && !isempty(&ses.writequeue)) {
@@ -168,7 +180,7 @@ void session_loop(void(*loophandler)()) {
 		FD_SET(ses.signal_pipe[0], &readfd);
 
 		/* set up for channels which can be read/written */
-		setchannelfds(&readfd, &writefd);
+		setchannelfds(&readfd, &writefd, writequeue_has_space);
 
 		/* Pending connections to test */
 		set_connect_fds(&writefd);
@@ -318,9 +330,7 @@ void session_cleanup() {
 void send_session_identification() {
 	buffer *writebuf = buf_new(strlen(LOCAL_IDENT "\r\n") + 1);
 	buf_putbytes(writebuf, LOCAL_IDENT "\r\n", strlen(LOCAL_IDENT "\r\n"));
-	buf_putbyte(writebuf, 0x0); /* packet type */
-	buf_setpos(writebuf, 0);
-	enqueue(&ses.writequeue, writebuf);
+	writebuf_enqueue(writebuf, 0);
 }
 
 static void read_session_identification() {