receive: hoist fpu outside of receive loop

Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>

3 files changed, 19 insertions, 18 deletions

diff --git a/src/crypto/chacha20poly1305.c b/src/crypto/chacha20poly1305.c
index 2ce7cbb..a9c3bf8 100644
--- a/src/crypto/chacha20poly1305.c
+++ b/src/crypto/chacha20poly1305.c
@@ -792,7 +792,8 @@ bool chacha20poly1305_decrypt(u8 *dst, const u8 *src, const size_t src_len,
 
 bool chacha20poly1305_decrypt_sg(struct scatterlist *dst, struct scatterlist *src, const size_t src_len,
 				 const u8 *ad, const size_t ad_len,
-				 const u64 nonce, const u8 key[CHACHA20POLY1305_KEYLEN])
+				 const u64 nonce, const u8 key[CHACHA20POLY1305_KEYLEN],
+				 bool have_simd)
 {
 	struct poly1305_ctx poly1305_state;
 	struct chacha20_ctx chacha20_state;
@@ -803,13 +804,10 @@ bool chacha20poly1305_decrypt_sg(struct scatterlist *dst, struct scatterlist *sr
 	size_t dst_len;
 	__le64 len;
 	__le64 le_nonce = cpu_to_le64(nonce);
-	bool have_simd;
 
 	if (unlikely(src_len < POLY1305_MAC_SIZE))
 		return false;
 
-	have_simd = chacha20poly1305_init_simd();
-
 	chacha20_keysetup(&chacha20_state, key, (u8 *)&le_nonce);
 
 	chacha20_crypt(&chacha20_state, block0, block0, sizeof(block0), have_simd);
@@ -856,7 +854,6 @@ err:
 	memzero_explicit(read_mac, POLY1305_MAC_SIZE);
 	memzero_explicit(computed_mac, POLY1305_MAC_SIZE);
 	memzero_explicit(&chacha20_state, sizeof(chacha20_state));
-	chacha20poly1305_deinit_simd(have_simd);
 	return !ret;
 }
 
diff --git a/src/crypto/chacha20poly1305.h b/src/crypto/chacha20poly1305.h
index c8885e0..991755d 100644
--- a/src/crypto/chacha20poly1305.h
+++ b/src/crypto/chacha20poly1305.h
@@ -20,17 +20,18 @@ void chacha20poly1305_encrypt(u8 *dst, const u8 *src, const size_t src_len,
 			      const u64 nonce, const u8 key[CHACHA20POLY1305_KEYLEN]);
 
 bool __must_check chacha20poly1305_encrypt_sg(struct scatterlist *dst, struct scatterlist *src, const size_t src_len,
-				 const u8 *ad, const size_t ad_len,
-				 const u64 nonce, const u8 key[CHACHA20POLY1305_KEYLEN],
-				 bool have_simd);
+					      const u8 *ad, const size_t ad_len,
+					      const u64 nonce, const u8 key[CHACHA20POLY1305_KEYLEN],
+					      bool have_simd);
 
 bool __must_check chacha20poly1305_decrypt(u8 *dst, const u8 *src, const size_t src_len,
-			      const u8 *ad, const size_t ad_len,
-			      const u64 nonce, const u8 key[CHACHA20POLY1305_KEYLEN]);
+					   const u8 *ad, const size_t ad_len,
+					   const u64 nonce, const u8 key[CHACHA20POLY1305_KEYLEN]);
 
 bool __must_check chacha20poly1305_decrypt_sg(struct scatterlist *dst, struct scatterlist *src, const size_t src_len,
-				 const u8 *ad, const size_t ad_len,
-				 const u64 nonce, const u8 key[CHACHA20POLY1305_KEYLEN]);
+					      const u8 *ad, const size_t ad_len,
+					      const u64 nonce, const u8 key[CHACHA20POLY1305_KEYLEN],
+					      bool have_simd);
 
 void xchacha20poly1305_encrypt(u8 *dst, const u8 *src, const size_t src_len,
 			       const u8 *ad, const size_t ad_len,
@@ -38,9 +39,9 @@ void xchacha20poly1305_encrypt(u8 *dst, const u8 *src, const size_t src_len,
 			       const u8 key[CHACHA20POLY1305_KEYLEN]);
 
 bool __must_check xchacha20poly1305_decrypt(u8 *dst, const u8 *src, const size_t src_len,
-			       const u8 *ad, const size_t ad_len,
-			       const u8 nonce[XCHACHA20POLY1305_NONCELEN],
-			       const u8 key[CHACHA20POLY1305_KEYLEN]);
+					    const u8 *ad, const size_t ad_len,
+					    const u8 nonce[XCHACHA20POLY1305_NONCELEN],
+					    const u8 key[CHACHA20POLY1305_KEYLEN]);
 
 #if defined(CONFIG_X86_64)
 #include <linux/version.h>
diff --git a/src/receive.c b/src/receive.c
index 2e91b67..7f9d1d7 100644
--- a/src/receive.c
+++ b/src/receive.c
@@ -196,7 +196,7 @@ static inline void keep_key_fresh(struct wireguard_peer *peer)
 	}
 }
 
-static inline bool skb_decrypt(struct sk_buff *skb, struct noise_symmetric_key *key)
+static inline bool skb_decrypt(struct sk_buff *skb, struct noise_symmetric_key *key, bool have_simd)
 {
 	struct scatterlist sg[MAX_SKB_FRAGS * 2 + 1];
 	struct sk_buff *trailer;
@@ -229,7 +229,7 @@ static inline bool skb_decrypt(struct sk_buff *skb, struct noise_symmetric_key *
 	if (skb_to_sgvec(skb, sg, 0, skb->len) <= 0)
 		return false;
 
-	if (!chacha20poly1305_decrypt_sg(sg, sg, skb->len, NULL, 0, PACKET_CB(skb)->nonce, key->key))
+	if (!chacha20poly1305_decrypt_sg(sg, sg, skb->len, NULL, 0, PACKET_CB(skb)->nonce, key->key, have_simd))
 		return false;
 
 	/* Another ugly situation of pushing and pulling the header so as to
@@ -411,12 +411,15 @@ void packet_decrypt_worker(struct work_struct *work)
 {
 	struct crypt_queue *queue = container_of(work, struct multicore_worker, work)->ptr;
 	struct sk_buff *skb;
+	bool have_simd = chacha20poly1305_init_simd();
 
 	while ((skb = ptr_ring_consume_bh(&queue->ring)) != NULL) {
-		enum packet_state state = likely(skb_decrypt(skb, &PACKET_CB(skb)->keypair->receiving)) ? PACKET_STATE_CRYPTED : PACKET_STATE_DEAD;
+		enum packet_state state = likely(skb_decrypt(skb, &PACKET_CB(skb)->keypair->receiving, have_simd)) ? PACKET_STATE_CRYPTED : PACKET_STATE_DEAD;
 
 		queue_enqueue_per_peer(&PACKET_PEER(skb)->rx_queue, skb, state);
 	}
+
+	chacha20poly1305_deinit_simd(have_simd);
 }
 
 static void packet_consume_data(struct wireguard_device *wg, struct sk_buff *skb)