From f5e8b4cd287f1f5f3abf12f1e9eb5b9504f999e7 Mon Sep 17 00:00:00 2001
From: Ryan Prichard <ryan.prichard@gmail.com>
Date: Mon, 28 Dec 2015 21:49:41 -0600
Subject: [PATCH] Fix a possible concurrency issue with CancelIo and winpty
 shutdown.

CancelIo requests that the I/O operation end, but the I/O operation may
still be going when the function returns.

https://blogs.msdn.microsoft.com/oldnewthing/20110202-00/?p=11613/
---
 src/agent/NamedPipe.cc            | 14 ++++++++++++++
 src/agent/NamedPipe.h             |  1 +
 src/unix-adapter/InputHandler.cc  |  1 +
 src/unix-adapter/OutputHandler.cc |  1 +
 4 files changed, 17 insertions(+)

diff --git a/src/agent/NamedPipe.cc b/src/agent/NamedPipe.cc
index 81413e2..b0d0fef 100644
--- a/src/agent/NamedPipe.cc
+++ b/src/agent/NamedPipe.cc
@@ -121,6 +121,18 @@ int NamedPipe::IoWorker::service()
     return progress;
 }
 
+// This function is called after CancelIo has returned.  We need to block until
+// the I/O operations have completed, which should happen very quickly.
+// https://blogs.msdn.microsoft.com/oldnewthing/20110202-00/?p=11613
+void NamedPipe::IoWorker::waitForCanceledIo()
+{
+    if (m_pending) {
+        DWORD actual = 0;
+        GetOverlappedResult(m_namedPipe->m_handle, &m_over, &actual, TRUE);
+        m_pending = false;
+    }
+}
+
 HANDLE NamedPipe::IoWorker::getWaitEvent()
 {
     return m_pending ? m_event : NULL;
@@ -247,6 +259,8 @@ void NamedPipe::closePipe()
     if (m_handle == NULL)
         return;
     CancelIo(m_handle);
+    m_inputWorker->waitForCanceledIo();
+    m_outputWorker->waitForCanceledIo();
     delete m_inputWorker;
     delete m_outputWorker;
     CloseHandle(m_handle);
diff --git a/src/agent/NamedPipe.h b/src/agent/NamedPipe.h
index f5d87cc..8a1126b 100644
--- a/src/agent/NamedPipe.h
+++ b/src/agent/NamedPipe.h
@@ -43,6 +43,7 @@ private:
         IoWorker(NamedPipe *namedPipe);
         virtual ~IoWorker();
         int service();
+        void waitForCanceledIo();
         HANDLE getWaitEvent();
     protected:
         NamedPipe *m_namedPipe;
diff --git a/src/unix-adapter/InputHandler.cc b/src/unix-adapter/InputHandler.cc
index 0ed35ed..3f1bed8 100755
--- a/src/unix-adapter/InputHandler.cc
+++ b/src/unix-adapter/InputHandler.cc
@@ -109,6 +109,7 @@ void InputHandler::threadProc() {
                 trace("InputHandler: shutting down, canceling I/O");
                 assert(m_shouldShutdown);
                 CancelIo(m_winpty);
+                GetOverlappedResult(m_winpty, &over, &written, TRUE);
                 break;
             }
             assert(waitRet == WAIT_OBJECT_0);
diff --git a/src/unix-adapter/OutputHandler.cc b/src/unix-adapter/OutputHandler.cc
index b103d5d..64b5163 100755
--- a/src/unix-adapter/OutputHandler.cc
+++ b/src/unix-adapter/OutputHandler.cc
@@ -83,6 +83,7 @@ void OutputHandler::threadProc() {
                 trace("OutputHandler: shutting down, canceling I/O");
                 assert(m_shouldShutdown);
                 CancelIo(m_winpty);
+                GetOverlappedResult(m_winpty, &over, &numRead, TRUE);
                 break;
             }
             assert(waitRet == WAIT_OBJECT_0);