From 29798a28aa126f88376400df484d3b9a22f44a01 Mon Sep 17 00:00:00 2001
From: KyungWoon Cho <cezanne@codemayo.com>
Date: Sat, 2 Jan 2021 21:38:27 +0900
Subject: [PATCH] enable to abort an actively running vhci(ude)

There was a race condition between WDF req cancellation and detach
command. urbr completion is decomposed into 2 steps: unmarking cancel
& completing WDF req.
---
 driver/vhci_ude/vhci_plugout.c |  2 ++
 driver/vhci_ude/vhci_read.c    |  5 ++++-
 driver/vhci_ude/vhci_urbr.c    | 33 +++++++++++++++++++++++++++------
 driver/vhci_ude/vhci_urbr.h    |  2 ++
 driver/vhci_ude/vhci_write.c   | 10 +++++++++-
 5 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/driver/vhci_ude/vhci_plugout.c b/driver/vhci_ude/vhci_plugout.c
index 4c9a13da..22a70a3d 100644
--- a/driver/vhci_ude/vhci_plugout.c
+++ b/driver/vhci_ude/vhci_plugout.c
@@ -38,6 +38,8 @@ abort_all_pending_urbrs(pctx_vusb_t vusb)
 		urbr = CONTAINING_RECORD(vusb->head_urbr.Flink, urb_req_t, list_all);
 		RemoveEntryListInit(&urbr->list_all);
 		RemoveEntryListInit(&urbr->list_state);
+		if (!unmark_cancelable_urbr(urbr))
+			continue;
 		WdfSpinLockRelease(vusb->spin_lock);
 
 		abort_pending_urbr(urbr);
diff --git a/driver/vhci_ude/vhci_read.c b/driver/vhci_ude/vhci_read.c
index 3a3cce19..4832c339 100644
--- a/driver/vhci_ude/vhci_read.c
+++ b/driver/vhci_ude/vhci_read.c
@@ -89,10 +89,13 @@ read_vusb(pctx_vusb_t vusb, WDFREQUEST req)
 	}
 
 	if (status != STATUS_SUCCESS) {
+		BOOLEAN	unmarked;
 		RemoveEntryListInit(&urbr->list_all);
+		unmarked = unmark_cancelable_urbr(urbr);
 		WdfSpinLockRelease(vusb->spin_lock);
 
-		complete_urbr(urbr, status);
+		if (unmarked)
+			complete_urbr(urbr, status);
 	}
 	else {
 		if (vusb->len_sent_partial == 0) {
diff --git a/driver/vhci_ude/vhci_urbr.c b/driver/vhci_ude/vhci_urbr.c
index 0bf20538..32728337 100644
--- a/driver/vhci_ude/vhci_urbr.c
+++ b/driver/vhci_ude/vhci_urbr.c
@@ -175,9 +175,14 @@ urbr_cancelled(_In_ WDFREQUEST req)
 	}
 	WdfSpinLockRelease(vusb->spin_lock);
 
-	submit_urbr_unlink(urbr->ep, urbr->seq_num);
-	TRD(URBR, "cancelled urbr destroyed: %!URBR!", urbr);
-	complete_urbr(urbr, STATUS_CANCELLED);
+	if (urbr != NULL && urbr->seq_num != 0) {
+		submit_urbr_unlink(urbr->ep, urbr->seq_num);
+		TRD(URBR, "cancelled urbr destroyed: %!URBR!", urbr);
+		complete_urbr(urbr, STATUS_CANCELLED);
+	}
+	else {
+		UdecxUrbCompleteWithNtStatus(req, STATUS_CANCELLED);
+	}
 }
 
 NTSTATUS
@@ -305,6 +310,23 @@ submit_req_reset_pipe(pctx_ep_t ep, WDFREQUEST req)
 	return submit_urbr_free(urbr);
 }
 
+BOOLEAN
+unmark_cancelable_urbr(purb_req_t urbr)
+{
+	WDFREQUEST	req;
+	NTSTATUS	status;
+
+	req = urbr->req;
+	if (req == NULL)
+		return TRUE;
+	if (urbr->type != URBR_TYPE_URB)
+		return TRUE;
+	status = WdfRequestUnmarkCancelable(req);
+	if (status == STATUS_CANCELLED)
+		return FALSE;
+	return TRUE;
+}
+
 void
 complete_urbr(purb_req_t urbr, NTSTATUS status)
 {
@@ -315,12 +337,11 @@ complete_urbr(purb_req_t urbr, NTSTATUS status)
 		if (urbr->type != URBR_TYPE_URB)
 			WdfRequestComplete(req, status);
 		else {
-			if (status != STATUS_CANCELLED)
-				WdfRequestUnmarkCancelable(req);
 			if (status == STATUS_SUCCESS)
 				UdecxUrbComplete(req, urbr->u.urb->UrbHeader.Status);
-			else 
+			else {
 				UdecxUrbCompleteWithNtStatus(req, status);
+			}
 		}
 	}
 	free_urbr(urbr);
diff --git a/driver/vhci_ude/vhci_urbr.h b/driver/vhci_ude/vhci_urbr.h
index f14cb7d4..d6876299 100644
--- a/driver/vhci_ude/vhci_urbr.h
+++ b/driver/vhci_ude/vhci_urbr.h
@@ -68,5 +68,7 @@ submit_req_reset_pipe(pctx_ep_t ep, WDFREQUEST req);
 extern NTSTATUS
 store_urbr(WDFREQUEST req_read, purb_req_t urbr);
 
+extern BOOLEAN
+unmark_cancelable_urbr(purb_req_t urbr);
 extern void
 complete_urbr(purb_req_t urbr, NTSTATUS status);
diff --git a/driver/vhci_ude/vhci_write.c b/driver/vhci_ude/vhci_write.c
index f1273635..0401104b 100644
--- a/driver/vhci_ude/vhci_write.c
+++ b/driver/vhci_ude/vhci_write.c
@@ -52,7 +52,15 @@ write_vusb(pctx_vusb_t vusb, WDFREQUEST req_write)
 	}
 
 	status = fetch_urbr(urbr, hdr);
-	complete_urbr(urbr, status);
+
+	WdfSpinLockAcquire(vusb->spin_lock);
+	if (unmark_cancelable_urbr(urbr)) {
+		WdfSpinLockRelease(vusb->spin_lock);
+		complete_urbr(urbr, status);
+	}
+	else {
+		WdfSpinLockRelease(vusb->spin_lock);
+	}
 out:
 	TRD(WRITE, "Leave: %!STATUS!", status);
 }