@@ -445,8 +445,32 @@ mt76_dma_get_buf(struct mt76_dev *dev, struct mt76_queue *q, int idx,
if (mt76_queue_is_wed_rx(q)) {
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;
+ bool found = false;
+ u32 id;
+
+ 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)) {
+ token = id;
+ found = true;
+
+ /* 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 (!found)
+ return NULL;
+ }
+ t = mt76_rx_token_release(dev, token);
if (!t)
return NULL;