--- 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