if_wg: remove M_WAITOK, check return codes on init

Here we remove all M_WAITOK checks, because we don't want to hang while
trying to allocate memory. It is better to return an error so the user
can try again later.

We also make sure to check all the return codes in peer and interface
allocation. The structure of those functions is:

1) Allocate all memory
2) Initialise fields in order of the struct
3) Cleanup gotos

Signed-off-by: Matt Dunwoodie <ncon@noconroy.net>

1 files changed, 90 insertions, 67 deletions

diff --git a/src/if_wg.c b/src/if_wg.c
index 0e5ea3f..dc6e7e4 100644
--- a/src/if_wg.c
+++ b/src/if_wg.c
@@ -379,8 +379,6 @@ static struct wg_packet *wg_queue_dequeue_serial(struct wg_queue *);
 static struct wg_packet *wg_queue_dequeue_parallel(struct wg_queue *);
 static void wg_input(struct mbuf *, int, struct inpcb *, const struct sockaddr *, void *);
 static void wg_peer_send_staged(struct wg_peer *);
-static void crypto_taskq_setup(struct wg_softc *);
-static void crypto_taskq_destroy(struct wg_softc *);
 static int wg_clone_create(struct if_clone *, int, caddr_t);
 static void wg_qflush(struct ifnet *);
 static int wg_transmit(struct ifnet *, struct mbuf *);
@@ -407,19 +405,25 @@ wg_peer_alloc(struct wg_softc *sc, const uint8_t pub_key[WG_KEY_SIZE])
 
 	sx_assert(&sc->sc_lock, SX_XLOCKED);
 
-	peer = malloc(sizeof(*peer), M_WG, M_WAITOK|M_ZERO);
-	peer->p_sc = sc;
+	if ((peer = malloc(sizeof(*peer), M_WG, M_NOWAIT|M_ZERO)) == NULL)
+		goto free_none;
+
+	if ((peer->p_remote = noise_remote_alloc(sc->sc_local, peer, pub_key)) == NULL)
+		goto free_peer;
+
+	if ((peer->p_tx_bytes = counter_u64_alloc(M_NOWAIT)) == NULL)
+		goto free_remote;
+
+	if ((peer->p_rx_bytes = counter_u64_alloc(M_NOWAIT)) == NULL)
+		goto free_tx_bytes;
+
 	peer->p_id = peer_counter++;
-	LIST_INIT(&peer->p_aips);
-	peer->p_aips_num = 0;
+	peer->p_sc = sc;
 
-	if ((peer->p_remote = noise_remote_alloc(sc->sc_local, peer, pub_key)) == NULL) {
-		free(peer, M_WG);
-		return NULL;
-	}
 	cookie_maker_init(&peer->p_cookie, pub_key);
 
 	rw_init(&peer->p_endpoint_lock, "wg_peer_endpoint");
+
 	wg_queue_init(&peer->p_stage_queue, "stageq");
 	wg_queue_init(&peer->p_encrypt_serial, "txq");
 	wg_queue_init(&peer->p_decrypt_serial, "rxq");
@@ -427,7 +431,6 @@ wg_peer_alloc(struct wg_softc *sc, const uint8_t pub_key[WG_KEY_SIZE])
 	peer->p_enabled = false;
 	peer->p_need_another_keepalive = false;
 	peer->p_persistent_keepalive_interval = 0;
-
 	callout_init(&peer->p_new_handshake, true);
 	callout_init(&peer->p_send_keepalive, true);
 	callout_init(&peer->p_retry_handshake, true);
@@ -443,10 +446,18 @@ wg_peer_alloc(struct wg_softc *sc, const uint8_t pub_key[WG_KEY_SIZE])
 	GROUPTASK_INIT(&peer->p_recv, 0, (gtask_fn_t *)wg_deliver_in, peer);
 	taskqgroup_attach(qgroup_wg_tqg, &peer->p_recv, peer, NULL, NULL, "wg recv");
 
-	peer->p_tx_bytes = counter_u64_alloc(M_WAITOK);
-	peer->p_rx_bytes = counter_u64_alloc(M_WAITOK);
+	LIST_INIT(&peer->p_aips);
+	peer->p_aips_num = 0;
 
 	return (peer);
+free_tx_bytes:
+	counter_u64_free(peer->p_tx_bytes);
+free_remote:
+	noise_remote_free(peer->p_remote, NULL);
+free_peer:
+	free(peer, M_WG);
+free_none:
+	return NULL;
 }
 
 static void
@@ -2221,9 +2232,10 @@ wgc_set(struct wg_softc *sc, struct wg_data_io *wgd)
 	if (wgd->wgd_size >= UINT32_MAX / 2)
 		return (E2BIG);
 
-	sx_xlock(&sc->sc_lock);
+	if ((nvlpacked = malloc(wgd->wgd_size, M_TEMP, M_NOWAIT)) == NULL)
+		return (ENOMEM);
 
