Fix sleep/resume and improve reliability

- Wait for T2 PCIe link re-training after S3 before MMIO access
- Fix wake mailbox command (0x18 -> 0x1b to match macOS)
- Reorder endpoint pause: drain pending output DMA before pausing
- Handle T2 port connect/resume/suspend event notifications
- Add command queue timeout (5s) to prevent infinite hangs
- Add device link to NVMe function 0 for runtime PM ordering
- Fix atomic counter underflow in transfer queue completions
- Add NULL checks in URB enqueue/dequeue to prevent crashes

Author: klizas

apple_bce.c

@@ -103,6 +103,14 @@ static int apple_bce_probe(struct pci_dev *dev, const struct pci_device_id *id)
 
     bce_vhci_create(bce, &bce->vhci);
 
+    /* The T2 chip requires function 0 (NVMe) to be a bus master for DMA
+     * on our function. Create a device link for runtime PM ordering.
+     * (System S3 ordering is already handled by PCI function numbering.) */
+    bce->pci0_link = device_link_add(&dev->dev, &bce->pci0->dev,
+                                     DL_FLAG_STATELESS | DL_FLAG_PM_RUNTIME);
+    if (!bce->pci0_link)
+        dev_warn(&dev->dev, "apple-bce: failed to create device link to function 0\n");
+
     return 0;
 
 fail_ts:
@@ -243,6 +251,9 @@ static void apple_bce_remove(struct pci_dev *dev)
 
     bce_vhci_destroy(&bce->vhci);
 
+    if (bce->pci0_link)
+        device_link_del(bce->pci0_link);
+
     bce_timestamp_stop(&bce->timestamp);
 #ifndef WITHOUT_NVME_PATCH
     pci_disable_device(bce->pci0);
@@ -356,6 +367,24 @@ static int apple_bce_resume(struct device *dev)
 {
     struct apple_bce_device *bce = pci_get_drvdata(to_pci_dev(dev));
     int status;
+    int i;
+    u16 vid;
+
+    /* Wait for T2 PCIe link to re-train after S3.
+     * MMIO to the T2 BARs will hang the CPU if the link is down.
+     * Config space reads go through the root port and return 0xFFFF safely.
+     * Poll aggressively first (link usually retrains in ~100-200ms),
+     * then back off to 50ms intervals. */
+    for (i = 0; i < 120; i++) {
+        pci_read_config_word(bce->pci, PCI_VENDOR_ID, &vid);
+        if (vid == PCI_VENDOR_ID_APPLE)
+            break;
+        msleep(i < 40 ? 5 : 50);
+    }
+    if (vid != PCI_VENDOR_ID_APPLE) {
+        pr_err("apple-bce: resume: T2 not accessible after timeout (vid=0x%04x)\n", vid);
+        return -ENODEV;
+    }
 
     pci_set_master(bce->pci);
     pci_set_master(bce->pci0);

apple_bce.h

@@ -15,6 +15,7 @@
 
 struct apple_bce_device {
     struct pci_dev *pci, *pci0;
+    struct device_link *pci0_link;
     dev_t devt;
     struct device *dev;
     void __iomem *reg_mem_mb;

mailbox.h

@@ -21,7 +21,7 @@ enum bce_message_type {
     BCE_MB_SLEEP_NO_STATE = 0x14,                // to-device
     BCE_MB_RESTORE_NO_STATE = 0x15,              // to-device
     BCE_MB_SAVE_STATE_AND_SLEEP = 0x17,          // to-device
-    BCE_MB_RESTORE_STATE_AND_WAKE = 0x18,        // to-device
+    BCE_MB_RESTORE_STATE_AND_WAKE = 0x1b,        // to-device (macOS: 0x6c00000000000000)
     BCE_MB_SAVE_STATE_AND_SLEEP_FAILURE = 0x19,  // from-device
     BCE_MB_SAVE_RESTORE_STATE_COMPLETE = 0x1A,   // from-device
 };

queue.c

@@ -250,18 +250,28 @@ static __always_inline void *bce_cmd_start(struct bce_queue_cmdq *cmdq, struct b
         return NULL;
 
     spin_lock(&cmdq->lck);
+    res->slot = cmdq->sq->tail;
     cmdq->tres[cmdq->sq->tail] = res;
     ret = bce_next_submission(cmdq->sq);
     return ret;
 }
 
-static __always_inline void bce_cmd_finish(struct bce_queue_cmdq *cmdq, struct bce_queue_cmdq_result_el *res)
+static __always_inline int bce_cmd_finish(struct bce_queue_cmdq *cmdq, struct bce_queue_cmdq_result_el *res)
 {
     bce_submit_to_device(cmdq->sq);
     spin_unlock(&cmdq->lck);
 
-    wait_for_completion(&res->cmpl);
+    if (!wait_for_completion_timeout(&res->cmpl, msecs_to_jiffies(5000))) {
+        pr_err("apple-bce: command queue timeout\n");
+        spin_lock(&cmdq->lck);
+        cmdq->tres[res->slot] = NULL;
+        spin_unlock(&cmdq->lck);
+        /* Reclaim the slot: advance head and wake any waiters */
+        bce_notify_submission_complete(cmdq->sq);
+        return -ETIMEDOUT;
+    }
     mb();
+    return 0;
 }
 
 u32 bce_cmd_register_queue(struct bce_queue_cmdq *cmdq, struct bce_queue_memcfg *cfg, const char *name, bool isdirout)
@@ -285,7 +295,8 @@ u32 bce_cmd_register_queue(struct bce_queue_cmdq *cmdq, struct bce_queue_memcfg
     cmd->addr = cfg->addr;
     cmd->length = cfg->length;
 
-    bce_cmd_finish(cmdq, &res);
+    if (bce_cmd_finish(cmdq, &res))
+        return (u32) -1;
     return res.status;
 }
 
@@ -298,7 +309,8 @@ u32 bce_cmd_unregister_memory_queue(struct bce_queue_cmdq *cmdq, u16 qid)
     cmd->cmd = BCE_CMD_UNREGISTER_MEMORY_QUEUE;
     cmd->flags = 0;
     cmd->qid = qid;
-    bce_cmd_finish(cmdq, &res);
+    if (bce_cmd_finish(cmdq, &res))
+        return (u32) -1;
     return res.status;
 }
 
@@ -311,7 +323,8 @@ u32 bce_cmd_flush_memory_queue(struct bce_queue_cmdq *cmdq, u16 qid)
     cmd->cmd = BCE_CMD_FLUSH_MEMORY_QUEUE;
     cmd->flags = 0;
     cmd->qid = qid;
-    bce_cmd_finish(cmdq, &res);
+    if (bce_cmd_finish(cmdq, &res))
+        return (u32) -1;
     return res.status;
 }
 

