timers: clear send_keepalive timer on sending handshake response

We reorganize this into also doing so on sending keepalives itself,
which means the state machine is much more consistent, even if this was
already implied.

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

3 files changed, 12 insertions, 6 deletions

diff --git a/src/send.c b/src/send.c
index 9b1aec0..6e04ad4 100644
--- a/src/send.c
+++ b/src/send.c
@@ -36,6 +36,7 @@ static void packet_send_handshake_initiation(struct wireguard_peer *peer)
 	if (noise_handshake_create_initiation(&packet, &peer->handshake)) {
 		cookie_add_mac_to_packet(&packet, sizeof(packet), peer);
 		timers_any_authenticated_packet_traversal(peer);
+		timers_any_authenticated_packet_sent(peer);
 		socket_send_buffer_to_peer(peer, &packet, sizeof(struct message_handshake_initiation), HANDSHAKE_DSCP);
 		timers_handshake_initiated(peer);
 	}
@@ -78,6 +79,7 @@ void packet_send_handshake_response(struct wireguard_peer *peer)
 		if (noise_handshake_begin_session(&peer->handshake, &peer->keypairs)) {
 			timers_session_derived(peer);
 			timers_any_authenticated_packet_traversal(peer);
+			timers_any_authenticated_packet_sent(peer);
 			socket_send_buffer_to_peer(peer, &packet, sizeof(struct message_handshake_response), HANDSHAKE_DSCP);
 		}
 	}
@@ -200,6 +202,7 @@ static void packet_create_data_done(struct sk_buff *first, struct wireguard_peer
 	bool is_keepalive, data_sent = false;
 
 	timers_any_authenticated_packet_traversal(peer);
+	timers_any_authenticated_packet_sent(peer);
 	skb_walk_null_queue_safe(first, skb, next) {
 		is_keepalive = skb->len == message_data_len(0);
 		if (likely(!socket_send_skb_to_peer(peer, skb, PACKET_CB(skb)->ds) && !is_keepalive))
diff --git a/src/timers.c b/src/timers.c
index 0e14098..bd22bf6 100644
--- a/src/timers.c
+++ b/src/timers.c
@@ -114,9 +114,6 @@ 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)))
-		del_timer(&peer->timer_send_keepalive);
-
 	if (likely(timers_active(peer)) && !timer_pending(&peer->timer_new_handshake))
 		mod_timer(&peer->timer_new_handshake, jiffies + KEEPALIVE_TIMEOUT + REKEY_TIMEOUT);
 }
@@ -132,6 +129,13 @@ 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);
+}
+
 /* 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)
 {
@@ -142,10 +146,8 @@ void timers_any_authenticated_packet_received(struct wireguard_peer *peer)
 /* Should be called after a handshake initiation message is sent. */
 void timers_handshake_initiated(struct wireguard_peer *peer)
 {
-	if (likely(timers_active(peer))) {
-		del_timer(&peer->timer_send_keepalive);
+	if (likely(timers_active(peer)))
 		mod_timer(&peer->timer_retransmit_handshake, jiffies + REKEY_TIMEOUT + prandom_u32_max(REKEY_TIMEOUT_JITTER_MAX));
-	}
 }
 
 /* Should be called after a handshake response message is received and processed or when getting key confirmation via the first data message. */
diff --git a/src/timers.h b/src/timers.h
index fd2fa79..6c67a6a 100644
--- a/src/timers.h
+++ b/src/timers.h
@@ -12,6 +12,7 @@ void timers_init(struct wireguard_peer *peer);
 void timers_stop(struct wireguard_peer *peer);
 void timers_data_sent(struct wireguard_peer *peer);
 void timers_data_received(struct wireguard_peer *peer);
+void timers_any_authenticated_packet_sent(struct wireguard_peer *peer);
 void timers_any_authenticated_packet_received(struct wireguard_peer *peer);
 void timers_handshake_initiated(struct wireguard_peer *peer);
 void timers_handshake_complete(struct wireguard_peer *peer);