peer: ensure destruction doesn't race

zx2c4 · zx2c4 · commit 2086e7003e28 · 2018-08-02T07:59:52.000+02:00
Completely rework peer removal to ensure peers don't jump between
contexts and create races.
diff --git a/src/cookie.c b/src/cookie.c
@@ -165,15 +165,9 @@ void cookie_message_consume(struct message_handshake_cookie *src, struct wiregua
 {
 	u8 cookie[COOKIE_LEN];
 	struct wireguard_peer *peer = NULL;
-	struct index_hashtable_entry *entry;
 	bool ret;
 
-	rcu_read_lock_bh();
-	entry = index_hashtable_lookup(&wg->index_hashtable, INDEX_HASHTABLE_HANDSHAKE | INDEX_HASHTABLE_KEYPAIR, src->receiver_index);
-	if (likely(entry))
-		peer = entry->peer;
-	rcu_read_unlock_bh();
-	if (unlikely(!peer))
+	if (unlikely(!index_hashtable_lookup(&wg->index_hashtable, INDEX_HASHTABLE_HANDSHAKE | INDEX_HASHTABLE_KEYPAIR, src->receiver_index, &peer)))
 		return;
 
 	down_read(&peer->latest_cookie.lock);
diff --git a/src/hashtables.c b/src/hashtables.c
@@ -152,7 +152,7 @@ void index_hashtable_remove(struct index_hashtable *table, struct index_hashtabl
 }
 
 /* Returns a strong reference to a entry->peer */
-struct index_hashtable_entry *index_hashtable_lookup(struct index_hashtable *table, const enum index_hashtable_type type_mask, const __le32 index)
+struct index_hashtable_entry *index_hashtable_lookup(struct index_hashtable *table, const enum index_hashtable_type type_mask, const __le32 index, struct wireguard_peer **peer)
 {
 	struct index_hashtable_entry *iter_entry, *entry = NULL;
 
@@ -166,7 +166,9 @@ struct index_hashtable_entry *index_hashtable_lookup(struct index_hashtable *tab
 	}
 	if (likely(entry)) {
 		entry->peer = peer_get_maybe_zero(entry->peer);
-		if (unlikely(!entry->peer))
+		if (likely(entry->peer))
+			*peer = entry->peer;
+		else
 			entry = NULL;
 	}
 	rcu_read_unlock_bh();
diff --git a/src/hashtables.h b/src/hashtables.h
@@ -47,6 +47,6 @@ void index_hashtable_init(struct index_hashtable *table);
 __le32 index_hashtable_insert(struct index_hashtable *table, struct index_hashtable_entry *entry);
 bool index_hashtable_replace(struct index_hashtable *table, struct index_hashtable_entry *old, struct index_hashtable_entry *new);
 void index_hashtable_remove(struct index_hashtable *table, struct index_hashtable_entry *entry);
-struct index_hashtable_entry *index_hashtable_lookup(struct index_hashtable *table, const enum index_hashtable_type type_mask, const __le32 index);
+struct index_hashtable_entry *index_hashtable_lookup(struct index_hashtable *table, const enum index_hashtable_type type_mask, const __le32 index, struct wireguard_peer **peer);
 
 #endif /* _WG_HASHTABLES_H */
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);
 }
@@ -542,7 +539,7 @@ bool noise_handshake_create_response(struct message_handshake_response *dst, str
 struct wireguard_peer *noise_handshake_consume_response(struct message_handshake_response *src, struct wireguard_device *wg)
 {
 	struct noise_handshake *handshake;
-	struct wireguard_peer *ret_peer = NULL;
+	struct wireguard_peer *peer = NULL, *ret_peer = NULL;
 	u8 key[NOISE_SYMMETRIC_KEY_LEN];
 	u8 hash[NOISE_HASH_LEN];
 	u8 chaining_key[NOISE_HASH_LEN];
@@ -556,7 +553,7 @@ struct wireguard_peer *noise_handshake_consume_response(struct message_handshake
 	if (unlikely(!wg->static_identity.has_identity))
 		goto out;
 
-	handshake = (struct noise_handshake *)index_hashtable_lookup(&wg->index_hashtable, INDEX_HASHTABLE_HANDSHAKE, src->receiver_index);
+	handshake = (struct noise_handshake *)index_hashtable_lookup(&wg->index_hashtable, INDEX_HASHTABLE_HANDSHAKE, src->receiver_index, &peer);
 	if (unlikely(!handshake))
 		goto out;
 
@@ -601,11 +598,11 @@ struct wireguard_peer *noise_handshake_consume_response(struct message_handshake
 	handshake->remote_index = src->sender_index;
 	handshake->state = HANDSHAKE_CONSUMED_RESPONSE;
 	up_write(&handshake->lock);
-	ret_peer = handshake->entry.peer;
+	ret_peer = peer;
 	goto out;
 
 fail:
-	peer_put(handshake->entry.peer);
+	peer_put(peer);
 out:
 	memzero_explicit(key, NOISE_SYMMETRIC_KEY_LEN);
 	memzero_explicit(hash, NOISE_HASH_LEN);
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,34 @@ 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, &peer);
+	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

Original file line number	Diff line number	Diff line change
`@@ -152,7 +152,7 @@ void index_hashtable_remove(struct index_hashtable *table, struct index_hashtabl`
`152`	`152`	`}`
`153`	`153`
`154`	`154`	`/* Returns a strong reference to a entry->peer */`
`155`		`-struct index_hashtable_entry index_hashtable_lookup(struct index_hashtable table, const enum index_hashtable_type type_mask, const __le32 index)`
	`155`	`+struct index_hashtable_entry index_hashtable_lookup(struct index_hashtable table, const enum index_hashtable_type type_mask, const __le32 index, struct wireguard_peer **peer)`
`156`	`156`	`{`
`157`	`157`	`struct index_hashtable_entry iter_entry, entry = NULL;`
`158`	`158`
`@@ -166,7 +166,9 @@ struct index_hashtable_entry index_hashtable_lookup(struct index_hashtable tab`
`166`	`166`	`}`
`167`	`167`	`if (likely(entry)) {`
`168`	`168`	`entry->peer = peer_get_maybe_zero(entry->peer);`
`169`		`- if (unlikely(!entry->peer))`
	`169`	`+ if (likely(entry->peer))`
	`170`	`+ *peer = entry->peer;`
	`171`	`+ else`
`170`	`172`	`entry = NULL;`
`171`	`173`	`}`
`172`	`174`	`rcu_read_unlock_bh();`