queue.h

@@ -56,6 +56,7 @@ struct bce_queue_cmdq_result_el {
     struct completion cmpl;
     u32 status;
     u64 result;
+    u32 slot;  /* queue slot index for O(1) timeout cleanup */
 };
 struct bce_queue_cmdq {
     struct bce_queue_sq *sq;

vhci/command.h

@@ -27,6 +27,7 @@ enum bce_vhci_command {
     BCE_VHCI_CMD_PORT_RESET = 0x14,
     BCE_VHCI_CMD_PORT_DISABLE = 0x15,
     BCE_VHCI_CMD_PORT_STATUS = 0x16,
+    BCE_VHCI_CMD_PORT_CONNECT = 0x18,   /* Undocumented - port connection notification from T2 */
 
     BCE_VHCI_CMD_DEVICE_CREATE = 0x30,
     BCE_VHCI_CMD_DEVICE_DESTROY = 0x31,

vhci/transfer.c

@@ -180,21 +180,21 @@ static void bce_vhci_transfer_queue_completion(struct bce_queue_sq *sq)
          */
         if (c->status == BCE_COMPLETION_ABORTED) { /* We flushed the queue */
             pr_debug("bce-vhci: [%02x] Got an abort completion\n", q->endp_addr);
-            if (is_sq_out && atomic_dec_and_test(&q->sq_out_pending))
+            if (is_sq_out && atomic_dec_if_positive(&q->sq_out_pending) == 0)
                 wake_up(&q->sq_out_wait_queue);
             bce_notify_submission_complete(sq);
             continue;
         }
         if (list_empty(&q->endp->urb_list)) {
             pr_err("bce-vhci: [%02x] Got a completion while no requests are pending\n", q->endp_addr);
-            if (is_sq_out && atomic_dec_and_test(&q->sq_out_pending))
+            if (is_sq_out && atomic_dec_if_positive(&q->sq_out_pending) == 0)
                 wake_up(&q->sq_out_wait_queue);
             continue;
         }
         pr_debug("bce-vhci: [%02x] Got a transfer queue completion\n", q->endp_addr);
         urb = list_first_entry(&q->endp->urb_list, struct urb, urb_list);
         bce_vhci_urb_transfer_completion(urb->hcpriv, c);
-        if (is_sq_out && atomic_dec_and_test(&q->sq_out_pending))
+        if (is_sq_out && atomic_dec_if_positive(&q->sq_out_pending) == 0)
             wake_up(&q->sq_out_wait_queue);
         bce_notify_submission_complete(sq);
     }
@@ -204,50 +204,66 @@ static void bce_vhci_transfer_queue_completion(struct bce_queue_sq *sq)
 }
 
 /* Timeout for waiting on pending output requests during pause */
