1 files changed, 9 insertions, 18 deletions

diff --git a/drivers/net/wireless/iwlwifi/iwl-testmode.c b/drivers/net/wireless/iwlwifi/iwl-testmode.c
index b06c6763cb7..76f7f925143 100644
--- a/drivers/net/wireless/iwlwifi/iwl-testmode.c
+++ b/drivers/net/wireless/iwlwifi/iwl-testmode.c
@@ -466,6 +466,7 @@ static int iwl_testmode_driver(struct ieee80211_hw *hw, struct nlattr **tb)
 	unsigned char *rsp_data_ptr = NULL;
 	int status = 0, rsp_data_len = 0;
 	u32 devid, inst_size = 0, data_size = 0;
+	const struct fw_img *img;
 
 	switch (nla_get_u32(tb[IWL_TM_ATTR_COMMAND])) {
 	case IWL_TM_CMD_APP2DEV_GET_DEVICENAME:
@@ -494,6 +495,7 @@ static int iwl_testmode_driver(struct ieee80211_hw *hw, struct nlattr **tb)
 
 	case IWL_TM_CMD_APP2DEV_CFG_INIT_CALIB:
 		iwl_testmode_cfg_init_calib(priv);
+		priv->ucode_loaded = false;
 		iwl_trans_stop_device(trans);
 		break;
 
@@ -512,6 +514,7 @@ static int iwl_testmode_driver(struct ieee80211_hw *hw, struct nlattr **tb)
 
 	case IWL_TM_CMD_APP2DEV_LOAD_WOWLAN_FW:
 		iwl_scan_cancel_timeout(priv, 200);
+		priv->ucode_loaded = false;
 		iwl_trans_stop_device(trans);
 		status = iwl_load_ucode_wait_alive(priv, IWL_UCODE_WOWLAN);
 		if (status) {
@@ -591,25 +594,13 @@ static int iwl_testmode_driver(struct ieee80211_hw *hw, struct nlattr **tb)
 			IWL_ERR(priv, "Memory allocation fail\n");
 			return -ENOMEM;
 		}
-		switch (priv->shrd->ucode_type) {
-		case IWL_UCODE_REGULAR:
-			inst_size = priv->fw->ucode_rt.code.len;
-			data_size = priv->fw->ucode_rt.data.len;
-			break;
-		case IWL_UCODE_INIT:
-			inst_size = priv->fw->ucode_init.code.len;
-			data_size = priv->fw->ucode_init.data.len;
-			break;
-		case IWL_UCODE_WOWLAN:
-			inst_size = priv->fw->ucode_wowlan.code.len;
-			data_size = priv->fw->ucode_wowlan.data.len;
-			break;
-		case IWL_UCODE_NONE:
+		if (!priv->ucode_loaded) {
 			IWL_ERR(priv, "No uCode has not been loaded\n");
-			break;
-		default:
-			IWL_ERR(priv, "Unsupported uCode type\n");
-			break;
+			return -EINVAL;
+		} else {
+			img = &priv->fw->img[priv->shrd->ucode_type];
+			inst_size = img->sec[IWL_UCODE_SECTION_INST].len;
+			data_size = img->sec[IWL_UCODE_SECTION_DATA].len;
 		}
 		NLA_PUT_U32(skb, IWL_TM_ATTR_FW_TYPE, priv->shrd->ucode_type);
 		NLA_PUT_U32(skb, IWL_TM_ATTR_FW_INST_SIZE, inst_size);