--- a/freeDiameter/cnxctx.c	Fri Feb 05 18:43:33 2010 +0900
+++ b/freeDiameter/cnxctx.c	Mon Feb 08 17:46:40 2010 +0900
@@ -230,6 +230,9 @@
 	cli->cc_socket = cli_sock;
 	cli->cc_proto = serv->cc_proto;
 	
+	/* Set the timeout */
+	fd_cnx_s_setto(cli->cc_socket);
+	
 	/* Generate the name for the connection object */
 	{
 		char addrbuf[INET6_ADDRSTRLEN];
@@ -296,6 +299,9 @@
 	cnx->cc_socket = sock;
 	cnx->cc_proto  = IPPROTO_TCP;
 	
+	/* Set the timeout */
+	fd_cnx_s_setto(cnx->cc_socket);
+	
 	/* Generate the names for the object */
 	{
 		char addrbuf[INET6_ADDRSTRLEN];
@@ -347,6 +353,9 @@
 	cnx->cc_socket = sock;
 	cnx->cc_proto  = IPPROTO_SCTP;
 	
+	/* Set the timeout */
+	fd_cnx_s_setto(cnx->cc_socket);
+	
 	/* Retrieve the number of streams and primary address */
 	CHECK_FCT_DO( fd_sctp_get_str_info( sock, &cnx->cc_sctp_para.str_in, &cnx->cc_sctp_para.str_out, &primary ), goto error );
 	if (cnx->cc_sctp_para.str_out > cnx->cc_sctp_para.str_in)
@@ -488,6 +497,45 @@
 /*     Use of a connection object     */
 /**************************************/
 
+/* Set the timeout option on the socket */
+void fd_cnx_s_setto(int sock) 
+{
+	struct timeval tv;
+	
+	/* Set a timeout on the socket so that in any case we are not stuck waiting for something */
+	memset(&tv, 0, sizeof(tv));
+	tv.tv_sec = 3;	/* allow 3 seconds timeout for TLS session cleanup */
+	CHECK_SYS_DO( setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)), /* best effort only */ );
+}	
+
+/* A recv-like function, taking a cnxctx object instead of socket as entry. Only used to filter timeouts error (GNUTLS does not like these...) */
+ssize_t fd_cnx_s_recv(struct cnxctx * conn, void *buffer, size_t length)
+{
+	ssize_t ret = 0;
+	int timedout = 0;
+again:
+	ret = recv(conn->cc_socket, buffer, length, 0);
+	/* Handle special case of timeout */
+	if ((ret < 0) && (errno == ETIMEDOUT)) {
+		if (!conn->cc_closing)
+			goto again; /* don't care, just ignore */
+		if (!timedout) {
+			timedout ++; /* allow for one timeout while closing */
+			goto again;
+		}
+		CHECK_SYS_DO(ret, /* continue */);
+		return 0; /* so that the connection appears closed */
+	}
+	
+	return ret;
+}
+
+/* Send */
+static ssize_t fd_cnx_s_send(struct cnxctx * conn, void *buffer, size_t length)
+{
+	return send(conn->cc_socket, buffer, length, 0);
+}
+
 /* Receiver thread (TCP & noTLS) : incoming message is directly saved into the target queue */
 static void * rcvthr_notls_tcp(void * arg)
 {
@@ -516,7 +564,7 @@
 		size_t	received = 0;
 
 		do {
-			ret = recv(conn->cc_socket, &header[received], sizeof(header) - received, 0);
+			ret = fd_cnx_s_recv(conn, &header[received], sizeof(header) - received);
 			if (ret <= 0) {
 				CHECK_SYS_DO(ret, /* continue */);
 				goto error; /* Stop the thread, the recipient of the event will cleanup */
@@ -541,7 +589,7 @@
 
 		while (received < length) {
 			pthread_cleanup_push(free, newmsg); /* In case we are canceled, clean the partialy built buffer */
-			ret = recv(conn->cc_socket, newmsg + received, length - received, 0);
+			ret = fd_cnx_s_recv(conn, newmsg + received, length - received);
 			pthread_cleanup_pop(0);
 
 			if (ret <= 0) {
@@ -591,7 +639,7 @@
 	ASSERT( Target_Queue(conn) );
 	
 	do {
-		CHECK_FCT_DO( fd_sctp_recvmeta(conn->cc_socket, NULL, &buf, &bufsz, &event), goto error );
+		CHECK_FCT_DO( fd_sctp_recvmeta(conn->cc_socket, NULL, &buf, &bufsz, &event, &conn->cc_closing), goto error );
 		if (event == FDEVP_CNX_ERROR) {
 			goto error;
 		}
@@ -1020,8 +1068,12 @@
 		CHECK_FCT( fd_sctps_init(conn) );
 #endif /* DISABLE_SCTP */
 	} else {
-		/* Set the socket info in the session */
-		gnutls_transport_set_ptr (conn->cc_tls_para.session, (gnutls_transport_ptr_t) (long) conn->cc_socket);
+		/* Set the transport pointer passed to push & pull callbacks */
+		gnutls_transport_set_ptr( conn->cc_tls_para.session, (gnutls_transport_ptr_t) conn );
+
+		/* Set the push and pull callbacks */
+		gnutls_transport_set_pull_function(conn->cc_tls_para.session, (void *)fd_cnx_s_recv);
+		gnutls_transport_set_push_function(conn->cc_tls_para.session, (void *)fd_cnx_s_send);
 	}
 
 	/* Handshake master session */
@@ -1253,15 +1305,6 @@
 	
 	conn->cc_closing = 1;
 	
-	/* Set a timeout on the socket so that in any case we are not stuck waiting for something */
-	if (conn->cc_socket > 0) {
-		struct timeval tv;
-		memset(&tv, 0, sizeof(tv));
-		tv.tv_sec = 3;	/* allow 3 seconds timeout for TLS session cleanup */
-		CHECK_SYS_DO( setsockopt(conn->cc_socket, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)), /* best effort only */ );
-		CHECK_SYS_DO( setsockopt(conn->cc_socket, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv)), /* best effort only */ );
-	}
-	
 	/* Initiate shutdown of the TLS session(s): call gnutls_bye(WR), then read until error */
 	if (conn->cc_tls) {
 #ifndef DISABLE_SCTP

--- a/freeDiameter/cnxctx.h	Fri Feb 05 18:43:33 2010 +0900
+++ b/freeDiameter/cnxctx.h	Mon Feb 08 17:46:40 2010 +0900
@@ -78,6 +78,10 @@
 	} 		cc_sctps_data;
 };
 
+/* Socket */
+ssize_t fd_cnx_s_recv(struct cnxctx * conn, void *buffer, size_t length);
+void fd_cnx_s_setto(int sock);
+
 /* TLS */
 int fd_tls_rcvthr_core(struct cnxctx * conn, gnutls_session_t session);
 int fd_tls_prepare(gnutls_session_t * session, int mode, char * priority, void * alt_creds);
@@ -99,7 +103,7 @@
 int fd_sctp_get_remote_ep(int sock, struct fd_list * list);
 int fd_sctp_get_str_info( int sock, uint16_t *in, uint16_t *out, sSS *primary );
 int fd_sctp_sendstr(int sock, uint16_t strid, uint8_t * buf, size_t len);
-int fd_sctp_recvmeta(int sock, uint16_t * strid, uint8_t ** buf, size_t * len, int *event);
+int fd_sctp_recvmeta(int sock, uint16_t * strid, uint8_t ** buf, size_t * len, int *event, int * cc_closing);
 
 /* TLS over SCTP (multi-stream) */
 struct sctps_ctx {

--- a/freeDiameter/sctp.c	Fri Feb 05 18:43:33 2010 +0900
+++ b/freeDiameter/sctp.c	Mon Feb 08 17:46:40 2010 +0900
@@ -1085,7 +1085,7 @@
 }
 
 /* Receive the next data from the socket, or next notification */
-int fd_sctp_recvmeta(int sock, uint16_t * strid, uint8_t ** buf, size_t * len, int *event)
+int fd_sctp_recvmeta(int sock, uint16_t * strid, uint8_t ** buf, size_t * len, int *event, int * cc_closing)
 {
 	ssize_t 		 ret = 0;
 	struct msghdr 		 mhdr;
@@ -1094,9 +1094,10 @@
 	uint8_t			*data = NULL;
 	size_t 			 bufsz = 0, datasize = 0;
 	size_t			 mempagesz = sysconf(_SC_PAGESIZE); /* We alloc buffer by memory pages for efficiency */
+	int 			 timedout = 0;
 	
-	TRACE_ENTRY("%d %p %p %p %p", sock, strid, buf, len, event);
-	CHECK_PARAMS( (sock > 0) && buf && len && event );
+	TRACE_ENTRY("%d %p %p %p %p %p", sock, strid, buf, len, event, cc_closing);
+	CHECK_PARAMS( (sock > 0) && buf && len && event && cc_closing );
 	
 	/* Cleanup out parameters */
 	*buf = NULL;
@@ -1123,12 +1124,24 @@
 	iov.iov_len  = bufsz - datasize;
 
 	/* Receive data from the socket */
+again:
 	pthread_cleanup_push(free, data);
 	ret = recvmsg(sock, &mhdr, 0);
 	pthread_cleanup_pop(0);
 	
+	/* First, handle timeouts (same as fd_cnx_s_recv) */
+	if ((ret < 0) && (errno == ETIMEDOUT)) {
+		if (!*cc_closing)
+			goto again; /* don't care, just ignore */
+		if (!timedout) {
+			timedout ++; /* allow for one timeout while closing */
+			goto again;
+		}
+		/* fallback to normal handling */
+	}
+	
 	/* Handle errors */
-	if (ret <= 0) { /* Socket is closed, or an error occurred */
+	if (ret <= 0) { /* Socket timedout, closed, or an error occurred */
 		CHECK_SYS_DO(ret, /* to log in case of error */);
 		free(data);
 		*event = FDEVP_CNX_ERROR;

--- a/freeDiameter/sctps.c	Fri Feb 05 18:43:33 2010 +0900
+++ b/freeDiameter/sctps.c	Mon Feb 08 17:46:40 2010 +0900
@@ -86,7 +86,7 @@
 	ASSERT( conn->cc_sctps_data.array );
 	
 	do {
-		CHECK_FCT_DO( fd_sctp_recvmeta(conn->cc_socket, &strid, &buf, &bufsz, &event), goto error );
+		CHECK_FCT_DO( fd_sctp_recvmeta(conn->cc_socket, &strid, &buf, &bufsz, &event, &conn->cc_closing), goto error );
 		switch (event) {
 			case FDEVP_CNX_MSG_RECV:
 				/* Demux this message in the appropriate fifo, another thread will pull, gnutls process, and send in target queue */