@@ -444,9 +444,32 @@ mt76_dma_get_buf(struct mt76_dev *dev, struct mt76_queue *q, int idx,
mt76_dma_should_drop_buf(drop, ctrl, buf1, desc_info);
if (mt76_queue_is_wed_rx(q)) {
+ u32 id, find = 0;
u32 token = FIELD_GET(MT_DMA_CTL_TOKEN, buf1);
- struct mt76_txwi_cache *t = mt76_rx_token_release(dev, token);
+ struct mt76_txwi_cache *t;
+
+ if (*more) {
+ spin_lock_bh(&dev->rx_token_lock);
+
+ idr_for_each_entry(&dev->rx_token, t, id) {
+ if (t->dma_addr == le32_to_cpu(desc->buf0)) {
+ find = 1;
+ token = id;
+
+ /* Write correct id back to DMA*/
+ u32p_replace_bits(&buf1, id,
+ MT_DMA_CTL_TOKEN);
+ WRITE_ONCE(desc->buf1, cpu_to_le32(buf1));
+ break;
+ }
+ }
+
+ spin_unlock_bh(&dev->rx_token_lock);
+ if (!find)
+ return NULL;
+ }
+ t = mt76_rx_token_release(dev, token);
if (!t)
return NULL;