From ec2d7ac18cb751acf53a0e7da88651cbf323ad4f Mon Sep 17 00:00:00 2001
From: tedu <tedu@openbsd.org>
Date: Fri, 16 Oct 2015 18:29:05 +0000
Subject: [PATCH] two phase handling for tcp so that slow connects don't stall
 the process

---
 usr.sbin/rebound/rebound.c | 74 +++++++++++++++++++++++++++++---------
 1 file changed, 58 insertions(+), 16 deletions(-)

diff --git a/usr.sbin/rebound/rebound.c b/usr.sbin/rebound/rebound.c
index ddba7053ac3..8c8a67ac2db 100644
--- a/usr.sbin/rebound/rebound.c
+++ b/usr.sbin/rebound/rebound.c
@@ -1,4 +1,4 @@
-/* $OpenBSD: rebound.c,v 1.20 2015/10/16 15:35:05 tedu Exp $ */
+/* $OpenBSD: rebound.c,v 1.21 2015/10/16 18:29:05 tedu Exp $ */
 /*
  * Copyright (c) 2015 Ted Unangst <tedu@openbsd.org>
  *
@@ -24,6 +24,7 @@
 #include <sys/signal.h>
 #include <sys/wait.h>
 
+#include <fcntl.h>
 #include <signal.h>
 #include <syslog.h>
 #include <stdlib.h>
@@ -78,6 +79,7 @@ static TAILQ_HEAD(, dnscache) cache;
 struct request {
 	int s;
 	int client;
+	int phase;
 	struct sockaddr from;
 	socklen_t fromlen;
 	struct timespec ts;
@@ -164,6 +166,9 @@ newrequest(int ud, struct sockaddr *remoteaddr)
 	if (!(req = calloc(1, sizeof(*req))))
 		return NULL;
 
+	req->ts = now;
+	req->ts.tv_sec += 30;
+
 	req->client = -1;
 	memcpy(&req->from, &from, fromlen);
 	req->fromlen = fromlen;
@@ -196,8 +201,6 @@ newrequest(int ud, struct sockaddr *remoteaddr)
 	}
 	if (send(req->s, buf, r, 0) != r)
 		goto fail;
-	req->ts = now;
-	req->ts.tv_sec += 30;
 
 	return req;
 fail:
@@ -263,6 +266,24 @@ freecacheent(struct dnscache *ent)
 	free(ent);
 }
 
+static struct request *
+tcpphasetwo(struct request *req)
+{
+	req->phase = 2;
+	if (setsockopt(req->client, SOL_SOCKET, SO_SPLICE, &req->s,
+	    sizeof(req->s)) == -1)
+		goto fail;
+	if (setsockopt(req->s, SOL_SOCKET, SO_SPLICE, &req->client,
+	    sizeof(req->client)) == -1)
+		goto fail;
+
+	return req;
+
+fail:
+	freerequest(req);
+	return NULL;
+}
+
 static struct request *
 newtcprequest(int ld, struct sockaddr *remoteaddr)
 {
@@ -271,6 +292,9 @@ newtcprequest(int ld, struct sockaddr *remoteaddr)
 	if (!(req = calloc(1, sizeof(*req))))
 		return NULL;
 
+	req->ts = now;
+	req->ts.tv_sec += 30;
+
 	req->s = -1;
 	req->fromlen = sizeof(req->from);
 	req->client = accept(ld, &req->from, &req->fromlen);
@@ -280,17 +304,18 @@ newtcprequest(int ld, struct sockaddr *remoteaddr)
 	req->s = socket(remoteaddr->sa_family, SOCK_STREAM, 0);
 	if (req->s == -1)
 		goto fail;
-	if (connect(req->s, remoteaddr, remoteaddr->sa_len) == -1)
-		goto fail;
-	if (setsockopt(req->client, SOL_SOCKET, SO_SPLICE, &req->s,
-	    sizeof(req->s)) == -1)
-		goto fail;
-	if (setsockopt(req->s, SOL_SOCKET, SO_SPLICE, &req->client,
-	    sizeof(req->client)) == -1)
+	if (fcntl(req->s, F_SETFL, O_NONBLOCK) == -1)
 		goto fail;
-	req->ts = now;
-	req->ts.tv_sec += 30;
+	req->phase = 1;
+	if (connect(req->s, remoteaddr, remoteaddr->sa_len) == -1) {
+		if (errno != EINPROGRESS)
+			goto fail;
+	} else {
+		TAILQ_INSERT_TAIL(&reqfifo, req, fifo);
+		return tcpphasetwo(req);
+	}
 
+	TAILQ_INSERT_TAIL(&reqfifo, req, fifo);
 	return req;
 
 fail:
@@ -333,7 +358,7 @@ static int
 launch(const char *confname, int ud, int ld, int kq)
 {
 	struct sockaddr_storage remoteaddr;
-	struct kevent chlist[1], kev[4];
+	struct kevent chlist[2], kev[4];
 	struct timespec ts, *timeout = NULL;
 	struct request *req;
 	struct dnscache *ent;
@@ -407,10 +432,27 @@ launch(const char *confname, int ud, int ld, int kq)
 				req = newtcprequest(ld,
 				    (struct sockaddr *)&remoteaddr);
 				if (req) {
-					EV_SET(&chlist[0], req->s, EVFILT_READ,
-					    EV_ADD, 0, 0, NULL);
+					EV_SET(&chlist[0], req->s,
+					    req->phase == 1 ? EVFILT_WRITE :
+					    EVFILT_READ, EV_ADD, 0, 0, NULL);
 					kevent(kq, chlist, 1, NULL, 0, NULL);
-					TAILQ_INSERT_TAIL(&reqfifo, req, fifo);
+				}
+			} else if (kev[i].filter == EVFILT_WRITE) {
+				req = TAILQ_FIRST(&reqfifo);
+				while (req) {
+					if (req->s == kev[i].ident)
+						break;
+					req = TAILQ_NEXT(req, fifo);
+				}
+				if (!req)
+					logerr("lost request");
+				req = tcpphasetwo(req);
+				if (req) {
+					EV_SET(&chlist[0], req->s, EVFILT_WRITE,
+					    EV_DELETE, 0, 0, NULL);
+					EV_SET(&chlist[1], req->s, EVFILT_READ,
+					    EV_ADD, 0, 0, NULL);
+					kevent(kq, chlist, 2, NULL, 0, NULL);
 				}
 			} else if (kev[i].filter == EVFILT_READ) {
 				/* use a tree here? */
-- 
2.20.1