>From c5613e859327a4cff8ac7a96f8f3550016efccc2 Mon Sep 17 00:00:00 2001
From: Ilan Peer <ilan.peer@intel.com>
Date: Wed, 16 Dec 2020 13:01:01 +0200
Subject: [PATCH 4/5] PASN: Add support for comeback flow to the wpa_supplicant

Signed-off-by: Ilan Peer <ilan.peer@intel.com>
---
 wpa_supplicant/ctrl_iface.c       | 26 +++++++++-
 wpa_supplicant/pasn_supplicant.c  | 84 +++++++++++++++++++++++++++----
 wpa_supplicant/wpa_supplicant_i.h |  6 ++-
 3 files changed, 103 insertions(+), 13 deletions(-)

diff --git a/wpa_supplicant/ctrl_iface.c b/wpa_supplicant/ctrl_iface.c
index da198b938ba3..80c1df8de205 100644
--- a/wpa_supplicant/ctrl_iface.c
+++ b/wpa_supplicant/ctrl_iface.c
@@ -10464,7 +10464,9 @@ static int wpas_ctrl_iface_pasn_start(struct wpa_supplicant *wpa_s, char *cmd)
 	u8 bssid[ETH_ALEN];
 	int akmp = -1, cipher = -1, got_bssid = 0;
 	u16 group = 0xFFFF;
-	int id = 0;
+	u8 *comeback = NULL;
+	size_t comeback_len = 0;
+	int id = 0, ret;
 
 	/*
 	 * Entry format: bssid=<BSSID> akmp=<AKMP> cipher=<CIPHER> group=<group>
@@ -10506,6 +10508,21 @@ static int wpas_ctrl_iface_pasn_start(struct wpa_supplicant *wpa_s, char *cmd)
 			group = atoi(token + 6);
 		} else if (os_strncmp(token, "nid=", 4) == 0) {
 			id = atoi(token + 4);
+		} else if (os_strncmp(token, "comeback=", 9) == 0) {
+			comeback_len = os_strlen(token + 9);
+
+			if (!comeback_len || comeback_len % 2)
+				return -1;
+
+			comeback_len /= 2;
+			comeback = os_malloc(comeback_len);
+			if (!comeback)
+				return -1;
+
+			if (hexstr2bin(token + 9, comeback, comeback_len)) {
+				os_free(comeback);
+				return -1;
+			}
 		} else {
 			wpa_printf(MSG_DEBUG,
 				   "CTRL: PASN Invalid parameter: '%s'",
@@ -10516,10 +10533,15 @@ static int wpas_ctrl_iface_pasn_start(struct wpa_supplicant *wpa_s, char *cmd)
 
 	if (!got_bssid || akmp == -1 || cipher == -1 || group == 0xFFFF) {
 		wpa_printf(MSG_DEBUG,"CTRL: PASN missing parameter");
+		os_free(comeback);
 		return -1;
 	}
 
-	return wpas_pasn_auth_start(wpa_s, bssid, akmp, cipher, group, id);
+	ret = wpas_pasn_auth_start(wpa_s, bssid, akmp, cipher, group, id,
+				   comeback, comeback_len);
+
+	os_free(comeback);
+	return ret;
 }
 #endif /* CONFIG_PASN */
 
diff --git a/wpa_supplicant/pasn_supplicant.c b/wpa_supplicant/pasn_supplicant.c
index 736a6c7d6f96..642bcea3c5ca 100644
--- a/wpa_supplicant/pasn_supplicant.c
+++ b/wpa_supplicant/pasn_supplicant.c
@@ -33,6 +33,7 @@ struct wpa_pasn_auth_work {
 	int cipher;
 	u16 group;
 	int network_id;
+	struct wpabuf *comeback;
 };
 
 
