netmgr: proper closing of TCP/TCPDNS connections, TCP quota support

lib/isc/include/isc/netmgr.h

@@ -193,6 +193,7 @@ isc_nm_tcp_listen(isc_nm_t *mgr,
 		  isc_nm_accept_cb_t cb,
 		  size_t extrahandlesize,
 		  void *cbarg,
+		  isc_quota_t *quota,
 		  isc_nmsocket_t **rv);
 
 
@@ -202,6 +203,7 @@ isc_nm_tcp_dnslisten(isc_nm_t *mgr,
 		     isc_nm_recv_cb_t cb,
 		     size_t extrahandlesize,
 		     void *arg,
+		     isc_quota_t *quota,
 		     isc_nmsocket_t **rv);
 
 void

lib/isc/netmgr/netmgr-int.h

@@ -112,6 +112,7 @@ typedef enum isc__netievent_type {
 	netievent_tcpstopread,
 	netievent_tcplisten,
 	netievent_tcpstoplisten,
+	netievent_tcpclose,
 } isc__netievent_type;
 
 typedef struct isc__netievent_stop {
@@ -207,9 +208,14 @@ typedef struct isc__netievent_tcplisten {
 typedef struct isc__netievent_tcpstoplisten {
 	isc__netievent_type	   type;
 	isc_nmsocket_t *	   socket;
-	isc__nm_uvreq_t *	   req;
 } isc__netievent_tcpstoplisten_t;
 
+typedef struct isc__netievent_tcpclose {
+	isc__netievent_type	   type;
+	isc_nmsocket_t *	   socket;
+	isc__nm_uvreq_t *	   req;
+} isc__netievent_tcpclose_t;
+
 typedef struct isc__netievent_tcpsend {
 	isc__netievent_type	   type;
 	isc_nmhandle_t		   handle;
@@ -273,29 +279,29 @@ typedef enum isc_nmsocket_type {
 #define VALID_NMSOCK(t)                 ISC_MAGIC_VALID(t, NMSOCK_MAGIC)
 struct isc_nmsocket {
 	/* Unlocked, RO */
-	int			 magic;
-	int			 tid;
-	isc_nmsocket_type	 type;
-	isc_nm_t *		 mgr;
-	isc_nmsocket_t *	 parent;
+	int			   magic;
+	int			   tid;
+	isc_nmsocket_type	   type;
+	isc_nm_t *		   mgr;
+	isc_nmsocket_t *	   parent;
+	isc_quota_t *		   quota;
+	bool			   overquota;
 	/* outer socket is for 'wrapped' sockets - e.g. tcpdns in tcp */
-	isc_nmsocket_t *	 outer;
-	isc_nmsocket_t *	 children;
-	int			 nchildren;
-	isc_nmiface_t *		 iface;
-	isc_nmhandle_t		 tcphandle;
+	isc_nmsocket_t *	   outer;
+	/* server socket for connections */
+	isc_nmsocket_t *	   server;
+	/* children sockets for multi-socket setups */
+	isc_nmsocket_t *	   children;
+	int			   nchildren;
+	isc_nmiface_t *		   iface;
+	isc_nmhandle_t		   tcphandle;
 
 	/* extra data allocated at the end of each isc_nmhandle_t */
-	size_t			 extrahandlesize;
+	size_t			   extrahandlesize;
 
 	/* libuv data */
-	uv_os_sock_t		 fd;
-	union {
-		uv_handle_t	   handle;
-		uv_stream_t	   stream;
-		uv_udp_t	   udp;
-		uv_tcp_t	   tcp;
-	} uv_handle;
+	uv_os_sock_t		   fd;
+	union uv_any_handle	   uv_handle;
 
 	/* Atomic */
 	/* Number of running (e.g. listening) children sockets */
@@ -386,6 +392,9 @@ isc__nmsocket_init(isc_nmsocket_t *socket,
 		   isc_nm_t *mgr,
 		   isc_nmsocket_type type);
 
+void
+isc__nmsocket_prep_destroy(isc_nmsocket_t *socket);
+
 /*
  * Send for UDP handle
  */
@@ -415,6 +424,9 @@ isc__nm_tcp_send(isc_nmhandle_t *handle,
 		 isc_nm_send_cb_t cb,
 		 void *cbarg);
 
+void
+isc__nm_tcp_close(isc_nmsocket_t *socket);
+
 
 /*
  * Async callbacks for TCP
@@ -431,7 +443,8 @@ void
 isc__nm_handle_tcpsend(isc__networker_t *worker, isc__netievent_t *ievent0);
 void
 isc__nm_handle_startread(isc__networker_t *worker, isc__netievent_t *ievent0);
-
+void
+isc__nm_handle_tcpclose(isc__networker_t *worker, isc__netievent_t *ievent0);
 /* static void
  * handle_stopread(isc__networker_t *worker, isc__netievent_t *ievent0);
  */
@@ -442,3 +455,5 @@ isc__nm_tcpdns_send(isc_nmhandle_t *handle,
 		    isc_region_t *region,
 		    isc_nm_send_cb_t cb,
 		    void *cbarg);
+void
+isc__nm_tcpdns_close(isc_nmsocket_t *socket);

lib/isc/netmgr/netmgr.c

@@ -22,6 +22,7 @@
 #include <isc/magic.h>
 #include <isc/mem.h>
 #include <isc/netmgr.h>
+#include <isc/quota.h>
 #include <isc/random.h>
 #include <isc/refcount.h>
 #include <isc/region.h>
@@ -319,6 +320,9 @@ async_cb(uv_async_t *handle) {
 		case netievent_tcpstoplisten:
 			isc__nm_handle_tcpstoplistening(worker, ievent);
 			break;
+		case netievent_tcpclose:
+			isc__nm_handle_tcpclose(worker, ievent);
+			break;
 		default:
 			INSIST(0);
 		}
@@ -387,10 +391,14 @@ isc__nmsocket_destroy(isc_nmsocket_t *socket, bool dofree) {
 		isc__nmsocket_destroy(&socket->children[i], false);
 	}
 
-	if (VALID_NMHANDLE(&socket->tcphandle) && socket->tcphandle.dofree != NULL) {
+	if (VALID_NMHANDLE(&socket->tcphandle) &&
+	    socket->tcphandle.dofree != NULL) {
 		socket->tcphandle.dofree(socket->tcphandle.opaque);
 	}
 	socket->tcphandle = (isc_nmhandle_t) {};
+	if (socket->quota != NULL) {
+		isc_quota_detach(&socket->quota);
+	}
 
 	ck_stack_entry_t *se;
 	while ((se = ck_stack_pop_mpmc(&socket->inactivehandles)) != NULL) {
@@ -425,28 +433,31 @@ isc__nmsocket_maybe_destroy(isc_nmsocket_t *socket) {
 
 /* The final external reference to the socket is gone, we can try destroying the
  * socket, but we have to wait for all the inflight handles to finish first. */
-static void
+void
 isc__nmsocket_prep_destroy(isc_nmsocket_t *socket) {
-	isc_mutex_lock(&socket->lock);
 	INSIST(socket->parent == NULL);
-	/* XXXWPK cmpxchg? */
-	INSIST(isc__nmsocket_active(socket));
-
 	atomic_store(&socket->active, false);
-	isc_mutex_unlock(&socket->lock);
 	/*
 	 * XXXWPK if we're here we already stopped listening, otherwise
 	 * we'd have a hanging reference from the listening process.
+	 *
+	 * If that's a regular socket we need to close it.
 	 */
-/*	switch (socket->type) {
- *      case isc_nm_udplistener:
- *              isc_nm_udp_stoplistening(socket);
- *              break;
- *      case isc_nm_udpsocket:
- *      default:
- *              break;
- *      } */
-	isc__nmsocket_maybe_destroy(socket);
+	if (atomic_load(&socket->closed)) {
+		isc__nmsocket_maybe_destroy(socket);
+	} else {
+		switch (socket->type) {
+		case isc_nm_tcpsocket:
+			isc__nm_tcp_close(socket);
+			break;
+		case isc_nm_tcpdnssocket:
+			isc__nm_tcpdns_close(socket);
+			break;
+		default: /* This socket does not require closing */
+			isc__nmsocket_maybe_destroy(socket);
+			break;
+		}
+	}
 }
 
 void
@@ -481,7 +492,8 @@ isc_nmsocket_detach(isc_nmsocket_t **socketp) {
 void
 isc__nmsocket_init(isc_nmsocket_t *socket,
 		   isc_nm_t *mgr,
-		   isc_nmsocket_type type) {
+		   isc_nmsocket_type type)
+{
 	*socket = (isc_nmsocket_t) {
 		.type = type,
 		.fd = -1

lib/isc/netmgr/tcp.c

@@ -22,6 +22,7 @@
 #include <isc/magic.h>
 #include <isc/mem.h>
 #include <isc/netmgr.h>
+#include <isc/quota.h>
 #include <isc/random.h>
 #include <isc/refcount.h>
 #include <isc/region.h>
@@ -151,10 +152,9 @@ tcp_connect_cb(uv_connect_t *uvreq, int status) {
 	if (status == 0) {
 		isc_sockaddr_t peer;
 		struct sockaddr_storage ss;
-		int l = sizeof(ss);
 		uv_tcp_getpeername(&socket->uv_handle.tcp,
 				   (struct sockaddr*) &ss,
-				   &l);
+				   &(int){sizeof(ss)});
 		isc_result_t result =
 			isc_sockaddr_fromsockaddr(&peer,
 						  (struct sockaddr*) &ss);
@@ -176,6 +176,7 @@ isc_nm_tcp_listen(isc_nm_t *mgr,
 		  isc_nm_accept_cb_t cb,
 		  size_t extrahandlesize,
 		  void *cbarg,
+		  isc_quota_t *quota,
 		  isc_nmsocket_t **rv)
 {
 	isc__netievent_tcplisten_t *ievent;
@@ -188,6 +189,16 @@ isc_nm_tcp_listen(isc_nm_t *mgr,
 	nsocket->rcb.accept = cb;
 	nsocket->rcbarg = cbarg;
 	nsocket->extrahandlesize = extrahandlesize;
+	if (quota != NULL) {
+		/*
+		 * We need to force it to make sure we get it attached.
+		 * An example failure mode would be server under attack
+		 * reconfiguring interfaces - that might cause weak attach
+		 * to fail and leave this listening socket without limits.
+		 * We can ignore the result.
+		 */
+		isc_quota_force(quota, &nsocket->quota);
+	}
 	nsocket->tid = isc_random_uniform(mgr->nworkers);
 	/*
 	 * Listening to TCP is rare enough not to care about the
@@ -277,9 +288,10 @@ isc_nm_read(isc_nmhandle_t *handle, isc_nm_recv_cb_t cb, void *cbarg) {
 	socket->rcb.recv = cb;
 	socket->rcbarg = cbarg; /* That's obviously broken... */
 	if (socket->tid == isc_nm_tid()) {
-		uv_read_start(&socket->uv_handle.stream,
-			      isc__nm_alloc_cb,
-			      read_cb);
+		int r = uv_read_start(&socket->uv_handle.stream,
+				      isc__nm_alloc_cb,
+				      read_cb);
+		INSIST(r == 0);
 	} else {
 		isc__netievent_startread_t *ievent =
 			isc__nm_get_ievent(socket->mgr,
@@ -309,7 +321,8 @@ read_cb(uv_stream_t *stream, ssize_t nread, const uv_buf_t*buf) {
 	if (nread < 0) {
 		isc__nm_free_uvbuf(socket, buf);
 		socket->rcb.recv(socket->rcbarg, &socket->tcphandle, NULL);
-		/* XXXWPK TODO clean up handles! */
+		/* XXXWPK TODO clean up handles, close the connection, reclaim
+		 * quota */
 		return;
 	}
 	isc_region_t region = { .base = (unsigned char *) buf->base,
@@ -322,46 +335,76 @@ read_cb(uv_stream_t *stream, ssize_t nread, const uv_buf_t*buf) {
 	isc__nm_free_uvbuf(socket, buf);
 }
 
-static void
-tcp_connection_cb(uv_stream_t *server, int status) {
-	(void) status; /* TODO */
-	isc_nmsocket_t *csocket, *ssocket;
-	isc__networker_t *worker;
-
-	ssocket = uv_handle_get_data((uv_handle_t*) server);
+static isc_result_t
+accept_connection(isc_nmsocket_t *ssocket) {
 	REQUIRE(VALID_NMSOCK(ssocket));
 	REQUIRE(ssocket->tid == isc_nm_tid());
 	if (!atomic_load_relaxed(&ssocket->active)) {
 		/* We're closing, bail */
-		return;
+		return (ISC_R_CANCELED);
 	}
 	INSIST(ssocket->rcb.accept != NULL);
 
-	worker = &ssocket->mgr->workers[isc_nm_tid()];
+	isc_quota_t *quota = NULL;
+	if (ssocket->quota != NULL) {
+		isc_result_t result = isc_quota_attach(ssocket->quota, &quota);
+		if (result != ISC_R_SUCCESS) {
+			return (result);
+		}
+	}
 
-	csocket = isc_mem_get(ssocket->mgr->mctx, sizeof(isc_nmsocket_t));
+	isc_nmsocket_t *csocket =
+		isc_mem_get(ssocket->mgr->mctx, sizeof(isc_nmsocket_t));
 	isc__nmsocket_init(csocket, ssocket->mgr, isc_nm_tcpsocket);
 	csocket->tid = isc_nm_tid();
 	csocket->extrahandlesize = ssocket->extrahandlesize;
+	csocket->quota = quota;
+	quota = NULL;
 
+	isc__networker_t *worker = &ssocket->mgr->workers[isc_nm_tid()];
 	uv_tcp_init(&worker->loop, &csocket->uv_handle.tcp);
-	int r = uv_accept(server, &csocket->uv_handle.stream);
+	int r = uv_accept(&ssocket->uv_handle.stream,
+			  &csocket->uv_handle.stream);
 	if (r != 0) {
-		return; /* XXXWPK TODO LOG */
+		if (csocket->quota != NULL) {
+			isc_quota_detach(&csocket->quota);
+		}
+		isc_mem_put(ssocket->mgr->mctx, csocket,
+			    sizeof(isc_nmsocket_t));
+		return (ISC_R_FAILURE); /* XXXWPK TODO translate ! */
 	}
+
+	isc_nmsocket_attach(ssocket, &csocket->server);
+
 	isc_sockaddr_t peer;
 	struct sockaddr_storage ss;
-	int l = sizeof(ss);
 	uv_tcp_getpeername(&csocket->uv_handle.tcp, (struct sockaddr*) &ss,
-			   &l);
+			   &(int){sizeof(ss)});
 	isc_result_t result =
 		isc_sockaddr_fromsockaddr(&peer, (struct sockaddr*) &ss);
 	INSIST(result == ISC_R_SUCCESS);
 	isc_nmhandle_t *handle = isc__nmhandle_get(csocket, &peer);
+
 	ssocket->rcb.accept(handle, ISC_R_SUCCESS, ssocket->rcbarg);
 	isc_nmsocket_detach(&csocket);
+	return (ISC_R_SUCCESS);
 }
 
+static void
+tcp_connection_cb(uv_stream_t *server, int status) {
+	(void) status;
+	isc_nmsocket_t *ssocket = uv_handle_get_data((uv_handle_t*) server);
+	isc_result_t result = accept_connection(ssocket);
+	if (result != ISC_R_SUCCESS) {
+		if (result == ISC_R_QUOTA || result == ISC_R_SOFTQUOTA) {
+			ssocket->overquota = true;
+		}
+		/* XXXWPK TODO LOG */
+	}
+}
+
+
+
 /*
  * isc__nm_tcp_send sends buf to a peer on a socket.
  */
@@ -424,8 +467,8 @@ tcp_send_cb(uv_write_t *req, int status) {
  */
 void
 isc__nm_handle_tcpsend(isc__networker_t *worker, isc__netievent_t *ievent0) {
-	isc__netievent_udpsend_t *ievent =
-		(isc__netievent_udpsend_t *) ievent0;
+	isc__netievent_tcpsend_t *ievent =
+		(isc__netievent_tcpsend_t *) ievent0;
 	INSIST(worker->id == ievent->handle.socket->tid);
 	isc_result_t result = tcp_send_direct(ievent->handle.socket,
 					      ievent->req);
@@ -454,3 +497,58 @@ tcp_send_direct(isc_nmsocket_t *socket, isc__nm_uvreq_t *req)
 		return (ISC_R_FAILURE);
 	}
 }
+
+static void
+tcp_close_cb(uv_handle_t *uvhandle) {
+	isc_nmsocket_t *socket = uvhandle->data;
+	INSIST(VALID_NMSOCK(socket));
+	atomic_store(&socket->closed, true);
+	isc__nmsocket_prep_destroy(socket);
+}
+
+static void
+tcp_close_direct(isc_nmsocket_t *socket) {
+	INSIST(VALID_NMSOCK(socket));
+	INSIST(socket->tid == isc_nm_tid());
+	INSIST(socket->type == isc_nm_tcpsocket);
+	if (socket->quota != NULL) {
+		isc_quota_detach(&socket->quota);
+		isc_nmsocket_t *ssocket = socket->server;
+		if (ssocket->overquota) {
+			/* XXXWPK TODO we should loop here */
+			isc_result_t result = accept_connection(ssocket);
+			if (result != ISC_R_QUOTA &&
+			    result != ISC_R_SOFTQUOTA) {
+				ssocket->overquota = false;
+			}
+		}
+	}
+	isc_nmsocket_detach(&socket->server);
+	uv_close(&socket->uv_handle.handle, tcp_close_cb);
+}
+
+void
+isc__nm_tcp_close(isc_nmsocket_t *socket) {
+	INSIST(VALID_NMSOCK(socket));
+	INSIST(socket->type == isc_nm_tcpsocket);
+	if (socket->tid == isc_nm_tid()) {
+		tcp_close_direct(socket);
+	} else {
+		/*
+		 * We need to create an event and pass it using async channel
+		 */
+		isc__netievent_tcpclose_t *ievent;
+		ievent = isc__nm_get_ievent(socket->mgr, netievent_tcpclose);
+		ievent->socket = socket;
+		isc__nm_enqueue_ievent(&socket->mgr->workers[socket->tid],
+				       (isc__netievent_t*) ievent);
+	}
+}
+
+void
+isc__nm_handle_tcpclose(isc__networker_t *worker, isc__netievent_t *ievent0) {
+	isc__netievent_tcpclose_t *ievent =
+		(isc__netievent_tcpclose_t *) ievent0;
+	INSIST(worker->id == ievent->socket->tid);
+	tcp_close_direct(ievent->socket);
+}

lib/isc/netmgr/tcpdns.c

@@ -71,6 +71,7 @@ dnslisten_readcb(void *arg, isc_nmhandle_t *handle, isc_region_t *region) {
 	isc_nmsocket_t *dnssocket = (isc_nmsocket_t*) arg;
 	if (region == NULL) {
 		/* Connection closed */
+		atomic_store(&dnssocket->closed, true);
 		isc_nmhandle_detach(&handle);
 		isc_nmsocket_detach(&dnssocket->outer);
 		isc_nmsocket_detach(&dnssocket);
@@ -98,6 +99,7 @@ isc_nm_tcp_dnslisten(isc_nm_t *mgr,
 		     isc_nm_recv_cb_t cb,
 		     size_t extrahandlesize,
 		     void *cbarg,
+		     isc_quota_t *quota,
 		     isc_nmsocket_t **rv)
 {
 	isc_result_t result;
@@ -117,6 +119,7 @@ isc_nm_tcp_dnslisten(isc_nm_t *mgr,
 				   dnslisten_acceptcb,
 				   extrahandlesize,
 				   dnslistensocket,
+				   quota,
 				   &dnslistensocket->outer);
 	dnslistensocket->listening = true;
 	*rv = dnslistensocket;
@@ -125,6 +128,7 @@ isc_nm_tcp_dnslisten(isc_nm_t *mgr,
 
 void
 isc_nm_tcpdns_stoplistening(isc_nmsocket_t *socket) {
+	INSIST(socket->type == isc_nm_tcpdnslistener);
 	isc_nm_tcp_stoplistening(socket->outer);
 	atomic_store(&socket->listening, false);
 	isc_nmsocket_detach(&socket->outer);
@@ -172,3 +176,10 @@ isc__nm_tcpdns_send(isc_nmhandle_t *handle,
 	isc_nmhandle_attach(handle, &t->orighandle);
 	return (isc__nm_tcp_send(t->handle, &t->region, tcpdnssend_cb, t));
 }
+
+void
+isc__nm_tcpdns_close(isc_nmsocket_t *socket) {
+	isc_nmsocket_detach(&socket->outer);
+	socket->closed = true;
+	isc__nmsocket_prep_destroy(socket);
+}

lib/isc/netmgr/udp.c

@@ -156,8 +156,10 @@ isc__nm_handle_udplisten(isc__networker_t *worker, isc__netievent_t *ievent0) {
 	uv_udp_bind(&socket->uv_handle.udp,
 		    &socket->parent->iface->addr.type.sa,
 		    0);
-	uv_recv_buffer_size(&socket->uv_handle.handle, &(int){16*1024*1024});
-	uv_send_buffer_size(&socket->uv_handle.handle, &(int){16*1024*1024});
+	uv_recv_buffer_size(&socket->uv_handle.handle,
+			    &(int){16 * 1024 * 1024});
+	uv_send_buffer_size(&socket->uv_handle.handle,
+			    &(int){16 * 1024 * 1024});
 	uv_udp_recv_start(&socket->uv_handle.udp, isc__nm_alloc_cb,
 			  udp_recv_cb);
 }

lib/ns/interfacemgr.c

@@ -463,6 +463,7 @@ ns_interface_listentcp(ns_interface_t *ifp) {
 				      ns__client_request,
 				      sizeof(ns_client_t),
 				      ifp,
+				      &ifp->mgr->sctx->tcpquota,
 				      &ifp->tcplistensocket);
 	if (result != ISC_R_SUCCESS) {
 		isc_log_write(IFMGR_COMMON_LOGARGS, ISC_LOG_ERROR,