peer: ensure destruction doesn't race

zx2c4 · zx2c4 · commit 7ebafe22dd9d · 2018-08-02T07:40:26.000+02:00
Completely rework peer removal to ensure peers don't jump between
contexts and create races.
diff --git a/src/noise.c b/src/noise.c
@@ -103,16 +103,13 @@ static struct noise_keypair *keypair_create(struct wireguard_peer *peer)
 
 static void keypair_free_rcu(struct rcu_head *rcu)
 {
-	struct noise_keypair *keypair = container_of(rcu, struct noise_keypair, rcu);
-
-	net_dbg_ratelimited("%s: Keypair %llu destroyed for peer %llu\n", keypair->entry.peer->device->dev->name, keypair->internal_id, keypair->entry.peer->internal_id);
-	kzfree(keypair);
+	kzfree(container_of(rcu, struct noise_keypair, rcu));
 }
 
 static void keypair_free_kref(struct kref *kref)
 {
 	struct noise_keypair *keypair = container_of(kref, struct noise_keypair, refcount);
-
+	net_dbg_ratelimited("%s: Keypair %llu destroyed for peer %llu\n", keypair->entry.peer->device->dev->name, keypair->internal_id, keypair->entry.peer->internal_id);
 	index_hashtable_remove(&keypair->entry.peer->device->index_hashtable, &keypair->entry);
 	call_rcu_bh(&keypair->rcu, keypair_free_rcu);
 }