@@ -56,8 +57,30 @@ static void wpas_pasn_cancel_auth_work(struct wpa_supplicant *wpa_s)
 
 
 static void wpas_pasn_auth_status(struct wpa_supplicant *wpa_s, const u8 *bssid,
-				  int akmp, int cipher, u8 status)
+				  int akmp, int cipher, u8 status,
+				  struct wpabuf *comeback,
+				  u16 comeback_after)
 {
+	if (comeback) {
+		size_t comeback_len = wpabuf_len(comeback);
+		char *comeback_txt = os_malloc(comeback_len * 2 + 1);
+
+		if (comeback_txt) {
+			wpa_snprintf_hex(comeback_txt, comeback_len * 2 + 1,
+					 wpabuf_head(comeback), comeback_len);
+
+			wpa_msg(wpa_s, MSG_INFO, PASN_AUTH_STATUS MACSTR
+				" akmp=%s, status=%u comeback_after=%u"
+				" comeback=%s",
+				MAC2STR(bssid),
+				wpa_key_mgmt_txt(akmp, WPA_PROTO_RSN),
+				status, comeback_after, comeback_txt);
+
+			os_free(comeback_txt);
+			return;
+		}
+	}
+
 	wpa_msg(wpa_s, MSG_INFO,
 		PASN_AUTH_STATUS MACSTR " akmp=%s, status=%u",
 		MAC2STR(bssid), wpa_key_mgmt_txt(akmp, WPA_PROTO_RSN),
@@ -612,7 +635,8 @@ static u8 wpas_pasn_get_wrapped_data_format(struct wpas_pasn *pasn)
 }
 
 
-static struct wpabuf * wpas_pasn_build_auth_1(struct wpa_supplicant *wpa_s)
+static struct wpabuf * wpas_pasn_build_auth_1(struct wpa_supplicant *wpa_s,
+					      struct wpabuf *comeback)
 {
 	struct wpas_pasn *pasn = &wpa_s->pasn;
 	struct wpabuf *buf, *pubkey = NULL, *wrapped_data_buf = NULL;
@@ -680,7 +704,7 @@ static struct wpabuf * wpas_pasn_build_auth_1(struct wpa_supplicant *wpa_s)
 		wrapped_data = WPA_PASN_WRAPPED_DATA_NO;
 
 	wpa_pasn_add_parameter_ie(buf, pasn->group, wrapped_data,
-				  pubkey, NULL, -1);
+				  pubkey, comeback, -1);
 
 	if (wpa_pasn_add_wrapped_data(buf, wrapped_data_buf) < 0)
 		goto fail;
@@ -821,6 +845,10 @@ static void wpas_pasn_reset(struct wpa_supplicant *wpa_s)
 	wpabuf_free(pasn->beacon_rsne_rsnxe);
 	pasn->beacon_rsne_rsnxe = NULL;
 
+	wpabuf_free(pasn->comeback);
+	pasn->comeback = NULL;
+	pasn->comeback_after = 0;
+
 #ifdef CONFIG_SAE
 	sae_clear_data(&pasn->sae);
 #endif /* CONFIG_SAE */
@@ -939,7 +967,7 @@ static int wpas_pasn_start(struct wpa_supplicant *wpa_s, const u8 *bssid,
 			   int akmp, int cipher, u16 group, int freq,
 			   const u8 *beacon_rsne, u8 beacon_rsne_len,
 			   const u8 *beacon_rsnxe, u8 beacon_rsnxe_len,
-			   int network_id)
+			   int network_id, struct wpabuf *comeback)
 {
 	struct wpas_pasn *pasn = &wpa_s->pasn;
 	struct wpa_ssid *ssid = NULL;
@@ -1016,7 +1044,7 @@ static int wpas_pasn_start(struct wpa_supplicant *wpa_s, const u8 *bssid,
 		   MAC2STR(pasn->bssid), pasn->akmp, pasn->cipher,
 		   pasn->group);
 
-	frame = wpas_pasn_build_auth_1(wpa_s);
+	frame = wpas_pasn_build_auth_1(wpa_s, comeback);
 	if (!frame) {
 		wpa_printf(MSG_DEBUG, "PASN: Failed building 1st auth frame");
 		goto fail;
@@ -1098,6 +1126,8 @@ static void wpas_pasn_auth_start_cb(struct wpa_radio_work *work, int deinit)
 					     wpa_s, NULL);
 			wpa_s->pasn_auth_work = NULL;
 		}
+
+		wpabuf_free(awork->comeback);
 		os_free(awork);
 		return;
 	}
@@ -1125,16 +1155,22 @@ static void wpas_pasn_auth_start_cb(struct wpa_radio_work *work, int deinit)
 	ret = wpas_pasn_start(wpa_s, awork->bssid, awork->akmp, awork->cipher,
 			      awork->group, bss->freq, rsne, *(rsne + 1) + 2,
 			      rsnxe, rsnxe ? *(rsnxe + 1) + 2 : 0,
-			      awork->network_id);
+			      awork->network_id, awork->comeback);
 	if (ret) {
 		wpa_printf(MSG_DEBUG,
 			   "PASN: Failed to start PASN authentication");
 		goto fail;
 	}
 
+	/* comeback token is no longer needed at this stage */
+	wpabuf_free(awork->comeback);
+	awork->comeback = NULL;
+
 	wpa_s->pasn_auth_work = work;
 	return;
 fail:
