From 02c1d2547c4dfdbfd1de46f85d344a9d11f41db0 Mon Sep 17 00:00:00 2001
From: Yoshinori Sato <yo-satoh@sios.com>
Date: Wed, 10 Jun 2015 14:51:15 +0900
Subject: [PATCH] Add time out for receiving data.

Signed-off-by: Yoshinori Sato <yo-satoh@sios.com>
---
 man/nbd-server.1.in.sgml |  8 +++++++
 nbd-server.c             | 58 ++++++++++++++++++++++++++++++++++++++----------
 nbdsrv.h                 |  1 +
 3 files changed, 55 insertions(+), 12 deletions(-)

diff --git a/man/nbd-server.1.in.sgml b/man/nbd-server.1.in.sgml
index 72a840b..ba7a55e 100644
--- a/man/nbd-server.1.in.sgml
+++ b/man/nbd-server.1.in.sgml
@@ -69,6 +69,7 @@ manpage.1: manpage.sgml
       <arg><option>-M <replaceable>max connections</replaceable></option></arg>
       <arg><option>-V</option></arg>
       <arg><option>-d</option></arg>
+      <arg><option>-T <replaceable>timeout seconds</replaceable></option></arg>
     </cmdsynopsis>
   </refsynopsisdiv>
   <refsect1>
@@ -237,6 +238,13 @@ manpage.1: manpage.sgml
 	</listitem>
       </varlistentry>
       <varlistentry>
+        <term><option>-T</option></term>
+	<listitem>
+	  <para>Timeout of data receiving. If set zero,
+	  waiting infinity.</para>
+	</listitem>
+      </varlistentry>
+      <varlistentry>
 	<term><option>host list filename</option></term>
 	<listitem>
 	  <para>This argument should contain a list of IP-addresses
