From 124af5c0388f70afb7680d5551b45deb3f232183 Mon Sep 17 00:00:00 2001
From: Sean Hefty <sean.hefty@intel.com>
Date: Fri, 10 Aug 2012 22:06:41 -0700
Subject: [PATCH] Refresh of fork-connect

---
 src/preload.c | 69 +++++++++++++++++++++++++++++++++------------------
 1 file changed, 45 insertions(+), 24 deletions(-)

diff --git a/src/preload.c b/src/preload.c
index b18d3100..c61d0921 100644
--- a/src/preload.c
+++ b/src/preload.c
@@ -99,12 +99,20 @@ static int fork_support;
 
 enum fd_type {
 	fd_normal,
-	fd_rsocket,
-	fd_fork
+	fd_rsocket
+};
+
+enum fd_fork_state {
+	fd_ready,
+	fd_fork,
+	fd_fork_listen,
+	fd_fork_active,
+	fd_fork_passive
 };
 
 struct fd_info {
 	enum fd_type type;
+	enum fd_fork_state state;
 	int fd;
 	int dupfd;
 	atomic_t refcnt;
@@ -143,13 +151,14 @@ err1:
 	return ret;
 }
 
-static void fd_store(int index, int fd, enum fd_type type)
+static void fd_store(int index, int fd, enum fd_type type, enum fd_fork_state state)
 {
 	struct fd_info *fdi;
 
 	fdi = idm_at(&idm, index);
 	fdi->fd = fd;
 	fdi->type = type;
+	fdi->state = state;
 }
 
 static inline enum fd_type fd_get(int index, int *fd)
@@ -175,6 +184,14 @@ static inline int fd_getd(int index)
 	return fdi ? fdi->fd : index;
 }
 
+static inline enum fd_state fd_gets(int index)
+{
+	struct fd_info *fdi;
+
+	fdi = idm_lookup(&idm, index);
+	return fdi ? fdi->state : fd_ready;
+}
+
 static inline enum fd_type fd_gett(int index)
 {
 	struct fd_info *fdi;
@@ -353,7 +370,7 @@ static int transpose_socket(int socket, enum fd_type new_type)
 	if (ret)
 		goto err;
 
-	fd_store(socket, dfd, new_type);
+	fd_store(socket, dfd, new_type, fd_ready);
 	return dfd;
 
 err:
@@ -398,9 +415,9 @@ int socket(int domain, int type, int protocol)
 			ret = real.socket(domain, type, protocol);
 			if (ret < 0)
 				return ret;
-			fd_store(index, ret, fd_fork);
+			fd_store(index, ret, fd_normal, fd_fork);
 		} else {
-			fd_store(index, ret, fd_rsocket);
+			fd_store(index, ret, fd_rsocket, fd_ready);
 			set_rsocket_options(ret);
 		}
 		return index;
@@ -427,22 +444,32 @@ int listen(int socket, int backlog)
 int accept(int socket, struct sockaddr *addr, socklen_t *addrlen)
 {
 	int fd, index, ret;
-	enum fd_type type;
 
-	type = fd_get(socket, &fd);
-	if (type == fd_rsocket || type == fd_fork) {
+	if (fd_get(socket, &fd) == fd_rsocket) {
+		index = fd_open();
+		if (index < 0)
+			return index;
+
+		ret = raccept(fd, addr, addrlen);
+		if (ret < 0) {
+			fd_close(index, &fd);
+			return ret;
+		}
+
+		fd_store(index, ret, fd_rsocket, fd_ready);
+		return index;
+	} else if (fd_gets(socket) == fd_fork_listen) {
 		index = fd_open();
 		if (index < 0)
 			return index;
 
-		ret = (type == fd_rsocket) ? raccept(fd, addr, addrlen) :
-					     real.accept(fd, addr, addrlen);
+		ret = real.accept(fd, addr, addrlen);
 		if (ret < 0) {
 			fd_close(index, &fd);
 			return ret;
 		}
 
-		fd_store(index, ret, type);
+		fd_store(index, ret, fd_normal, fd_fork_passive);
 		return index;
 	} else {
 		return real.accept(fd, addr, addrlen);
@@ -471,7 +498,7 @@ static int fork_active(int socket, const struct sockaddr *addr, socklen_t addrle
 
 	ret = real.recv(fd, &msg, sizeof msg, MSG_PEEK);
 	if ((ret != sizeof msg) || msg) {
-		fd_store(socket, fd, fd_normal);
+		fd_store(socket, fd, fd_normal, fd_ready);
 		return 0;
 	}
 
@@ -544,7 +571,7 @@ static void fork_passive(int socket)
 	copysockopts(dfd, sfd, &rs, &real);
 	real.shutdown(sfd, SHUT_RDWR);
 	real.close(sfd);
-	fd_store(socket, dfd, fd_rsocket);
+	fd_store(socket, dfd, fd_rsocket, fd_ready);
 
 lclose:
 	rclose(lfd);
@@ -553,7 +580,7 @@ sclose:
 	sem_close(sem);
 out:
 	if (ret)
-		fd_store(socket, sfd, fd_normal);
+		fd_store(socket, sfd, fd_normal, fd_ready);
 }
 
 static inline enum fd_type fd_fork_get(int index, int *fd)
@@ -577,10 +604,7 @@ int connect(int socket, const struct sockaddr *addr, socklen_t addrlen)
 {
 	int fd, ret;
 
-	switch (fd_get(socket, &fd)) {
-	case fd_fork:
-		return fork_active(socket, addr, addrlen);
-	case fd_rsocket:
+	if (fd_get(socket, &fd) == fd_rsocket) {
 		ret = rconnect(fd, addr, addrlen);
 		if (!ret || errno == EINPROGRESS)
 			return ret;
@@ -591,12 +615,9 @@ int connect(int socket, const struct sockaddr *addr, socklen_t addrlen)
 
 		rclose(fd);
 		fd = ret;
-		break;
-	default:
-		break;
+	} else {
+		return real.connect(fd, addr, addrlen);
 	}
-
-	return real.connect(fd, addr, addrlen);
 }
 
 ssize_t recv(int socket, void *buf, size_t len, int flags)
-- 
2.45.2