-	nvlpacked = malloc(wgd->wgd_size, M_TEMP, M_WAITOK);
+	sx_xlock(&sc->sc_lock);
 	err = copyin(wgd->wgd_data, nvlpacked, wgd->wgd_size);
 	if (err)
 		goto out;
@@ -2345,7 +2357,10 @@ wgc_get(struct wg_softc *sc, struct wg_data_io *wgd)
 	}
 	peer_count = sc->sc_peers_num;
 	if (peer_count) {
-		nvl_peers = mallocarray(peer_count, sizeof(void *), M_NVLIST, M_WAITOK | M_ZERO);
+		if ((nvl_peers = mallocarray(peer_count, sizeof(void *), M_NVLIST, M_NOWAIT | M_ZERO)) == NULL) {
+			err = ENOMEM;
+			goto err;
+		}
 		i = 0;
 		TAILQ_FOREACH(peer, &sc->sc_peers, p_entry) {
 			if (i >= peer_count)
@@ -2374,7 +2389,10 @@ wgc_get(struct wg_softc *sc, struct wg_data_io *wgd)
 
 			aip_count = peer->p_aips_num;
 			if (aip_count) {
-				nvl_aips = mallocarray(aip_count, sizeof(void *), M_NVLIST, M_WAITOK | M_ZERO);
+				if ((nvl_aips = mallocarray(aip_count, sizeof(void *), M_NVLIST, M_NOWAIT | M_ZERO)) == NULL) {
+					err = ENOMEM;
+					goto err_peer;
+				}
 				j = 0;
 				LIST_FOREACH(aip, &peer->p_aips, a_entry) {
 					if (j >= aip_count)
@@ -2571,71 +2589,67 @@ wg_down(struct wg_softc *sc)
 	sx_xunlock(&sc->sc_lock);
 }
 
-static void
-crypto_taskq_setup(struct wg_softc *sc)
-{
-
-	sc->sc_encrypt = mallocarray(sizeof(struct grouptask), mp_ncpus, M_WG, M_WAITOK);
-	sc->sc_decrypt = mallocarray(sizeof(struct grouptask), mp_ncpus, M_WG, M_WAITOK);
-
-	for (int i = 0; i < mp_ncpus; i++) {
-		GROUPTASK_INIT(&sc->sc_encrypt[i], 0,
-		     (gtask_fn_t *)wg_softc_encrypt, sc);
-		taskqgroup_attach_cpu(qgroup_wg_tqg, &sc->sc_encrypt[i], sc, i, NULL, NULL, "wg encrypt");
-		GROUPTASK_INIT(&sc->sc_decrypt[i], 0,
-		    (gtask_fn_t *)wg_softc_decrypt, sc);
-		taskqgroup_attach_cpu(qgroup_wg_tqg, &sc->sc_decrypt[i], sc, i, NULL, NULL, "wg decrypt");
-	}
-}
-
-static void
-crypto_taskq_destroy(struct wg_softc *sc)
-{
-	for (int i = 0; i < mp_ncpus; i++) {
-		taskqgroup_detach(qgroup_wg_tqg, &sc->sc_encrypt[i]);
-		taskqgroup_detach(qgroup_wg_tqg, &sc->sc_decrypt[i]);
-	}
-	free(sc->sc_encrypt, M_WG);
-	free(sc->sc_decrypt, M_WG);
-}
-
 static int
 wg_clone_create(struct if_clone *ifc, int unit, caddr_t params)
 {
 	struct wg_softc *sc;
 	struct ifnet *ifp;
 
-	sc = malloc(sizeof(*sc), M_WG, M_NOWAIT | M_ZERO);
-	if (!sc)
-		goto err_mem;
-	sc->sc_local = noise_local_alloc(sc);
-	if (!sc->sc_local)
+	if ((sc = malloc(sizeof(*sc), M_WG, M_NOWAIT | M_ZERO)) == NULL)
+		goto free_none;
+
+	if ((sc->sc_local = noise_local_alloc(sc)) == NULL)
 		goto free_sc;
+
+	if ((sc->sc_encrypt = mallocarray(sizeof(struct grouptask), mp_ncpus, M_WG, M_NOWAIT)) == NULL)
+		goto free_local;
+
+	if ((sc->sc_decrypt = mallocarray(sizeof(struct grouptask), mp_ncpus, M_WG, M_NOWAIT)) == NULL)
+		goto free_encrypt;
+
 	if (!rn_inithead((void **)&sc->sc_aip4, offsetof(struct aip_addr, in) * NBBY))
-		goto free_noise_local;
+		goto free_decrypt;
+
 	if (!rn_inithead((void **)&sc->sc_aip6, offsetof(struct aip_addr, in6) * NBBY))
 		goto free_aip4;
 
+	atomic_add_int(&clone_count, 1);
+	ifp = sc->sc_ifp = if_alloc(IFT_WIREGUARD);
+
 	sc->sc_ucred = crhold(curthread->td_ucred);
 	sc->sc_socket.so_fibnum = curthread->td_proc->p_fibnum;
-	ifp = sc->sc_ifp = if_alloc(IFT_WIREGUARD);
-	ifp->if_softc = sc;
-	if_initname(ifp, wgname, unit);
+	sc->sc_socket.so_port = 0;
+
 	TAILQ_INIT(&sc->sc_peers);
 	sc->sc_peers_num = 0;
+
 	cookie_checker_init(&sc->sc_cookie);
-	sc->sc_socket.so_port = 0;
-	atomic_add_int(&clone_count, 1);
-	ifp->if_capabilities = ifp->if_capenable = WG_CAPS;
-	wg_queue_init(&sc->sc_handshake_queue, "hsq");
-	sx_init(&sc->sc_lock, "wg softc lock");
+
+	RADIX_NODE_HEAD_LOCK_INIT(sc->sc_aip4);
+	RADIX_NODE_HEAD_LOCK_INIT(sc->sc_aip6);
+
 	GROUPTASK_INIT(&sc->sc_handshake, 0, (gtask_fn_t *)wg_softc_handshake_receive, sc);
 	taskqgroup_attach(qgroup_wg_tqg, &sc->sc_handshake, sc, NULL, NULL, "wg tx initiation");
-	crypto_taskq_setup(sc);
+	wg_queue_init(&sc->sc_handshake_queue, "hsq");
+
+	for (int i = 0; i < mp_ncpus; i++) {
+		GROUPTASK_INIT(&sc->sc_encrypt[i], 0,
+		     (gtask_fn_t *)wg_softc_encrypt, sc);
+		taskqgroup_attach_cpu(qgroup_wg_tqg, &sc->sc_encrypt[i], sc, i, NULL, NULL, "wg encrypt");
+		GROUPTASK_INIT(&sc->sc_decrypt[i], 0,
+		    (gtask_fn_t *)wg_softc_decrypt, sc);
+		taskqgroup_attach_cpu(qgroup_wg_tqg, &sc->sc_decrypt[i], sc, i, NULL, NULL, "wg decrypt");
+	}
+
 	wg_queue_init(&sc->sc_encrypt_parallel, "encp");
 	wg_queue_init(&sc->sc_decrypt_parallel, "decp");
-	RADIX_NODE_HEAD_LOCK_INIT(sc->sc_aip4);
-	RADIX_NODE_HEAD_LOCK_INIT(sc->sc_aip6);
+
+	sx_init(&sc->sc_lock, "wg softc lock");
+
+	ifp->if_softc = sc;
+	ifp->if_capabilities = ifp->if_capenable = WG_CAPS;
+	if_initname(ifp, wgname, unit);
+
 	if_setmtu(ifp, DEFAULT_MTU);
 	ifp->if_flags = IFF_NOARP | IFF_MULTICAST;
 	ifp->if_init = wg_init;
@@ -2654,14 +2668,18 @@ wg_clone_create(struct if_clone *ifc, int unit, caddr_t params)
 	LIST_INSERT_HEAD(&wg_list, sc, sc_entry);
 	sx_xunlock(&wg_sx);
 	return (0);
-
 free_aip4:
-	rn_detachhead((void **)&sc->sc_aip4);
-free_noise_local:
+	RADIX_NODE_HEAD_DESTROY(sc->sc_aip4);
+	free(sc->sc_aip4, M_RTABLE);
+free_decrypt:
+	free(sc->sc_decrypt, M_WG);
+free_encrypt:
+	free(sc->sc_encrypt, M_WG);
+free_local:
 	noise_local_free(sc->sc_local, NULL);
 free_sc:
 	free(sc, M_WG);
-err_mem:
+free_none:
 	return (ENOMEM);
 }
 
@@ -2712,7 +2730,12 @@ wg_clone_destroy(struct ifnet *ifp)
 	sx_xunlock(&sc->sc_lock);
 	sx_destroy(&sc->sc_lock);
 	taskqgroup_detach(qgroup_wg_tqg, &sc->sc_handshake);
-	crypto_taskq_destroy(sc);
+	for (int i = 0; i < mp_ncpus; i++) {
+		taskqgroup_detach(qgroup_wg_tqg, &sc->sc_encrypt[i]);
+		taskqgroup_detach(qgroup_wg_tqg, &sc->sc_decrypt[i]);
+	}
+	free(sc->sc_encrypt, M_WG);
+	free(sc->sc_decrypt, M_WG);
 	wg_queue_deinit(&sc->sc_handshake_queue);
 	wg_queue_deinit(&sc->sc_encrypt_parallel);
 	wg_queue_deinit(&sc->sc_decrypt_parallel);