diff --git a/nbd-server.c b/nbd-server.c
index d51075d..c940ad2 100644
--- a/nbd-server.c
+++ b/nbd-server.c
@@ -271,11 +271,30 @@ static inline const char * getcommandname(uint64_t command) {
  * @param f a file descriptor
  * @param buf a buffer
  * @param len the number of bytes to be read
+ * @param tout timeout value of receive
  **/
-static inline void readit(int f, void *buf, size_t len) {
+static inline int readit(int f, void *buf, size_t len, int tout) {
 	ssize_t res;
 	while (len > 0) {
 		DEBUG("*");
+		if (tout > 0) {
+			fd_set rfd;
+			struct timeval tv;
+
+			FD_ZERO(&rfd);
+			FD_SET(f, &rfd);
+			tv.tv_sec = tout;
+			tv.tv_usec = 0;
+			errno = 0;
+			if (select(f + 1, &rfd, NULL, NULL, &tv) < 1) {
+				if (errno == EINTR)
+					continue;	/* Ignore signal */
+				/* Timeout occured */
+				err_nonfatal("communication timed out");
+				return -1;
+			}
+		}
+			
 		if ((res = read(f, buf, len)) <= 0) {
 			if(errno != EAGAIN) {
 				err("Read failed: %m");
@@ -285,6 +304,7 @@ static inline void readit(int f, void *buf, size_t len) {
 			buf += res;
 		}
 	}
+	return 0;
 }
 
 /**
@@ -295,11 +315,11 @@ static inline void readit(int f, void *buf, size_t len) {
  * @param len the number of bytes to consume
  * @param bufsiz the size of the buffer
  **/
-static inline void consume(int f, void * buf, size_t len, size_t bufsiz) {
+static inline void consume(int f, void * buf, size_t len, size_t bufsiz, int tout) {
 	size_t curlen;
 	while (len>0) {
 		curlen = (len>bufsiz)?bufsiz:len;
-		readit(f, buf, curlen);
+		readit(f, buf, curlen, tout);
 		len -= curlen;
 	}
 }
@@ -485,6 +505,7 @@ SERVER* cmdline(int argc, char *argv[]) {
 		{"output-config", required_argument, NULL, 'o'},
 		{"max-connection", required_argument, NULL, 'M'},
 		{"version", no_argument, NULL, 'V'},
+		{"timeout", required_argument, NULL, 'T'},
 		{0,0,0,0}
 	};
 	SERVER *serve;
@@ -589,6 +610,9 @@ SERVER* cmdline(int argc, char *argv[]) {
 			printf("This is nbd-server version " VERSION "\n");
 			exit(EXIT_SUCCESS);
 			break;
+		case 'T':
+			serve->timeout_data = strtol(optarg, NULL, 0);
+			break;
 		default:
 			usage();
 			exit(EXIT_FAILURE);
@@ -739,6 +763,7 @@ GArray* parse_cfile(gchar* f, struct generic_conf *const genconf, bool expect_ge
 		{ "trim",	FALSE,  PARAM_BOOL,	&(s.flags),		F_TRIM },
 		{ "listenaddr", FALSE,  PARAM_STRING,   &(s.listenaddr),	0 },
 		{ "maxconnections", FALSE, PARAM_INT,	&(s.max_connections),	0 },
+		{ "timeout",    FALSE,  PARAM_INT,      &(s.timeout_data),      0 },
 	};
 	const int lp_size=sizeof(lp)/sizeof(PARAM);
         struct generic_conf genconftmp;
@@ -1612,7 +1637,7 @@ int mainloop(CLIENT *client) {
 		i++;
 		printf("%d: ", i);
 #endif
-		readit(client->net, &request, sizeof(request));
+		readit(client->net, &request, sizeof(request), 0);
 		if (client->transactionlogfd != -1)
 			writeit(client->transactionlogfd, &request, sizeof(request));
 
@@ -1670,20 +1695,29 @@ int mainloop(CLIENT *client) {
 		case NBD_CMD_WRITE:
 			DEBUG("wr: net->buf, ");
 			while(len > 0) {
-				readit(client->net, buf, currlen);
-				DEBUG("buf->exp, ");
-				if ((client->server->flags & F_READONLY) ||
+				if (readit(client->net, buf, currlen, 
+					   client->server->timeout_data) == 0) {
+					DEBUG("buf->exp, ");
+					if ((client->server->flags & F_READONLY) ||
 				    (client->server->flags & F_AUTOREADONLY)) {
-					DEBUG("[WRITE to READONLY!]");
-					ERROR(client, reply, EPERM);
-					consume(client->net, buf, len-currlen, BUFSIZE);
-					continue;
+						DEBUG("[WRITE to READONLY!]");
+						ERROR(client, reply, EPERM);
+						consume(client->net, buf, 
+							len-currlen, BUFSIZE,
+							client->server->timeout_data);
+						continue;
+					}
+				} else {
+					/* Timeout occuerd */
+					ERROR(client, reply, EIO);
 				}
 				if (expwrite(request.from, buf, currlen, client,
 					     request.type & NBD_CMD_FLAG_FUA)) {
 					DEBUG("Write failed: %m" );
 					ERROR(client, reply, errno);
-					consume(client->net, buf, len-currlen, BUFSIZE);
+					consume(client->net, buf, 
+						len-currlen, BUFSIZE,
+						client->server->timeout_data);
 					continue;
 				}
 				len -= currlen;
diff --git a/nbdsrv.h b/nbdsrv.h
index f80b9e9..404f980 100644
--- a/nbdsrv.h
+++ b/nbdsrv.h
@@ -47,6 +47,7 @@ typedef struct {
 	int max_connections; /**< maximum number of opened connections */
 	gchar* transactionlog;/**< filename for transaction log */
 	gchar* cowdir;	     /**< directory for copy-on-write diff files. */
+	int timeout_data;    /**< timeout for receive dara. */
 } SERVER;
 
 /**
-- 
2.1.4