2 files changed, 32 insertions, 12 deletions

diff --git a/drivers/net/wireless/mediatek/mt76/mcu.c b/drivers/net/wireless/mediatek/mt76/mcu.c
index 2bbe41e3b78f..76dadef7ccfe 100644
--- a/drivers/net/wireless/mediatek/mt76/mcu.c
+++ b/drivers/net/wireless/mediatek/mt76/mcu.c
@@ -51,8 +51,8 @@ void mt76_mcu_rx_event(struct mt76_dev *dev, struct sk_buff *skb)
 }
 EXPORT_SYMBOL_GPL(mt76_mcu_rx_event);
 
-int mt76_mcu_send_msg(struct mt76_dev *dev, int cmd, const void *data,
-		      int len, bool wait_resp)
+int mt76_mcu_send_and_get_msg(struct mt76_dev *dev, int cmd, const void *data,
+			      int len, bool wait_resp, struct sk_buff **ret_skb)
 {
 	struct sk_buff *skb;
 
@@ -63,16 +63,20 @@ int mt76_mcu_send_msg(struct mt76_dev *dev, int cmd, const void *data,
 	if (!skb)
 		return -ENOMEM;
 
-	return mt76_mcu_skb_send_msg(dev, skb, cmd, wait_resp);
+	return mt76_mcu_skb_send_and_get_msg(dev, skb, cmd, wait_resp, ret_skb);
 }
-EXPORT_SYMBOL_GPL(mt76_mcu_send_msg);
+EXPORT_SYMBOL_GPL(mt76_mcu_send_and_get_msg);
 
-int mt76_mcu_skb_send_msg(struct mt76_dev *dev, struct sk_buff *skb,
-			  int cmd, bool wait_resp)
+int mt76_mcu_skb_send_and_get_msg(struct mt76_dev *dev, struct sk_buff *skb,
+				  int cmd, bool wait_resp,
+				  struct sk_buff **ret_skb)
 {
 	unsigned long expires;
 	int ret, seq;
 
+	if (ret_skb)
+		*ret_skb = NULL;
+
 	mutex_lock(&dev->mcu.mutex);
 
 	ret = dev->mcu_ops->mcu_skb_send_msg(dev, skb, cmd, &seq);
@@ -89,7 +93,10 @@ int mt76_mcu_skb_send_msg(struct mt76_dev *dev, struct sk_buff *skb,
 	do {
 		skb = mt76_mcu_get_response(dev, expires);
 		ret = dev->mcu_ops->mcu_parse_response(dev, cmd, skb, seq);
-		dev_kfree_skb(skb);
+		if (!ret && ret_skb)
+			*ret_skb = skb;
+		else
+			dev_kfree_skb(skb);
 	} while (ret == -EAGAIN);
 
 out:
@@ -97,4 +104,4 @@ out:
 
 	return ret;
 }
-EXPORT_SYMBOL_GPL(mt76_mcu_skb_send_msg);
+EXPORT_SYMBOL_GPL(mt76_mcu_skb_send_and_get_msg);
diff --git a/drivers/net/wireless/mediatek/mt76/mt76.h b/drivers/net/wireless/mediatek/mt76/mt76.h
index daa3d0b482a6..e1a3338c5723 100644
--- a/drivers/net/wireless/mediatek/mt76/mt76.h
+++ b/drivers/net/wireless/mediatek/mt76/mt76.h
@@ -1067,10 +1067,23 @@ mt76_mcu_msg_alloc(struct mt76_dev *dev, const void *data,
 void mt76_mcu_rx_event(struct mt76_dev *dev, struct sk_buff *skb);
 struct sk_buff *mt76_mcu_get_response(struct mt76_dev *dev,
 				      unsigned long expires);
-int mt76_mcu_send_msg(struct mt76_dev *dev, int cmd, const void *data,
-		      int len, bool wait_resp);
-int mt76_mcu_skb_send_msg(struct mt76_dev *dev, struct sk_buff *skb,
-			  int cmd, bool wait_resp);
+int mt76_mcu_send_and_get_msg(struct mt76_dev *dev, int cmd, const void *data,
+			      int len, bool wait_resp, struct sk_buff **ret);
+int mt76_mcu_skb_send_and_get_msg(struct mt76_dev *dev, struct sk_buff *skb,
+				  int cmd, bool wait_resp, struct sk_buff **ret);
+static inline int
+mt76_mcu_send_msg(struct mt76_dev *dev, int cmd, const void *data, int len,
+		  bool wait_resp)
+{
+	return mt76_mcu_send_and_get_msg(dev, cmd, data, len, wait_resp, NULL);
+}
+
+static inline int
+mt76_mcu_skb_send_msg(struct mt76_dev *dev, struct sk_buff *skb, int cmd,
+		      bool wait_resp)
+{
+	return mt76_mcu_skb_send_and_get_msg(dev, skb, cmd, wait_resp, NULL);
+}
 
 void mt76_set_irq_mask(struct mt76_dev *dev, u32 addr, u32 clear, u32 set);