From 332b108f41bdc096773b880741c0675442fa6e6d Mon Sep 17 00:00:00 2001
From: Justin Kotalik <jkotalik12@gmail.com>
Date: Wed, 25 Oct 2017 11:46:34 -0700
Subject: [PATCH] Changes PostCompletion to handle OnAsyncCompletion after
 managed request has completed.  (#212)

---
 src/AspNetCore/Inc/inprocessstoredcontext.h   | 23 +++++++++++++
 src/AspNetCore/Src/inprocessapplication.cxx   | 16 +++++++--
 src/AspNetCore/Src/inprocessstoredcontext.cxx | 34 +++++++++++++++++++
 src/AspNetCore/Src/managedexports.cxx         | 24 +++++++++++++
 4 files changed, 94 insertions(+), 3 deletions(-)

diff --git a/src/AspNetCore/Inc/inprocessstoredcontext.h b/src/AspNetCore/Inc/inprocessstoredcontext.h
index 01a6045609..64244692d4 100644
--- a/src/AspNetCore/Inc/inprocessstoredcontext.h
+++ b/src/AspNetCore/Inc/inprocessstoredcontext.h
@@ -9,6 +9,7 @@ public:
         IHttpContext* pHttpContext,
         PVOID pvManagedContext
     );
+
     ~IN_PROCESS_STORED_CONTEXT();
 
     virtual
@@ -46,6 +47,26 @@ public:
         VOID
     );
 
+    BOOL
+    QueryIsManagedRequestComplete(
+        VOID
+    );
+
+    VOID
+    IndicateManagedRequestComplete(
+        VOID
+    ); 
+
+    REQUEST_NOTIFICATION_STATUS
+    QueryAsyncCompletionStatus(
+        VOID
+    );
+
+    VOID
+    SetAsyncCompletionStatus(
+        REQUEST_NOTIFICATION_STATUS requestNotificationStatus
+    );
+
     static
     HRESULT
     GetInProcessStoredContext(
@@ -63,5 +84,7 @@ public:
 private:
     PVOID m_pManagedHttpContext;
     IHttpContext* m_pHttpContext;
+    BOOL m_fManagedRequestComplete;
+    REQUEST_NOTIFICATION_STATUS m_requestNotificationStatus;
 };
 