-#define BCE_VHCI_PAUSE_TIMEOUT_MS 5000
+#define BCE_VHCI_PAUSE_TIMEOUT_MS 2000
 
 int bce_vhci_transfer_queue_do_pause(struct bce_vhci_transfer_queue *q)
 {
     unsigned long flags;
     int status;
     int pending;
     long timeout;
-    u8 endp_addr = (u8) (q->endp->desc.bEndpointAddress & 0x8F);
+
+    pr_info("bce-vhci: [%02x] pause: starting (dev=%d)\n", q->endp_addr, q->dev_addr);
 
     spin_lock_irqsave(&q->urb_lock, flags);
     q->active = false;
     spin_unlock_irqrestore(&q->urb_lock, flags);
     bce_vhci_transfer_queue_remove_pending(q);
 
-    if ((status = bce_vhci_cmd_endpoint_set_state(
-            &q->vhci->cq, q->dev_addr, endp_addr, BCE_VHCI_ENDPOINT_PAUSED, &q->state)))
-        return status;
-    if (q->state != BCE_VHCI_ENDPOINT_PAUSED)
-        return -EINVAL;
-
-    if (q->sq_in)
-        bce_cmd_flush_memory_queue(q->vhci->dev->cmd_cmdq, (u16) q->sq_in->qid);
-
+    /*
+     * Wait for pending output transfers to COMPLETE before pausing/flushing.
+     * This ensures commands like keyboard backlight off actually reach the T2
+     * before we abort remaining transfers. Without this, the backlight command
+     * gets aborted by flush and the keyboard stays lit during suspend.
+     */
     if (q->sq_out) {
-        bce_cmd_flush_memory_queue(q->vhci->dev->cmd_cmdq, (u16) q->sq_out->qid);
-        /*
-         * Suspend/resume fix: Wait for all pending output DMA transfers to
-         * complete before returning. This prevents use-after-free when the
-         * queue is destroyed while transfers are still in flight.
-         */
         pending = atomic_read(&q->sq_out_pending);
+        pr_info("bce-vhci: [%02x] pause: %d pending outputs, waiting for completion\n",
+                q->endp_addr, pending);
         if (pending > 0) {
             timeout = wait_event_timeout(q->sq_out_wait_queue,
                     atomic_read(&q->sq_out_pending) == 0,
                     msecs_to_jiffies(BCE_VHCI_PAUSE_TIMEOUT_MS));
             if (timeout == 0) {
                 pending = atomic_read(&q->sq_out_pending);
                 if (pending > 0)
-                    pr_warn("bce-vhci: [%02x] Timeout waiting for %d pending output requests\n",
+                    pr_warn("bce-vhci: [%02x] pause: TIMEOUT waiting for %d pending outputs\n",
                             q->endp_addr, pending);
+            } else {
+                pr_info("bce-vhci: [%02x] pause: pending outputs completed\n", q->endp_addr);
             }
         }
     }
+
+    pr_info("bce-vhci: [%02x] pause: setting endpoint state to PAUSED\n", q->endp_addr);
+    if ((status = bce_vhci_cmd_endpoint_set_state(
+            &q->vhci->cq, q->dev_addr, q->endp_addr, BCE_VHCI_ENDPOINT_PAUSED, &q->state))) {
+        pr_err("bce-vhci: [%02x] pause: set_state failed with %d\n", q->endp_addr, status);
+        return status;
+    }
+    if (q->state != BCE_VHCI_ENDPOINT_PAUSED) {
+        pr_err("bce-vhci: [%02x] pause: unexpected state %d\n", q->endp_addr, q->state);
+        return -EINVAL;
+    }
+
+    pr_info("bce-vhci: [%02x] pause: flushing queues\n", q->endp_addr);
+    if (q->sq_in)
+        bce_cmd_flush_memory_queue(q->vhci->dev->cmd_cmdq, (u16) q->sq_in->qid);
+
+    if (q->sq_out)
+        bce_cmd_flush_memory_queue(q->vhci->dev->cmd_cmdq, (u16) q->sq_out->qid);
+
+    pr_info("bce-vhci: [%02x] pause: done\n", q->endp_addr);
     return 0;
 }
 
@@ -259,10 +275,9 @@ int bce_vhci_transfer_queue_do_resume(struct bce_vhci_transfer_queue *q)
     int status;
     struct urb *urb, *urbt;
     struct bce_vhci_urb *vurb;
-    u8 endp_addr = (u8) (q->endp->desc.bEndpointAddress & 0x8F);
 
     if ((status = bce_vhci_cmd_endpoint_set_state(
-            &q->vhci->cq, q->dev_addr, endp_addr, BCE_VHCI_ENDPOINT_ACTIVE, &q->state)))
+            &q->vhci->cq, q->dev_addr, q->endp_addr, BCE_VHCI_ENDPOINT_ACTIVE, &q->state)))
         return status;
     if (q->state != BCE_VHCI_ENDPOINT_ACTIVE)
         return -EINVAL;