From e7fac7c5c91fbabafa530abdf32cf11dd36834d7 Mon Sep 17 00:00:00 2001
From: Shayne Chen <shayne.chen@mediatek.com>
Date: Fri, 25 Feb 2022 09:36:01 +0800
Subject: [PATCH 2/7] mt76: testmode: rework tx antenna setting

---
 mt7915/mcu.c      |  7 +------
 mt7915/testmode.c | 13 ++++++-------
 testmode.c        |  3 +--
 3 files changed, 8 insertions(+), 15 deletions(-)

diff --git a/mt7915/mcu.c b/mt7915/mcu.c
index 3586a8e3..d9c9805b 100644
--- a/mt7915/mcu.c
+++ b/mt7915/mcu.c
@@ -2642,14 +2642,9 @@ int mt7915_mcu_set_chan_info(struct mt7915_phy *phy, int cmd)
 
 #ifdef CONFIG_NL80211_TESTMODE
 	if (phy->mt76->test.tx_antenna_mask &&
-	    (phy->mt76->test.state == MT76_TM_STATE_TX_FRAMES ||
-	     phy->mt76->test.state == MT76_TM_STATE_RX_FRAMES ||
-	     phy->mt76->test.state == MT76_TM_STATE_TX_CONT)) {
+	    mt76_testmode_enabled(phy->mt76)) {
 		req.tx_streams_num = fls(phy->mt76->test.tx_antenna_mask);
 		req.rx_streams = phy->mt76->test.tx_antenna_mask;
-
-		if (phy != &dev->phy)
-			req.rx_streams >>= dev->chainshift;
 	}
 #endif
 
diff --git a/mt7915/testmode.c b/mt7915/testmode.c
index 01866028..41b99496 100644
--- a/mt7915/testmode.c
+++ b/mt7915/testmode.c
@@ -474,11 +474,7 @@ mt7915_tm_set_tx_frames(struct mt7915_phy *phy, bool en)
 		if (td->tx_spe_idx) {
 			phy->test.spe_idx = td->tx_spe_idx;
 		} else {
-			u8 tx_ant = td->tx_antenna_mask;
-
-			if (phy != &dev->phy)
-				tx_ant >>= dev->chainshift;
-			phy->test.spe_idx = spe_idx_map[tx_ant];
+			phy->test.spe_idx = spe_idx_map[td->tx_antenna_mask];
 		}
 	}
 
@@ -720,8 +716,10 @@ mt7915_tm_set_params(struct mt76_phy *mphy, struct nlattr **tb,
 {
 	struct mt76_testmode_data *td = &mphy->test;
 	struct mt7915_phy *phy = mphy->priv;
-	u32 changed = 0;
+	struct mt7915_dev *dev = phy->dev;
+	u32 chainmask = mphy->chainmask, changed = 0;
 	int i;
+	bool ext_phy = phy != &dev->phy;
 
 	BUILD_BUG_ON(NUM_TM_CHANGED >= 32);
 
@@ -729,7 +727,8 @@ mt7915_tm_set_params(struct mt76_phy *mphy, struct nlattr **tb,
 	    td->state == MT76_TM_STATE_OFF)
 		return 0;
 
-	if (td->tx_antenna_mask & ~mphy->chainmask)
+	chainmask = ext_phy ? chainmask >> dev->chainshift : chainmask;
+	if (td->tx_antenna_mask > chainmask)
 		return -EINVAL;
 
 	for (i = 0; i < ARRAY_SIZE(tm_change_map); i++) {
diff --git a/testmode.c b/testmode.c
index 0accc71a..57cdfdf6 100644
--- a/testmode.c
+++ b/testmode.c
@@ -447,8 +447,7 @@ int mt76_testmode_cmd(struct ieee80211_hw *hw, struct ieee80211_vif *vif,
 	    mt76_tm_get_u8(tb[MT76_TM_ATTR_TX_RATE_LDPC], &td->tx_rate_ldpc, 0, 1) ||
 	    mt76_tm_get_u8(tb[MT76_TM_ATTR_TX_RATE_STBC], &td->tx_rate_stbc, 0, 1) ||
 	    mt76_tm_get_u8(tb[MT76_TM_ATTR_TX_LTF], &td->tx_ltf, 0, 2) ||
-	    mt76_tm_get_u8(tb[MT76_TM_ATTR_TX_ANTENNA],
-			   &td->tx_antenna_mask, 0, 0xff) ||
+	    mt76_tm_get_u8(tb[MT76_TM_ATTR_TX_ANTENNA], &td->tx_antenna_mask, 1, 0xff) ||
 	    mt76_tm_get_u8(tb[MT76_TM_ATTR_TX_SPE_IDX], &td->tx_spe_idx, 0, 27) ||
 	    mt76_tm_get_u8(tb[MT76_TM_ATTR_TX_DUTY_CYCLE],
 			   &td->tx_duty_cycle, 0, 99) ||
-- 
2.25.1