diff --git a/src/AspNetCore/Src/inprocessapplication.cxx b/src/AspNetCore/Src/inprocessapplication.cxx
index 7224125635..f57518b7e4 100644
--- a/src/AspNetCore/Src/inprocessapplication.cxx
+++ b/src/AspNetCore/Src/inprocessapplication.cxx
@@ -30,6 +30,7 @@ IN_PROCESS_APPLICATION::OnAsyncCompletion(
 {
     HRESULT hr;
     IN_PROCESS_STORED_CONTEXT* pInProcessStoredContext = NULL;
+    REQUEST_NOTIFICATION_STATUS dwRequestNotificationStatus = RQ_NOTIFICATION_CONTINUE;
 
     hr = IN_PROCESS_STORED_CONTEXT::GetInProcessStoredContext(pHttpContext, &pInProcessStoredContext);
     if (FAILED(hr))
@@ -38,9 +39,18 @@ IN_PROCESS_APPLICATION::OnAsyncCompletion(
         pHttpContext->GetResponse()->SetStatus(500, "Internal Server Error", 19, hr);
         return RQ_NOTIFICATION_FINISH_REQUEST;
     }
-
-    // Call the managed handler for async completion.
-    return m_AsyncCompletionHandler(pInProcessStoredContext->QueryManagedHttpContext(), hrCompletionStatus, cbCompletion);
+    else if (pInProcessStoredContext->QueryIsManagedRequestComplete())
+    {
+        // means PostCompletion has been called and this is the associated callback.
+        dwRequestNotificationStatus = pInProcessStoredContext->QueryAsyncCompletionStatus();
+        // TODO cleanup whatever disconnect listener there is
+        return dwRequestNotificationStatus;
+    }
+    else
+    {
+        // Call the managed handler for async completion.
+        return m_AsyncCompletionHandler(pInProcessStoredContext->QueryManagedHttpContext(), hrCompletionStatus, cbCompletion);
+    }
 }
 
 BOOL
diff --git a/src/AspNetCore/Src/inprocessstoredcontext.cxx b/src/AspNetCore/Src/inprocessstoredcontext.cxx
index 27704b6c96..c0b2d8ff16 100644
--- a/src/AspNetCore/Src/inprocessstoredcontext.cxx
+++ b/src/AspNetCore/Src/inprocessstoredcontext.cxx
@@ -8,8 +8,10 @@ IN_PROCESS_STORED_CONTEXT::IN_PROCESS_STORED_CONTEXT(
     PVOID pMangedHttpContext
 )
 {
+    // TODO if we want to go by IIS patterns, we should have these in a separate initialize function 
     m_pManagedHttpContext = pMangedHttpContext;
     m_pHttpContext = pHttpContext;
+    m_fManagedRequestComplete = FALSE;
 }
 
 IN_PROCESS_STORED_CONTEXT::~IN_PROCESS_STORED_CONTEXT()
@@ -32,6 +34,38 @@ IN_PROCESS_STORED_CONTEXT::QueryHttpContext(
     return m_pHttpContext;
 }
 
+BOOL
+IN_PROCESS_STORED_CONTEXT::QueryIsManagedRequestComplete(
+    VOID 
+)
+{
+    return m_fManagedRequestComplete;
+}
+
+VOID
+IN_PROCESS_STORED_CONTEXT::IndicateManagedRequestComplete(
+    VOID
+)
+{
+    m_fManagedRequestComplete = TRUE;
+}
+
+REQUEST_NOTIFICATION_STATUS
+IN_PROCESS_STORED_CONTEXT::QueryAsyncCompletionStatus(
+    VOID
+)
+{
+    return m_requestNotificationStatus;
+}
+
+VOID
+IN_PROCESS_STORED_CONTEXT::SetAsyncCompletionStatus(
+    REQUEST_NOTIFICATION_STATUS requestNotificationStatus
+)
+{
+    m_requestNotificationStatus = requestNotificationStatus;
+}
+
 HRESULT
 IN_PROCESS_STORED_CONTEXT::GetInProcessStoredContext(
     IHttpContext*               pHttpContext,
diff --git a/src/AspNetCore/Src/managedexports.cxx b/src/AspNetCore/Src/managedexports.cxx
index 4af8feadab..f9dd492fcc 100644
--- a/src/AspNetCore/Src/managedexports.cxx
+++ b/src/AspNetCore/Src/managedexports.cxx
@@ -62,6 +62,30 @@ http_post_completion(
     return pHttpContext->PostCompletion(cbBytes);
 }
 
+EXTERN_C __MIDL_DECLSPEC_DLLEXPORT
+HRESULT
+http_set_completion_status(
+    _In_ IHttpContext* pHttpContext,
+    REQUEST_NOTIFICATION_STATUS requestNotificationStatus
+)
+{
+    HRESULT hr = S_OK;
+    IN_PROCESS_STORED_CONTEXT* pInProcessStoredContext = NULL;
+
+    hr = IN_PROCESS_STORED_CONTEXT::GetInProcessStoredContext(
+        pHttpContext,
+        &pInProcessStoredContext
+    );
+
+    if (FAILED(hr))
+    {
+        return hr;
+    }
+    pInProcessStoredContext->IndicateManagedRequestComplete();
+    pInProcessStoredContext->SetAsyncCompletionStatus(requestNotificationStatus);
+    return hr;
+}
+
 EXTERN_C __MIDL_DECLSPEC_DLLEXPORT
 HRESULT
 http_set_managed_context(