From e336e49df3e63ff6ff0e101a4803bf6c0f9a3e02 Mon Sep 17 00:00:00 2001
From: "Jason A. Donenfeld" <Jason@zx2c4.com>
Date: Mon, 25 Jun 2018 06:05:14 +0200
Subject: send: wait for suspend to complete before sending handshake

Otherwise the transmission might be dropped on Android devices.
---
 src/device.c | 58 ++++++++++++++++++++++++++++++++++++----------------------
 src/device.h |  5 +++++
 src/send.c   |  6 ++++++
 3 files changed, 47 insertions(+), 22 deletions(-)

diff --git a/src/device.c b/src/device.c
index df9189b..5f5ef8b 100644
--- a/src/device.c
+++ b/src/device.c
@@ -67,31 +67,44 @@ static int open(struct net_device *dev)
 	return 0;
 }
 
-#if defined(CONFIG_PM_SLEEP) && !defined(CONFIG_ANDROID)
+#ifdef CONFIG_PM_SLEEP
+static DECLARE_COMPLETION(pm_awake);
+void device_wait_for_awake(void)
+{
+	wait_for_completion(&pm_awake);
+}
 static int pm_notification(struct notifier_block *nb, unsigned long action, void *data)
 {
-	struct wireguard_device *wg;
-	struct wireguard_peer *peer;
-
-	if (action != PM_HIBERNATION_PREPARE && action != PM_SUSPEND_PREPARE)
-		return 0;
-
-	rtnl_lock();
-	list_for_each_entry(wg, &device_list, device_list) {
-		mutex_lock(&wg->device_update_lock);
-		list_for_each_entry(peer, &wg->peer_list, peer_list) {
-			noise_handshake_clear(&peer->handshake);
-			noise_keypairs_clear(&peer->keypairs);
-			if (peer->timers_enabled)
-				del_timer(&peer->timer_zero_key_material);
+	switch (action) {
+	case PM_HIBERNATION_PREPARE:
+	case PM_SUSPEND_PREPARE: {
+#ifndef CONFIG_ANDROID
+		struct wireguard_device *wg;
+		struct wireguard_peer *peer;
+		rtnl_lock();
+		list_for_each_entry(wg, &device_list, device_list) {
+			mutex_lock(&wg->device_update_lock);
+			list_for_each_entry(peer, &wg->peer_list, peer_list) {
+				noise_handshake_clear(&peer->handshake);
+				noise_keypairs_clear(&peer->keypairs);
+				if (peer->timers_enabled)
+					del_timer(&peer->timer_zero_key_material);
+			}
+			mutex_unlock(&wg->device_update_lock);
 		}
-		mutex_unlock(&wg->device_update_lock);
+		rtnl_unlock();
+		rcu_barrier_bh();
+#endif
+		reinit_completion(&pm_awake);
 	}
-	rtnl_unlock();
-	rcu_barrier_bh();
+	case PM_POST_HIBERNATION:
+	case PM_POST_SUSPEND:
+		complete_all(&pm_awake);
+	}
+
 	return 0;
 }
-static struct notifier_block pm_notifier = { .notifier_call = pm_notification };
+static struct notifier_block pm_notifier = { .notifier_call = pm_notification, .priority = INT_MIN };
 #endif
 
 static int stop(struct net_device *dev)
@@ -390,7 +403,8 @@ int __init device_init(void)
 {
 	int ret;
 
-#if defined(CONFIG_PM_SLEEP) && !defined(CONFIG_ANDROID)
+#ifdef CONFIG_PM_SLEEP
+	complete_all(&pm_awake);
 	ret = register_pm_notifier(&pm_notifier);
 	if (ret)
 		return ret;
@@ -409,7 +423,7 @@ int __init device_init(void)
 error_netdevice:
 	unregister_netdevice_notifier(&netdevice_notifier);
 error_pm:
-#if defined(CONFIG_PM_SLEEP) && !defined(CONFIG_ANDROID)
+#ifdef CONFIG_PM_SLEEP
 	unregister_pm_notifier(&pm_notifier);
 #endif
 	return ret;
@@ -419,7 +433,7 @@ void device_uninit(void)
 {
 	rtnl_link_unregister(&link_ops);
 	unregister_netdevice_notifier(&netdevice_notifier);
-#if defined(CONFIG_PM_SLEEP) && !defined(CONFIG_ANDROID)
+#ifdef CONFIG_PM_SLEEP
 	unregister_pm_notifier(&pm_notifier);
 #endif
 	rcu_barrier_bh();
diff --git a/src/device.h b/src/device.h
index 2a0e2c7..da8a72a 100644
--- a/src/device.h
+++ b/src/device.h
@@ -60,5 +60,10 @@ struct wireguard_device {
 
 int device_init(void);
 void device_uninit(void);
+#if defined(CONFIG_PM_SLEEP)
+void device_wait_for_awake(void);
+#else
+static inline void device_wait_for_awake(void) { }
+#endif
 
 #endif /* _WG_DEVICE_H */
diff --git a/src/send.c b/src/send.c
index 7a8bea1..0b3ff48 100644
--- a/src/send.c
+++ b/src/send.c
@@ -31,6 +31,8 @@ static void packet_send_handshake_initiation(struct wireguard_peer *peer)
 	peer->last_sent_handshake = ktime_get_boot_fast_ns();
 	up_write(&peer->handshake.lock);
 
+	device_wait_for_awake();
+
 	net_dbg_ratelimited("%s: Sending handshake initiation to peer %llu (%pISpfsc)\n", peer->device->dev->name, peer->internal_id, &peer->endpoint.addr);
 
 	if (noise_handshake_create_initiation(&packet, &peer->handshake)) {
@@ -71,6 +73,8 @@ void packet_send_handshake_response(struct wireguard_peer *peer)
 {
 	struct message_handshake_response packet;
 
+	device_wait_for_awake();
+
 	net_dbg_ratelimited("%s: Sending handshake response to peer %llu (%pISpfsc)\n", peer->device->dev->name, peer->internal_id, &peer->endpoint.addr);
 	peer->last_sent_handshake = ktime_get_boot_fast_ns();
 
@@ -89,6 +93,8 @@ void packet_send_handshake_cookie(struct wireguard_device *wg, struct sk_buff *i
 {
 	struct message_handshake_cookie packet;
 
+	device_wait_for_awake();
+
 	net_dbg_skb_ratelimited("%s: Sending cookie response for denied handshake message for %pISpfsc\n", wg->dev->name, initiating_skb);
 	cookie_message_create(&packet, initiating_skb, sender_index, &wg->cookie_checker);
 	socket_send_buffer_as_reply_to_skb(wg, initiating_skb, &packet, sizeof(packet));
-- 
cgit v1.2.3-59-g8ed1b