diff --git a/src/peer.c b/src/peer.c
@@ -48,6 +48,7 @@ struct wireguard_peer *peer_create(struct wireguard_device *wg, const u8 public_
 	spin_lock_init(&peer->keypairs.keypair_update_lock);
 	INIT_WORK(&peer->transmit_handshake_work, packet_handshake_send_worker);
 	rwlock_init(&peer->endpoint_lock);
+	atomic_set(&peer->dead_count, 1);
 	kref_init(&peer->refcount);
 	skb_queue_head_init(&peer->staged_packet_queue);
 	peer->last_sent_handshake = ktime_get_boot_fast_ns() - (u64)(REKEY_TIMEOUT + 1) * NSEC_PER_SEC;
@@ -86,27 +87,47 @@ void peer_remove(struct wireguard_peer *peer)
 	if (unlikely(!peer))
 		return;
 	lockdep_assert_held(&peer->device->device_update_lock);
+
+	/* Remove from configuration-time lookup structures so new packets can't enter. */
+	list_del_init(&peer->peer_list);
 	allowedips_remove_by_peer(&peer->device->peer_allowedips, peer, &peer->device->device_update_lock);
 	pubkey_hashtable_remove(&peer->device->peer_hashtable, peer);
-	skb_queue_purge(&peer->staged_packet_queue);
+
+	/* Mark as dead, so that we don't allow jumping contexts after. */
+	while (atomic_cmpxchg(&peer->dead_count, 1, 0) != 1) cpu_relax();
+
+	/* The transition between packet encryption/decryption queues isn't guarded
+	 * by the dead_count, but each reference's life is strictly bounded by
+	 * two generations: once for parallel crypto and once for serial ingestion,
+	 * so we can simply flush twice, and be sure that we no longer have references
+	 * inside these queues.
+	 */
+
+	/* The first flush is for encrypt/decrypt. */
+	flush_workqueue(peer->device->packet_crypt_wq);
+	/* The second.1 flush is for send (but not receive, since that's napi). */
+	flush_workqueue(peer->device->packet_crypt_wq);
+	/* The second.2 flush is for receive (but not send, since that's wq). */
+	napi_disable(&peer->napi);
+	/* It's now safe to remove the napi struct, which must be done here from process context. */
+	netif_napi_del(&peer->napi);
+	/* Ensure any workstructs we own (like transmit_handshake_work or clear_peer_work) no longer are in use. */
+	flush_workqueue(peer->device->handshake_send_wq);
+
+	/* Remove keys and handshakes from memory. Handshake removal must be done here from process context. */
 	noise_handshake_clear(&peer->handshake);
 	noise_keypairs_clear(&peer->keypairs);
-	list_del_init(&peer->peer_list);
+
+	/* Destroy all ongoing timers that were in-flight at the beginning of this function. */
 	timers_stop(peer);
-	flush_workqueue(peer->device->packet_crypt_wq); /* The first flush is for encrypt/decrypt. */
-	flush_workqueue(peer->device->packet_crypt_wq); /* The second.1 flush is for send (but not receive, since that's napi). */
-	napi_disable(&peer->napi); /* The second.2 flush is for receive (but not send, since that's wq). */
-	flush_workqueue(peer->device->handshake_send_wq);
-	netif_napi_del(&peer->napi);
+
 	--peer->device->num_peers;
 	peer_put(peer);
 }
 
 static void rcu_release(struct rcu_head *rcu)
 {
 	struct wireguard_peer *peer = container_of(rcu, struct wireguard_peer, rcu);
-
-	pr_debug("%s: Peer %llu (%pISpfsc) destroyed\n", peer->device->dev->name, peer->internal_id, &peer->endpoint.addr);
 	dst_cache_destroy(&peer->endpoint_cache);
 	packet_queue_free(&peer->rx_queue, false);
 	packet_queue_free(&peer->tx_queue, false);
@@ -116,9 +137,12 @@ static void rcu_release(struct rcu_head *rcu)
 static void kref_release(struct kref *refcount)
 {
 	struct wireguard_peer *peer = container_of(refcount, struct wireguard_peer, refcount);
-
+	pr_debug("%s: Peer %llu (%pISpfsc) destroyed\n", peer->device->dev->name, peer->internal_id, &peer->endpoint.addr);
+	/* Remove ourself from dynamic runtime lookup structures, now that the last reference is gone. */
 	index_hashtable_remove(&peer->device->index_hashtable, &peer->handshake.entry);
+	/* Remove any lingering packets that didn't have a chance to be transmitted. */
 	skb_queue_purge(&peer->staged_packet_queue);
+	/* Free the memory used. */
 	call_rcu_bh(&peer->rcu, rcu_release);
 }
 
diff --git a/src/peer.h b/src/peer.h
@@ -54,6 +54,7 @@ struct wireguard_peer {
 	bool timers_enabled, timer_need_another_keepalive, sent_lastminute_handshake;
 	struct timespec walltime_last_handshake;
 	struct kref refcount;
+	atomic_t dead_count;
 	struct rcu_head rcu;
 	struct list_head peer_list;
 	u64 internal_id;
diff --git a/src/receive.c b/src/receive.c
@@ -120,6 +120,8 @@ static void receive_handshake_packet(struct wireguard_device *wg, struct sk_buff
 			net_dbg_skb_ratelimited("%s: Invalid handshake initiation from %pISpfsc\n", wg->dev->name, skb);
 			return;
 		}
+		if (unlikely(!atomic_inc_not_zero(&peer->dead_count)))
+			goto err_dead;
 		socket_set_peer_endpoint_from_skb(peer, skb);
 		net_dbg_ratelimited("%s: Receiving handshake initiation from peer %llu (%pISpfsc)\n", wg->dev->name, peer->internal_id, &peer->endpoint.addr);
 		packet_send_handshake_response(peer);
@@ -137,6 +139,8 @@ static void receive_handshake_packet(struct wireguard_device *wg, struct sk_buff
 			net_dbg_skb_ratelimited("%s: Invalid handshake response from %pISpfsc\n", wg->dev->name, skb);
 			return;
 		}
+		if (unlikely(!atomic_inc_not_zero(&peer->dead_count)))
+			goto err_dead;
 		socket_set_peer_endpoint_from_skb(peer, skb);
 		net_dbg_ratelimited("%s: Receiving handshake response from peer %llu (%pISpfsc)\n", wg->dev->name, peer->internal_id, &peer->endpoint.addr);
 		if (noise_handshake_begin_session(&peer->handshake, &peer->keypairs)) {
@@ -164,6 +168,8 @@ static void receive_handshake_packet(struct wireguard_device *wg, struct sk_buff
 
 	timers_any_authenticated_packet_received(peer);
 	timers_any_authenticated_packet_traversal(peer);
+	atomic_dec(&peer->dead_count);
+err_dead:
 	peer_put(peer);
 }
 
@@ -282,9 +288,9 @@ static inline bool counter_validate(union noise_counter *counter, u64 their_coun
 }
 #include "selftest/counter.h"
 
-static void packet_consume_data_done(struct sk_buff *skb, struct endpoint *endpoint)
+static void packet_consume_data_done(struct wireguard_peer *peer, struct sk_buff *skb, struct endpoint *endpoint)
 {
-	struct wireguard_peer *peer = PACKET_PEER(skb), *routed_peer;
+	struct wireguard_peer *routed_peer;
 	struct net_device *dev = peer->device->dev;
 	unsigned int len, len_before_trim;
 
@@ -400,7 +406,7 @@ int packet_rx_poll(struct napi_struct *napi, int budget)
 			goto next;
 
 		skb_reset(skb);
-		packet_consume_data_done(skb, &endpoint);
+		packet_consume_data_done(peer, skb, &endpoint);
 		free = false;
 
 next:
@@ -436,32 +442,36 @@ void packet_decrypt_worker(struct work_struct *work)
 
 static void packet_consume_data(struct wireguard_device *wg, struct sk_buff *skb)
 {
-	struct wireguard_peer *peer;
+	struct wireguard_peer *peer = NULL;
 	__le32 idx = ((struct message_data *)skb->data)->key_idx;
 	int ret;
 
 	rcu_read_lock_bh();
-	PACKET_CB(skb)->keypair = noise_keypair_get((struct noise_keypair *)index_hashtable_lookup(&wg->index_hashtable, INDEX_HASHTABLE_KEYPAIR, idx));
-	rcu_read_unlock_bh();
-	if (unlikely(!PACKET_CB(skb)->keypair)) {
-		dev_kfree_skb(skb);
-		return;
+	PACKET_CB(skb)->keypair = (struct noise_keypair *)index_hashtable_lookup(&wg->index_hashtable, INDEX_HASHTABLE_KEYPAIR, idx);
+	if (likely(PACKET_CB(skb)->keypair))
+		peer = PACKET_PEER(skb);
+	if (unlikely(!noise_keypair_get(PACKET_CB(skb)->keypair))) {
+		rcu_read_unlock_bh();
+		goto err_keypair;
 	}
+	rcu_read_unlock_bh();
 
-	/* The call to index_hashtable_lookup gives us a reference to its underlying peer, so we don't need to call peer_get(). */
-	peer = PACKET_PEER(skb);
+	if (unlikely(list_empty(&peer->peer_list) || !atomic_inc_not_zero(&peer->dead_count)))
+		goto err;
 
+	peer_get(peer);
 	ret = queue_enqueue_per_device_and_peer(&wg->decrypt_queue, &peer->rx_queue, skb, wg->packet_crypt_wq, &wg->decrypt_queue.last_cpu);
-	if (likely(!ret))
-		return; /* Successful. No need to drop references below. */
-
-	if (ret == -EPIPE)
+	if (unlikely(ret == -EPIPE))
 		queue_enqueue_per_peer(&peer->rx_queue, skb, PACKET_STATE_DEAD);
-	else {
-		peer_put(peer);
-		noise_keypair_put(PACKET_CB(skb)->keypair);
-		dev_kfree_skb(skb);
-	}
+	atomic_dec(&peer->dead_count);
+	peer_put(peer);
+	if (likely(!ret || ret == -EPIPE))
+		return;
+err:
+	noise_keypair_put(PACKET_CB(skb)->keypair);
+err_keypair:
+	peer_put(peer);
+	dev_kfree_skb(skb);
 }
 
 void packet_receive(struct wireguard_device *wg, struct sk_buff *skb)
diff --git a/src/send.c b/src/send.c
@@ -58,13 +58,14 @@ void packet_send_queued_handshake_initiation(struct wireguard_peer *peer, bool i
 	/* First checking the timestamp here is just an optimization; it will
 	 * be caught while properly locked inside the actual work queue.
 	 */
-	if (!has_expired(peer->last_sent_handshake, REKEY_TIMEOUT))
+	if (!has_expired(peer->last_sent_handshake, REKEY_TIMEOUT) || unlikely(!atomic_inc_not_zero(&peer->dead_count)))
 		return;
 
 	peer_get(peer);
 	/* Queues up calling packet_send_queued_handshakes(peer), where we do a peer_put(peer) after: */
 	if (!queue_work(peer->device->handshake_send_wq, &peer->transmit_handshake_work))
 		peer_put(peer); /* If the work was already queued, we want to drop the extra reference */
+	atomic_dec(&peer->dead_count);
 }
 
 void packet_send_handshake_response(struct wireguard_peer *peer)
@@ -268,17 +269,21 @@ static void packet_create_data(struct sk_buff *first)
 	struct wireguard_device *wg = peer->device;
 	int ret;
 
-	ret = queue_enqueue_per_device_and_peer(&wg->encrypt_queue, &peer->tx_queue, first, wg->packet_crypt_wq, &wg->encrypt_queue.last_cpu);
-	if (likely(!ret))
-		return; /* Successful. No need to fall through to drop references below. */
+	if (unlikely(!atomic_inc_not_zero(&peer->dead_count)))
+		goto err;
 
-	if (ret == -EPIPE)
+	peer_get(peer);
+	ret = queue_enqueue_per_device_and_peer(&wg->encrypt_queue, &peer->tx_queue, first, wg->packet_crypt_wq, &wg->encrypt_queue.last_cpu);
+	if (unlikely(ret == -EPIPE))
 		queue_enqueue_per_peer(&peer->tx_queue, first, PACKET_STATE_DEAD);
-	else {
-		peer_put(peer);
-		noise_keypair_put(PACKET_CB(first)->keypair);
-		skb_free_null_queue(first);
-	}
+	atomic_dec(&peer->dead_count);
+	peer_put(peer);
+	if (likely(!ret || ret == -EPIPE))
+		return;
+err:
+	noise_keypair_put(PACKET_CB(first)->keypair);
+	peer_put(peer);
+	skb_free_null_queue(first);
 }
 
 void packet_send_staged_packets(struct wireguard_peer *peer)
diff --git a/src/timers.c b/src/timers.c
@@ -27,9 +27,20 @@
 	if (unlikely(!peer)) \
 		return;
 
-static inline bool timers_active(struct wireguard_peer *peer)
+static inline void mod_peer_timer(struct wireguard_peer *peer, struct timer_list *timer, unsigned long expires)
 {
-	return netif_running(peer->device->dev) && !list_empty(&peer->peer_list);
+	if (unlikely(!netif_running(peer->device->dev) || !atomic_inc_not_zero(&peer->dead_count)))
+		return;
+	mod_timer(timer, expires);
+	atomic_dec(&peer->dead_count);
+}
+
+static inline void del_peer_timer(struct wireguard_peer *peer, struct timer_list *timer)
+{
+	if (unlikely(!netif_running(peer->device->dev) || !atomic_inc_not_zero(&peer->dead_count)))
+		return;
+	del_timer(timer);
+	atomic_dec(&peer->dead_count);
 }
 
 static void expired_retransmit_handshake(struct timer_list *timer)
@@ -39,8 +50,7 @@ static void expired_retransmit_handshake(struct timer_list *timer)
 	if (peer->timer_handshake_attempts > MAX_TIMER_HANDSHAKES) {
 		pr_debug("%s: Handshake for peer %llu (%pISpfsc) did not complete after %d attempts, giving up\n", peer->device->dev->name, peer->internal_id, &peer->endpoint.addr, MAX_TIMER_HANDSHAKES + 2);
 
-		if (likely(timers_active(peer)))
-			del_timer(&peer->timer_send_keepalive);
+		del_peer_timer(peer, &peer->timer_send_keepalive);
 		/* We drop all packets without a keypair and don't try again,
 		 * if we try unsuccessfully for too long to make a handshake.
 		 */
@@ -49,8 +59,8 @@ static void expired_retransmit_handshake(struct timer_list *timer)
 		/* We set a timer for destroying any residue that might be left
 		 * of a partial exchange.
 		 */
-		if (likely(timers_active(peer)) && !timer_pending(&peer->timer_zero_key_material))
-			mod_timer(&peer->timer_zero_key_material, jiffies + REJECT_AFTER_TIME * 3 * HZ);
+		if (!timer_pending(&peer->timer_zero_key_material))
+			mod_peer_timer(peer, &peer->timer_zero_key_material, jiffies + REJECT_AFTER_TIME * 3 * HZ);
 	} else {
 		++peer->timer_handshake_attempts;
 		pr_debug("%s: Handshake for peer %llu (%pISpfsc) did not complete after %d seconds, retrying (try %d)\n", peer->device->dev->name, peer->internal_id, &peer->endpoint.addr, REKEY_TIMEOUT, peer->timer_handshake_attempts + 1);
@@ -70,8 +80,7 @@ static void expired_send_keepalive(struct timer_list *timer)
 	packet_send_keepalive(peer);
 	if (peer->timer_need_another_keepalive) {
 		peer->timer_need_another_keepalive = false;
-		if (likely(timers_active(peer)))
-			mod_timer(&peer->timer_send_keepalive, jiffies + KEEPALIVE_TIMEOUT * HZ);
+		mod_peer_timer(peer, &peer->timer_send_keepalive, jiffies + KEEPALIVE_TIMEOUT * HZ);
 	}
 	peer_put(peer);
 }
@@ -91,8 +100,14 @@ static void expired_zero_key_material(struct timer_list *timer)
 {
 	peer_get_from_timer(timer_zero_key_material);
 
-	if (!queue_work(peer->device->handshake_send_wq, &peer->clear_peer_work)) /* Takes our reference. */
+	if (unlikely(!atomic_inc_not_zero(&peer->dead_count)))
+		goto out;
+	peer_get(peer);
+	if (!queue_work(peer->device->handshake_send_wq, &peer->clear_peer_work))
 		peer_put(peer); /* If the work was already on the queue, we want to drop the extra reference */
+	atomic_dec(&peer->dead_count);
+out:
+	peer_put(peer);
 }
 static void queued_expired_zero_key_material(struct work_struct *work)
 {
@@ -116,16 +131,16 @@ static void expired_send_persistent_keepalive(struct timer_list *timer)
 /* Should be called after an authenticated data packet is sent. */
 void timers_data_sent(struct wireguard_peer *peer)
 {
-	if (likely(timers_active(peer)) && !timer_pending(&peer->timer_new_handshake))
-		mod_timer(&peer->timer_new_handshake, jiffies + (KEEPALIVE_TIMEOUT + REKEY_TIMEOUT) * HZ);
+	if (!timer_pending(&peer->timer_new_handshake))
+		mod_peer_timer(peer, &peer->timer_new_handshake, jiffies + (KEEPALIVE_TIMEOUT + REKEY_TIMEOUT) * HZ);
 }
 
 /* Should be called after an authenticated data packet is received. */
 void timers_data_received(struct wireguard_peer *peer)
 {
-	if (likely(timers_active(peer))) {
+	if (likely(netif_running(peer->device->dev))) {
 		if (!timer_pending(&peer->timer_send_keepalive))
-			mod_timer(&peer->timer_send_keepalive, jiffies + KEEPALIVE_TIMEOUT * HZ);
+			mod_peer_timer(peer, &peer->timer_send_keepalive, jiffies + KEEPALIVE_TIMEOUT * HZ);
 		else
 			peer->timer_need_another_keepalive = true;
 	}
@@ -134,29 +149,25 @@ void timers_data_received(struct wireguard_peer *peer)
 /* Should be called after any type of authenticated packet is sent -- keepalive, data, or handshake. */
 void timers_any_authenticated_packet_sent(struct wireguard_peer *peer)
 {
-	if (likely(timers_active(peer)))
-		del_timer(&peer->timer_send_keepalive);
+	del_peer_timer(peer, &peer->timer_send_keepalive);
 }
 
 /* Should be called after any type of authenticated packet is received -- keepalive, data, or handshake. */
 void timers_any_authenticated_packet_received(struct wireguard_peer *peer)
 {
-	if (likely(timers_active(peer)))
-		del_timer(&peer->timer_new_handshake);
+	del_peer_timer(peer, &peer->timer_new_handshake);
 }
 
 /* Should be called after a handshake initiation message is sent. */
 void timers_handshake_initiated(struct wireguard_peer *peer)
 {
-	if (likely(timers_active(peer)))
-		mod_timer(&peer->timer_retransmit_handshake, jiffies + REKEY_TIMEOUT * HZ + prandom_u32_max(REKEY_TIMEOUT_JITTER_MAX_JIFFIES));
+	mod_peer_timer(peer, &peer->timer_retransmit_handshake, jiffies + REKEY_TIMEOUT * HZ + prandom_u32_max(REKEY_TIMEOUT_JITTER_MAX_JIFFIES));
 }
 
 /* Should be called after a handshake response message is received and processed or when getting key confirmation via the first data message. */
 void timers_handshake_complete(struct wireguard_peer *peer)
 {
-	if (likely(timers_active(peer)))
-		del_timer(&peer->timer_retransmit_handshake);
+	del_peer_timer(peer, &peer->timer_retransmit_handshake);
 	peer->timer_handshake_attempts = 0;
 	peer->sent_lastminute_handshake = false;
 	getnstimeofday(&peer->walltime_last_handshake);
@@ -165,15 +176,14 @@ void timers_handshake_complete(struct wireguard_peer *peer)
 /* Should be called after an ephemeral key is created, which is before sending a handshake response or after receiving a handshake response. */
 void timers_session_derived(struct wireguard_peer *peer)
 {
-	if (likely(timers_active(peer)))
-		mod_timer(&peer->timer_zero_key_material, jiffies + REJECT_AFTER_TIME * 3 * HZ);
+	mod_peer_timer(peer, &peer->timer_zero_key_material, jiffies + REJECT_AFTER_TIME * 3 * HZ);
 }
 
 /* Should be called before a packet with authentication -- keepalive, data, or handshake -- is sent, or after one is received. */
 void timers_any_authenticated_packet_traversal(struct wireguard_peer *peer)
 {
-	if (peer->persistent_keepalive_interval && likely(timers_active(peer)))
-		mod_timer(&peer->timer_persistent_keepalive, jiffies + peer->persistent_keepalive_interval * HZ);
+	if (peer->persistent_keepalive_interval)
+		mod_peer_timer(peer, &peer->timer_persistent_keepalive, jiffies + peer->persistent_keepalive_interval * HZ);
 }
 
 void timers_init(struct wireguard_peer *peer)

Original file line number	Diff line number	Diff line change
`@@ -103,16 +103,13 @@ static struct noise_keypair keypair_create(struct wireguard_peer peer)`
`103`	`103`
`104`	`104`	`static void keypair_free_rcu(struct rcu_head *rcu)`
`105`	`105`	`{`
`106`		`- struct noise_keypair *keypair = container_of(rcu, struct noise_keypair, rcu);`
`107`		`-`
`108`		`- net_dbg_ratelimited("%s: Keypair %llu destroyed for peer %llu\n", keypair->entry.peer->device->dev->name, keypair->internal_id, keypair->entry.peer->internal_id);`
`109`		`- kzfree(keypair);`
	`106`	`+ kzfree(container_of(rcu, struct noise_keypair, rcu));`
`110`	`107`	`}`
`111`	`108`
`112`	`109`	`static void keypair_free_kref(struct kref *kref)`
`113`	`110`	`{`
`114`	`111`	`struct noise_keypair *keypair = container_of(kref, struct noise_keypair, refcount);`
`115`		`-`
	`112`	`+ net_dbg_ratelimited("%s: Keypair %llu destroyed for peer %llu\n", keypair->entry.peer->device->dev->name, keypair->internal_id, keypair->entry.peer->internal_id);`
`116`	`113`	`index_hashtable_remove(&keypair->entry.peer->device->index_hashtable, &keypair->entry);`
`117`	`114`	`call_rcu_bh(&keypair->rcu, keypair_free_rcu);`
`118`	`115`	`}`