+	wpabuf_free(awork->comeback);
+	awork->comeback = NULL;
 	os_free(awork);
 	work->ctx = NULL;
 	radio_work_done(work);
@@ -1142,7 +1178,8 @@ fail:
 
 
 int wpas_pasn_auth_start(struct wpa_supplicant *wpa_s, const u8 *bssid,
-			 int akmp, int cipher, u16 group, int network_id)
+			 int akmp, int cipher, u16 group, int network_id,
+			 const u8 *comeback, size_t comeback_len)
 {
 	struct wpa_pasn_auth_work *awork;
 	struct wpa_bss *bss;
@@ -1188,8 +1225,19 @@ int wpas_pasn_auth_start(struct wpa_supplicant *wpa_s, const u8 *bssid,
 	awork->group = group;
 	awork->network_id = network_id;
 
+	if (comeback && comeback_len) {
+		awork->comeback = wpabuf_alloc(comeback_len);
+		if (!awork->comeback) {
+			os_free(awork);
+			return -1;
+		}
+
+		wpabuf_put_data(awork->comeback, comeback, comeback_len);
+	}
+
 	if (radio_add_work(wpa_s, bss->freq, "pasn-start-auth", 1,
 			   wpas_pasn_auth_start_cb, awork) < 0) {
+		wpabuf_free(awork->comeback);
 		os_free(awork);
 		return -1;
 	}
@@ -1209,7 +1257,8 @@ void wpas_pasn_auth_stop(struct wpa_supplicant *wpa_s)
 	wpa_printf(MSG_DEBUG, "PASN: Stopping authentication");
 
 	wpas_pasn_auth_status(wpa_s, pasn->bssid, pasn->akmp, pasn->cipher,
-			      pasn->status);
+			      pasn->status, pasn->comeback,
+			      pasn->comeback_after);
 
 	wpas_pasn_reset(wpa_s);
 }
@@ -1308,10 +1357,21 @@ int wpas_pasn_auth_rx(struct wpa_supplicant *wpa_s,
 		goto fail;
 	}
 
-	/* TODO: handle comeback flow */
 	if (status == WLAN_STATUS_ASSOC_REJECTED_TEMPORARILY) {
 		wpa_printf(MSG_DEBUG,
 			   "PASN: Authentication temporarily rejected");
+
+		if (pasn_params.comeback && pasn_params.comeback_len) {
+			pasn->comeback = wpabuf_alloc(pasn_params.comeback_len);
+			if (pasn->comeback) {
+				wpabuf_put_data(pasn->comeback,
+						pasn_params.comeback,
+						pasn_params.comeback_len);
+				pasn->comeback_after = pasn_params.after;
+			}
+		}
+
+		pasn->status = status;
 		goto fail;
 	}
 
@@ -1437,7 +1497,11 @@ fail:
 	 * the frame and terminate the authentication exchange. However, better
 	 * reply to the AP with an error status.
 	 */
-	pasn->status = WLAN_STATUS_UNSPECIFIED_FAILURE;
+	if (status == WLAN_STATUS_SUCCESS)
+		pasn->status = WLAN_STATUS_UNSPECIFIED_FAILURE;
+	else
+		pasn->status = status;
+
 	wpas_pasn_auth_stop(wpa_s);
 	return -1;
 }
diff --git a/wpa_supplicant/wpa_supplicant_i.h b/wpa_supplicant/wpa_supplicant_i.h
index 0ce7447fe728..bea940642f00 100644
--- a/wpa_supplicant/wpa_supplicant_i.h
+++ b/wpa_supplicant/wpa_supplicant_i.h
@@ -546,6 +546,9 @@ struct wpas_pasn {
 	struct wpa_ptk ptk;
 	struct crypto_ecdh *ecdh;
 
+	struct wpabuf *comeback;
+	u16 comeback_after;
+
 #ifdef CONFIG_SAE
 	struct sae_data sae;
 #endif /* CONFIG_SAE */
@@ -1712,7 +1715,8 @@ void wpas_handle_assoc_resp_mscs(struct wpa_supplicant *wpa_s, const u8 *bssid,
 
 int wpas_pasn_auth_start(struct wpa_supplicant *wpa_s,
 			 const u8 *bssid, int akmp, int cipher,
-			 u16 group, int network_id);
+			 u16 group, int network_id,
+			 const u8 *comeback, size_t comeback_len);
 void wpas_pasn_auth_stop(struct wpa_supplicant *wpa_s);
 int wpas_pasn_auth_tx_status(struct wpa_supplicant *wpa_s,
 			     const u8 *data, size_t data_len, u8 acked);
-- 
2.25.1