have the M3 and M4 correctly trigger

12 files changed, 544 insertions, 287 deletions

diff --git a/network/wlan/WIFICX/drivercode/transition_steps.hpp b/network/wlan/WIFICX/drivercode/transition_steps.hpp
new file mode 100644
index 00000000..dc04eab2
--- /dev/null
+++ b/network/wlan/WIFICX/drivercode/transition_steps.hpp
@@ -0,0 +1,17 @@
+template<typename Param, typename... StepFns>
+NTSTATUS ExecuteSteps(TransitionContext& ctx, Param& p, StepFns... fns)
+{
+    NTSTATUS status = STATUS_SUCCESS;
+    auto seq = { (status = (fns(ctx, p)), status == STATUS_SUCCESS ? 0 : 1)... };
+    UNREFERENCED_PARAMETER(seq);
+    return status;
+}
+
+// In a traits Handle():
+static NTSTATUS Handle(TransitionContext& ctx, ParamType& params)
+{
+    return ExecuteSteps(ctx, params,
+        [](TransitionContext& c, ParamType& par){ return c.DevCtx->wifiHAL->WifiIhvSetRadioState(par, c.Header); },
+        [](TransitionContext& c, ParamType&){ return STATUS_SUCCESS; } // extra step
+    );
+}
\ No newline at end of file
diff --git a/network/wlan/WIFICX/drivercode/umkmfusion.h b/network/wlan/WIFICX/drivercode/umkmfusion.h
index f94454f8..35e5a6cf 100644
--- a/network/wlan/WIFICX/drivercode/umkmfusion.h
+++ b/network/wlan/WIFICX/drivercode/umkmfusion.h
@@ -19,17 +19,23 @@ namespace Wifi
         }
 
         // Trace the HRESULT value for diagnostics
-
-
         return SUCCEEDED(hr) ? STATUS_SUCCESS : STATUS_UNSUCCESSFUL;
     }
 
-    _inline
-    NTSTATUS ConvertNDISSTATUSoNTSTATUS(NDIS_STATUS ndisStatus) {
-        
-        return ndisStatus > 0 ? STATUS_SUCCESS : STATUS_UNSUCCESSFUL;
+    __inline
+    NTSTATUS ConvertNDISSTATUSToNTSTATUS(NDIS_STATUS ndisStatus)
+    {
+        // NDIS_STATUS_SUCCESS is 0x00000000
+        // Any non-success NDIS status should map to STATUS_UNSUCCESSFUL
+        return (ndisStatus == NDIS_STATUS_SUCCESS) ? STATUS_SUCCESS : STATUS_UNSUCCESSFUL;
     }
 
+    __inline
+    NDIS_STATUS ConvertNTSTATUSToNDISSTATUS(NTSTATUS ntStatus)
+    {
+        // Use NT_SUCCESS macro to check NTSTATUS
+        return NT_SUCCESS(ntStatus) ? NDIS_STATUS_SUCCESS : NDIS_STATUS_FAILURE;
+    }
 
     _Must_inspect_result_
     _inline
diff --git a/network/wlan/WIFICX/drivercode/wifihal.cpp b/network/wlan/WIFICX/drivercode/wifihal.cpp
index 824ab6ac..edeb8fe0 100644
--- a/network/wlan/WIFICX/drivercode/wifihal.cpp
+++ b/network/wlan/WIFICX/drivercode/wifihal.cpp
@@ -6,7 +6,7 @@
 #include "WifiHal.h"
 #include "WifiHal.tmh"
 
-NTSTATUS WifiHAL::WifiIhvSetDeviceCapabilities(WDFDEVICE Device)
+NTSTATUS WifiHAL::WifiIhvSetDeviceCapabilities()
 {
     WIFI_DEVICE_CAPABILITIES deviceCaps = {};
     WIFI_DEVICE_CAPABILITIES_INIT(&deviceCaps);
@@ -35,7 +35,7 @@ NTSTATUS WifiHAL::WifiIhvSetDeviceCapabilities(WDFDEVICE Device)
     deviceCaps.BeaconReportsImplemented = FALSE;
 
     WX_RETURN_NTSTATUS_IF_NOT_NT_SUCCESS_MSG(
-        WifiDeviceSetDeviceCapabilities(Device, &deviceCaps), 
+        WifiDeviceSetDeviceCapabilities(m_Device, &deviceCaps), 
         "Failed to set device capabilities");
 
     WIFI_STATION_CAPABILITIES StationCaps = {};
@@ -133,7 +133,7 @@ NTSTATUS WifiHAL::WifiIhvSetDeviceCapabilities(WDFDEVICE Device)
     }
 
     WX_RETURN_NTSTATUS_IF_NOT_NT_SUCCESS_MSG(
-        WifiDeviceSetStationCapabilities(Device, &StationCaps), 
+        WifiDeviceSetStationCapabilities(m_Device, &StationCaps), 
         "Failed to set station capabilities");
 
     WIFI_BAND_CAPABILITIES BandCaps = {};
@@ -392,7 +392,7 @@ NTSTATUS WifiHAL::WifiIhvSetDeviceCapabilities(WDFDEVICE Device)
     BandCaps.BandInfoList = BandInfo;
 
     WX_RETURN_NTSTATUS_IF_NOT_NT_SUCCESS_MSG(
-        WifiDeviceSetBandCapabilities(Device, &BandCaps), 
+        WifiDeviceSetBandCapabilities(m_Device, &BandCaps), 
         "Failed to set band capabilities");
 
     WIFI_PHY_CAPABILITIES PhyCaps = {};
@@ -426,13 +426,13 @@ NTSTATUS WifiHAL::WifiIhvSetDeviceCapabilities(WDFDEVICE Device)
     PhyCaps.PhyInfoList = PhyInfoList;
 
     WX_RETURN_NTSTATUS_IF_NOT_NT_SUCCESS_MSG(
-        WifiDeviceSetPhyCapabilities(Device, &PhyCaps), 
+        WifiDeviceSetPhyCapabilities(m_Device, &PhyCaps), 
         "Failed to set PHY capabilities");
 
     return STATUS_SUCCESS;
 }
 
-NTSTATUS WifiHAL::WifiIhvReset(WDFDEVICE Device, const WDI_TASK_DOT11_RESET_PARAMETERS& ResetParameters, const PWDI_MESSAGE_HEADER pWdiHeader, PCTLV_CONTEXT)
+NTSTATUS WifiHAL::WifiIhvReset(const WDI_TASK_DOT11_RESET_PARAMETERS& ResetParameters, const PWDI_MESSAGE_HEADER)
 {
     if (0 == ResetParameters.Optional.ResetMACAddress_IsPresent)
     {
@@ -453,13 +453,12 @@ NTSTATUS WifiHAL::WifiIhvReset(WDFDEVICE Device, const WDI_TASK_DOT11_RESET_PARA
 
     // Reset the connection ID in case the previous connection attempt did not complete
     m_LastConnectEntryId = 0;
-    WifiIhvSendIndicationToOs(Device, pWdiHeader, WDI_INDICATION_DOT11_RESET_COMPLETE, pWdiHeader->TransactionId, nullptr, 0);
-
+    
     return STATUS_SUCCESS;
 }
 
 _Use_decl_annotations_
-NTSTATUS WifiHAL::WifiIhvSetRadioState(WDFDEVICE Device, const WDI_SET_RADIO_STATE_PARAMETERS& RadioState, const PWDI_MESSAGE_HEADER pWdiHeader, PCTLV_CONTEXT Context)
+NTSTATUS WifiHAL::WifiIhvSetRadioState(const WDI_SET_RADIO_STATE_PARAMETERS& RadioState, const PWDI_MESSAGE_HEADER pWdiHeader)
 {
     if (RadioState.SoftwareRadioState != m_CurrentRadioState)
     {
@@ -473,20 +472,18 @@ NTSTATUS WifiHAL::WifiIhvSetRadioState(WDFDEVICE Device, const WDI_SET_RADIO_STA
 
         RadioStatusParams.RadioState.HardwareState = TRUE;
         RadioStatusParams.RadioState.SoftwareState = m_CurrentRadioState;
-        if (GenerateWdiIndicationRadioStatus(&RadioStatusParams, 0, Context, &cbOutput, &pOutput) == NDIS_STATUS_SUCCESS)
+        if (GenerateWdiIndicationRadioStatus(&RadioStatusParams, 0, m_TlvContext, &cbOutput, &pOutput) == NDIS_STATUS_SUCCESS)
         {
-            WifiIhvSendIndicationToOs(Device, pWdiHeader, WDI_INDICATION_RADIO_STATUS, 0, pOutput, cbOutput);
+            WifiIhvSendUnsolicitedIndicationToOs(m_Device, pWdiHeader, WDI_INDICATION_RADIO_STATUS, pOutput, cbOutput);
             FreeGenerated(pOutput);
         }
     }
 
-    WifiIhvSendIndicationToOs(Device, pWdiHeader, WDI_INDICATION_SET_RADIO_STATE_COMPLETE, pWdiHeader->TransactionId, nullptr, 0);
-
     return STATUS_SUCCESS;
 }
 
 _Use_decl_annotations_
-NTSTATUS WifiHAL::WifiIhvScan(WDFDEVICE Device, WDI_SCAN_PARAMETERS& ScanParameters, const PWDI_MESSAGE_HEADER pWdiHeader, PCTLV_CONTEXT)
+NTSTATUS WifiHAL::WifiIhvScan(WDI_SCAN_PARAMETERS& ScanParameters, const PWDI_MESSAGE_HEADER pWdiHeader)
 {
     for (UINT8 connectEntry = 1; connectEntry < ConnectEntryId_MAX; connectEntry++)
     {
@@ -510,11 +507,10 @@ NTSTATUS WifiHAL::WifiIhvScan(WDFDEVICE Device, WDI_SCAN_PARAMETERS& ScanParamet
             }
 
             // Send the BSS entry indication
-            WifiIhvSendIndicationToOs(
-                Device,
+            WifiIhvSendUnsolicitedIndicationToOs(
+                m_Device,
                 pWdiHeader,
                 WDI_INDICATION_BSS_ENTRY_LIST,
-                0,
                 g_ConnectEntries[connectEntry].pTlvBssEntry,
                 g_ConnectEntries[connectEntry].TlvBssEntrySize);
         }
@@ -527,16 +523,14 @@ NTSTATUS WifiHAL::WifiIhvScan(WDFDEVICE Device, WDI_SCAN_PARAMETERS& ScanParamet
         (ScanParameters.SSIDList.pElements[0].pElements[4] == 'H') && (ScanParameters.SSIDList.pElements[0].pElements[5] == 'I') &&
         (ScanParameters.SSIDList.pElements[0].pElements[6] == 'D') && (ScanParameters.SSIDList.pElements[0].pElements[7] == 'E'))
     {
-        WifiIhvSendIndicationToOs(
-            Device, pWdiHeader, WDI_INDICATION_BSS_ENTRY_LIST, 0, s_TLV_BSS_Entry_ProbeResponse_8_Hidden, sizeof(s_TLV_BSS_Entry_ProbeResponse_8_Hidden));
+        WifiIhvSendUnsolicitedIndicationToOs(
+            m_Device, pWdiHeader, WDI_INDICATION_BSS_ENTRY_LIST, s_TLV_BSS_Entry_ProbeResponse_8_Hidden, sizeof(s_TLV_BSS_Entry_ProbeResponse_8_Hidden));
     }
 
-    WifiIhvSendIndicationToOs(Device, pWdiHeader, WDI_INDICATION_SCAN_COMPLETE, pWdiHeader->TransactionId, nullptr, 0);
-
     return STATUS_SUCCESS;
 }
 
-NTSTATUS WifiHAL::WifiIhvConnect(WDFDEVICE Device, WDI_TASK_CONNECT_PARAMETERS& ConnectParameters, const PWDI_MESSAGE_HEADER pWdiHeader, PCTLV_CONTEXT)
+NTSTATUS WifiHAL::WifiIhvConnect(WDI_TASK_CONNECT_PARAMETERS& ConnectParameters, const PWDI_MESSAGE_HEADER pWdiHeader)
 {
     NT_ASSERT(m_LastConnectEntryId == 0);
     if (m_LastConnectEntryId != 0) // Not Disconnected State
@@ -548,26 +542,25 @@ NTSTATUS WifiHAL::WifiIhvConnect(WDFDEVICE Device, WDI_TASK_CONNECT_PARAMETERS&
     }
 
     WX_RETURN_NTSTATUS_IF_NOT_NT_SUCCESS_MSG(WifiIhvPerformAssociation(
-        Device, &ConnectParameters.PreferredBSSEntryList, &ConnectParameters.ConnectParameters.AuthenticationAlgorithms, pWdiHeader),
+        &ConnectParameters.PreferredBSSEntryList, &ConnectParameters.ConnectParameters.AuthenticationAlgorithms, pWdiHeader),
         "Failed to perform association");
 
     //
     // WPA3-SAE requires the SAE Exchange, so do not complete the Connection request until the SAE exchange is complete
     //
-    if (WDI_AUTH_ALGO_WPA3_SAE != m_LastAuthAlgo)
+    if (WDI_AUTH_ALGO_WPA3_SAE == m_LastAuthAlgo)
     {
-        WifiIhvSendIndicationToOs(Device, pWdiHeader, WDI_INDICATION_CONNECT_COMPLETE, pWdiHeader->TransactionId, nullptr, 0);
+        return STATUS_PENDING;
     }
-
     return STATUS_SUCCESS;
 }
 
 _Use_decl_annotations_
-NTSTATUS WifiHAL::WifiIhvSendLinkStateIndication(_In_ WDFDEVICE Device, _In_ PWDI_MESSAGE_HEADER pWdiHeader, ULONG numLinks)
+NTSTATUS WifiHAL::WifiIhvSendLinkStateIndication(_In_ PWDI_MESSAGE_HEADER pWdiHeader, ULONG numLinks)
 {
     // Report link quality
     NTSTATUS ntStatus = STATUS_SUCCESS;
-    PWIFI_IHV_DEVICE_CONTEXT pDeviceContext = WifiGetIhvDeviceContext(Device);
+    PWIFI_IHV_DEVICE_CONTEXT pDeviceContext = WifiGetIhvDeviceContext(m_Device);
     WDI_INDICATION_LINK_STATE_CHANGE_PARAMETERS linkStateChangeParameters = {};
     WDI_LINK_INFO_CONTAINER pLinkInfo[2] = {};
     UINT8* pOutput = nullptr;
@@ -612,7 +605,7 @@ NTSTATUS WifiHAL::WifiIhvSendLinkStateIndication(_In_ WDFDEVICE Device, _In_ PWD
     ntStatus = GenerateWdiIndicationLinkStateChangeFromIhv(&linkStateChangeParameters, 0, &pDeviceContext->TlvContext, &cbOutput, &pOutput);
     if (STATUS_SUCCESS == ntStatus)
     {
-        WifiIhvSendIndicationToOs(Device, pWdiHeader, WDI_INDICATION_LINK_STATE_CHANGE, 0, pOutput, cbOutput);
+        WifiIhvSendUnsolicitedIndicationToOs(m_Device, pWdiHeader, WDI_INDICATION_LINK_STATE_CHANGE, pOutput, cbOutput);
         FreeGenerated(pOutput);
     }
     else
@@ -625,7 +618,6 @@ NTSTATUS WifiHAL::WifiIhvSendLinkStateIndication(_In_ WDFDEVICE Device, _In_ PWD
 
 _Use_decl_annotations_
 NTSTATUS WifiHAL::WifiIhvPerformAssociation(
-    WDFDEVICE Device,
     struct ArrayOfElements<WDI_CONNECT_BSS_ENTRY_CONTAINER>* pPreferredBSSEntryList,
     struct ArrayOfElements<WDI_AUTH_ALGORITHM>* pAuthenticationAlgorithms,
     PWDI_MESSAGE_HEADER pWdiHeader)
@@ -640,7 +632,7 @@ NTSTATUS WifiHAL::WifiIhvPerformAssociation(
 #ifdef WIFI_IHV_NETV
     WifiNetvDevice* pDeviceContext = WifiNetvDeviceGetContext(Device);
 #else
-    PWIFI_IHV_DEVICE_CONTEXT pDeviceContext = WifiGetIhvDeviceContext(Device);
+    PWIFI_IHV_DEVICE_CONTEXT pDeviceContext = WifiGetIhvDeviceContext(m_Device);
 #endif // WIFI_IHV_NETV
     ULONG assocStatus = WDI_ASSOC_STATUS_SUCCESS;
 
@@ -693,7 +685,7 @@ NTSTATUS WifiHAL::WifiIhvPerformAssociation(
                         {
                             WFCInfo("[SAE] Indicating request for COMMIT_REQUEST_PARAMS_NEEDED ...");
 
-                            WifiIhvSendIndicationToOs(Device, pWdiHeader, WDI_INDICATION_SAE_AUTH_PARAMS_NEEDED, 0, pOutput, cbOutput);
+                            WifiIhvSendUnsolicitedIndicationToOs(m_Device, pWdiHeader, WDI_INDICATION_SAE_AUTH_PARAMS_NEEDED, pOutput, cbOutput);
 
                             FreeGenerated(pOutput);
 
@@ -705,11 +697,10 @@ NTSTATUS WifiHAL::WifiIhvPerformAssociation(
                     RtlCopyMemory(&assocStatus, &pAssociationResult[18], sizeof(ULONG));
 
                     // Send the association indication
-                    WifiIhvSendIndicationToOs(
-                        Device,
+                    WifiIhvSendUnsolicitedIndicationToOs(
+                        m_Device,
                         pWdiHeader,
                         WDI_INDICATION_ASSOCIATION_RESULT,
-                        0,
                         g_ConnectEntries[connectEntry].pTlvAssociationResult,
                         g_ConnectEntries[connectEntry].TlvAssociationResultSize);
 
@@ -743,14 +734,22 @@ NTSTATUS WifiHAL::WifiIhvPerformAssociation(
     if (STATUS_SUCCESS == ntStatus && (WDI_AUTH_ALGO_WPA3_SAE != m_LastAuthAlgo))
     {
         // Report link quality
-        ntStatus = WifiIhvSendLinkStateIndication(Device, pWdiHeader, 2);
+        ntStatus = WifiIhvSendLinkStateIndication(pWdiHeader, 2);
     }
 
     return ntStatus;
 }
 
-NTSTATUS WifiHAL::WifiIhvSetSaeAuthParams(WDFDEVICE Device, const WDI_SET_SAE_AUTH_PARAMS_COMMAND& setSAEAuthParams, const PWDI_MESSAGE_HEADER pWdiHeader, PCTLV_CONTEXT Context)
+NTSTATUS WifiHAL::WifiIhvSetSaeAuthParams(const WDI_SET_SAE_AUTH_PARAMS_COMMAND& setSAEAuthParams, const PWDI_MESSAGE_HEADER pWdiHeader)
 {
+    //Since this is DIRECT OID, need to check the m_LastConnectTransactionId match
+    if (pWdiHeader->TransactionId != m_LastConnectTransactionId) 
+    {
+        WFCError("WDI_SET_SAE_AUTH_PARAMS called with invalid TransactionId: %llu, expected: %llu\n",
+            pWdiHeader->TransactionId, m_LastConnectTransactionId);
+        return STATUS_INVALID_DEVICE_REQUEST;   
+    }
+
     WDI_INDICATION_SAE_AUTH_PARAMS_NEEDED_PARAMETERS SAEAuthParamsNeeded{};
     UINT8* pOutput = nullptr;
     ULONG cbOutput = 0;
@@ -826,21 +825,20 @@ NTSTATUS WifiHAL::WifiIhvSetSaeAuthParams(WDFDEVICE Device, const WDI_SET_SAE_AU
             g_ConnectEntries[m_LastConnectEntryId].pTlvAssociationResult[27] = (UCHAR)WDI_AUTH_ALGO_RSNA_PSK;
         }
 
-        WifiIhvSendIndicationToOs(
-            Device,
+        WifiIhvSendUnsolicitedIndicationToOs(
+            m_Device,
             pWdiHeader,
             WDI_INDICATION_ASSOCIATION_RESULT,
-            0,
             g_ConnectEntries[m_LastConnectEntryId].pTlvAssociationResult,
             g_ConnectEntries[m_LastConnectEntryId].TlvAssociationResultSize);
         RtlCopyMemory(&m_ConnectedPeer, &g_ConnectEntries[m_LastConnectEntryId].pMacAddress, sizeof(DOT11_MAC_ADDRESS));
 
         // Report link quality
-        WX_RETURN_NTSTATUS_IF_NOT_NT_SUCCESS_MSG(WifiIhvSendLinkStateIndication(Device, pWdiHeader, 1),
+        WX_RETURN_NTSTATUS_IF_NOT_NT_SUCCESS_MSG(WifiIhvSendLinkStateIndication(pWdiHeader, 1),
             "Failed WifiIhvSendLinkStateIndication");
 
-        WifiIhvSendIndicationToOs(
-            Device, pWdiHeader, WDI_INDICATION_CONNECT_COMPLETE, m_LastConnectTransactionId, nullptr, 0);
+        WifiIhvSendM4IndicationToOs(
+            m_Device, WDI_INDICATION_CONNECT_COMPLETE, pWdiHeader, STATUS_SUCCESS);
 #ifdef WIFI_IHV_HANDSHAKE
         //
         // Receive M1 frame of 4-way handshake
@@ -857,10 +855,10 @@ NTSTATUS WifiHAL::WifiIhvSetSaeAuthParams(WDFDEVICE Device, const WDI_SET_SAE_AU
     RtlCopyMemory(SAEAuthParamsNeeded.BssId.Address, &g_ConnectEntries[m_LastConnectEntryId].pMacAddress, sizeof(DOT11_MAC_ADDRESS));
 
     auto ndisStatus =
-        GenerateWdiIndicationSaeAuthParamsNeeded(&SAEAuthParamsNeeded, 0, Context, &cbOutput, &pOutput);
+        GenerateWdiIndicationSaeAuthParamsNeeded(&SAEAuthParamsNeeded, 0, m_TlvContext, &cbOutput, &pOutput);
     if (ndisStatus == NDIS_STATUS_SUCCESS)
     {
-        WifiIhvSendIndicationToOs(Device, pWdiHeader, WDI_INDICATION_SAE_AUTH_PARAMS_NEEDED, 0, pOutput, cbOutput);
+        WifiIhvSendUnsolicitedIndicationToOs(m_Device, pWdiHeader, WDI_INDICATION_SAE_AUTH_PARAMS_NEEDED, pOutput, cbOutput);
 
         FreeGenerated(pOutput);
     }
@@ -873,7 +871,7 @@ NTSTATUS WifiHAL::WifiIhvSetSaeAuthParams(WDFDEVICE Device, const WDI_SET_SAE_AU
 }
 
 _Use_decl_annotations_
-void WifiHAL::WifiIhvPerformDisassociation(_In_ WDFDEVICE Device, _In_ PWDI_MESSAGE_HEADER pWdiHeader, _In_ WDI_ASSOC_STATUS DisassocStatus)
+void WifiHAL::WifiIhvPerformDisassociation(_In_ PWDI_MESSAGE_HEADER pWdiHeader, _In_ WDI_ASSOC_STATUS DisassocStatus)
 {
     UCHAR s_TLV_Disassociation[] =
     {
@@ -895,13 +893,17 @@ void WifiHAL::WifiIhvPerformDisassociation(_In_ WDFDEVICE Device, _In_ PWDI_MESS
     RtlCopyMemory(&s_TLV_Disassociation[10], &DisassocStatus, sizeof(ULONG));
 
     // Send the disassociation indication
-    WifiIhvSendIndicationToOs(Device, pWdiHeader, WDI_INDICATION_DISASSOCIATION, 0, s_TLV_Disassociation, sizeof(s_TLV_Disassociation));
+    WifiIhvSendUnsolicitedIndicationToOs(m_Device, pWdiHeader, WDI_INDICATION_DISASSOCIATION, s_TLV_Disassociation, sizeof(s_TLV_Disassociation));
 
     m_LastConnectEntryId = 0; // Disconnected State
 
 #ifdef WIFI_IHV_NETV
     DeleteDatapathPeer(0x33);
 #endif
+}
 
-    WifiIhvSendIndicationToOs(Device, pWdiHeader, WDI_INDICATION_DISCONNECT_COMPLETE, pWdiHeader->TransactionId, nullptr, 0);
-}
\ No newline at end of file
+WifiHAL::WifiHAL(WDFDEVICE Device)
+{
+    m_Device = Device;
+    m_TlvContext = &WifiGetIhvDeviceContext(Device)->TlvContext;
+}
diff --git a/network/wlan/WIFICX/drivercode/wifihal.h b/network/wlan/WIFICX/drivercode/wifihal.h
index 5f1ec299..d2d770d9 100644
--- a/network/wlan/WIFICX/drivercode/wifihal.h
+++ b/network/wlan/WIFICX/drivercode/wifihal.h
@@ -5,17 +5,21 @@
 class WifiHAL
 {
 public:
-    NTSTATUS WifiIhvSetDeviceCapabilities(_In_ WDFDEVICE Device);
-    NTSTATUS WifiIhvReset(_In_ WDFDEVICE Device,_In_ const WDI_TASK_DOT11_RESET_PARAMETERS& ResetParameters, _In_ const PWDI_MESSAGE_HEADER pWdiHeader, _In_ PCTLV_CONTEXT Context);
-    NTSTATUS WifiIhvSetRadioState(_In_ WDFDEVICE Device, _In_ const WDI_SET_RADIO_STATE_PARAMETERS& RadioState, _In_ const PWDI_MESSAGE_HEADER pWdiHeader, _In_ PCTLV_CONTEXT Context);
-    NTSTATUS WifiIhvScan(_In_ WDFDEVICE Device, _In_ WDI_SCAN_PARAMETERS& ScanParameters, _In_ const PWDI_MESSAGE_HEADER pWdiHeader, _In_ PCTLV_CONTEXT Context);
-    NTSTATUS WifiIhvConnect(_In_ WDFDEVICE Device, _In_ WDI_TASK_CONNECT_PARAMETERS& ConnectParameters, _In_ const PWDI_MESSAGE_HEADER pWdiHeader, _In_ PCTLV_CONTEXT Context);
-    NTSTATUS WifiIhvPerformAssociation(_In_ WDFDEVICE Device, _In_ struct ArrayOfElements<WDI_CONNECT_BSS_ENTRY_CONTAINER>* pPreferredBSSEntryList, _In_ struct ArrayOfElements<WDI_AUTH_ALGORITHM>* pAuthenticationAlgorithms, _In_ PWDI_MESSAGE_HEADER pWdiHeader);
-    NTSTATUS WifiIhvSetSaeAuthParams(_In_ WDFDEVICE Device, _In_ const WDI_SET_SAE_AUTH_PARAMS_COMMAND& setSAEAuthParams, _In_ const PWDI_MESSAGE_HEADER pWdiHeader, _In_ PCTLV_CONTEXT Context);
-    void WifiIhvPerformDisassociation(_In_ WDFDEVICE Device, _In_ PWDI_MESSAGE_HEADER pWdiHeader, _In_ WDI_ASSOC_STATUS DisassocStatus);
+    WifiHAL(_In_ WDFDEVICE Device);
+    NTSTATUS WifiIhvSetDeviceCapabilities();
+    NTSTATUS WifiIhvReset(_In_ const WDI_TASK_DOT11_RESET_PARAMETERS& ResetParameters, _In_ const PWDI_MESSAGE_HEADER pWdiHeader);
+    NTSTATUS WifiIhvSetRadioState(_In_ const WDI_SET_RADIO_STATE_PARAMETERS& RadioState, _In_ const PWDI_MESSAGE_HEADER pWdiHeader);
+    NTSTATUS WifiIhvScan(_In_ WDI_SCAN_PARAMETERS& ScanParameters, _In_ const PWDI_MESSAGE_HEADER pWdiHeader);
+    NTSTATUS WifiIhvConnect(_In_ WDI_TASK_CONNECT_PARAMETERS& ConnectParameters, _In_ const PWDI_MESSAGE_HEADER pWdiHeader);
+    NTSTATUS WifiIhvPerformAssociation(_In_ struct ArrayOfElements<WDI_CONNECT_BSS_ENTRY_CONTAINER>* pPreferredBSSEntryList, _In_ struct ArrayOfElements<WDI_AUTH_ALGORITHM>* pAuthenticationAlgorithms, _In_ PWDI_MESSAGE_HEADER pWdiHeader);
+    NTSTATUS WifiIhvSetSaeAuthParams(_In_ const WDI_SET_SAE_AUTH_PARAMS_COMMAND& setSAEAuthParams, _In_ const PWDI_MESSAGE_HEADER pWdiHeader);
+    void WifiIhvPerformDisassociation(_In_ PWDI_MESSAGE_HEADER pWdiHeader, _In_ WDI_ASSOC_STATUS DisassocStatus);
 private:
-    NTSTATUS WifiIhvSendLinkStateIndication(_In_ WDFDEVICE Device, _In_ PWDI_MESSAGE_HEADER pWdiHeader, ULONG numLinks);
-
+    NTSTATUS WifiIhvSendLinkStateIndication(_In_ PWDI_MESSAGE_HEADER pWdiHeader, ULONG numLinks);
+    
+    WDFDEVICE m_Device{};
+    PCTLV_CONTEXT m_TlvContext{};
+    
     UCHAR m_CurrentRadioState{};
     UINT32 m_LastConnectEntryId{};
     UINT32 m_LastConnectTransactionId{};
diff --git a/network/wlan/WIFICX/drivercode/wifirequest.cpp b/network/wlan/WIFICX/drivercode/wifirequest.cpp
index 1fefe040..525fa14c 100644
--- a/network/wlan/WIFICX/drivercode/wifirequest.cpp
+++ b/network/wlan/WIFICX/drivercode/wifirequest.cpp
@@ -1,31 +1,43 @@
 // Copyright (c) Microsoft Corporation.  All rights reserved.
 
 #include "precomp.h"
-#include "wifiHAL.h"
+#include "wifitransition.h"
 #include "wifirequest.h"
 #include "wifirequest.tmh"
 
 _Use_decl_annotations_
 void EvtWifiDeviceSendCommand(WDFDEVICE Device, WIFIREQUEST SendRequest)
 {
-    NTSTATUS Status = STATUS_SUCCESS;
     UINT InBufferLen = 0;
     UINT OutBufferLen = 0;
-    UINT BytesWritten = 0;
     void* Buffer = WifiRequestGetInOutBuffer(SendRequest, &InBufferLen, &OutBufferLen);
     UINT16 MessageId = WifiRequestGetMessageId(SendRequest);
 
-    Status = ProcessWifiRequest(Device, MessageId, Buffer, InBufferLen, OutBufferLen, &BytesWritten);
+    TransitionContext tctx{
+        Device,
+        WifiGetIhvDeviceContext(Device),
+        SendRequest,
+        static_cast<PWDI_MESSAGE_HEADER>(Buffer),
+        Buffer,
+        InBufferLen,
+        OutBufferLen
+    };
 
-    if (Status != STATUS_PENDING)
+    if(!NT_SUCCESS(RunTransitionByMessage(tctx, MessageId)))
     {
-        WifiRequestComplete(SendRequest, Status, BytesWritten);
+        WFCError("RunTransitionByMessage failed for MessageId: 0x%04X", MessageId);
+        return;
     }
-    return;
 }
 
-_Use_decl_annotations_
-void WifiIhvSendIndicationToOs(WDFDEVICE Device, PWDI_MESSAGE_HEADER pWdiHeader, UINT16 MessageId, UINT32 TransactionId, PUCHAR pTlvData, UINT32 TlvDataSize)
+void WifiIhvSendIndicationToOs(
+    _In_ WDFDEVICE Device, 
+    _In_ const PWDI_MESSAGE_HEADER pOriginalWdiHeader, 
+    _In_ UINT16 WifiRequestMessageId, 
+    _In_ UINT32 WifiRequestTransactionId, 
+    _In_ NTSTATUS WifiRquestM4Status,
+    _In_ PUCHAR pTlvData, 
+    _In_ UINT32 TlvDataSize)
 {
     WDFMEMORY data = WDF_NO_HANDLE;
     PUCHAR pIndicationBuffer = nullptr;
@@ -45,8 +57,9 @@ void WifiIhvSendIndicationToOs(WDFDEVICE Device, PWDI_MESSAGE_HEADER pWdiHeader,
     RtlZeroMemory(pIndicationBuffer, indicationSize);
     pIndicationHeader = reinterpret_cast<PWDI_MESSAGE_HEADER>(pIndicationBuffer);
 
-    RtlCopyMemory(pIndicationHeader, pWdiHeader, sizeof(WDI_MESSAGE_HEADER));
-    pIndicationHeader->TransactionId = TransactionId;
+    RtlCopyMemory(pIndicationHeader, pOriginalWdiHeader, sizeof(WDI_MESSAGE_HEADER));
+    pIndicationHeader->TransactionId = WifiRequestTransactionId;
+    pIndicationHeader->Status = Wifi::ConvertNDISSTATUSToNTSTATUS(WifiRquestM4Status);
 
     if (TlvDataSize > 0)
     {
@@ -54,193 +67,31 @@ void WifiIhvSendIndicationToOs(WDFDEVICE Device, PWDI_MESSAGE_HEADER pWdiHeader,
     }
 
     // Send the indication up to WifiCx
-    WifiDeviceReceiveIndication(Device, MessageId, data);
+    WifiDeviceReceiveIndication(Device, WifiRequestMessageId, data);
 
     // Don't need to keep this around
     WdfObjectDelete(data);
 }
 
 _Use_decl_annotations_
-static NTSTATUS ProcessWifiRequest(WDFDEVICE Device, UINT16 MessageId, void* Buffer, UINT InBufferLen, UINT OutBufferLen, UINT* pBytesWritten)
+void WifiIhvSendUnsolicitedIndicationToOs(WDFDEVICE Device, PWDI_MESSAGE_HEADER pWdiHeader, UINT16 MessageId, PUCHAR pTlvData, UINT32 TlvDataSize)
 {
-    NTSTATUS ntStatus = STATUS_SUCCESS;
-    PWIFI_IHV_DEVICE_CONTEXT deviceContext = WifiGetIhvDeviceContext(Device);
-    PWDI_MESSAGE_HEADER pWdiHeader = static_cast<PWDI_MESSAGE_HEADER>(Buffer);
-    UINT BytesWritten = sizeof(WDI_MESSAGE_HEADER);
-
-    if (InBufferLen < sizeof(WDI_MESSAGE_HEADER))
-    {
-        return STATUS_INVALID_PARAMETER;
-    }
-
-    if (OutBufferLen < sizeof(WDI_MESSAGE_HEADER))
-    {
-        return STATUS_INVALID_PARAMETER;
-    }
-
-    auto const buffer = static_cast<UCHAR*>(Buffer);
-
-    // Check for integer overflow
-    if (!(buffer + InBufferLen > buffer))
-    {
-        return STATUS_INVALID_PARAMETER;
-    }
-
-    switch (MessageId)
-    {
-    case WDI_TASK_SET_RADIO_STATE:
-    {
-        WDI_SET_RADIO_STATE_PARAMETERS RadioStateParams = {};
-        ULONG TLVStreamLength = static_cast<ULONG>(InBufferLen - sizeof(WDI_MESSAGE_HEADER));
-        auto TLVByteStream = buffer +sizeof(WDI_MESSAGE_HEADER);
-
-        DumpMessageTlvByteStream(
-            WDI_TASK_SET_RADIO_STATE, TRUE, deviceContext->TlvContext.PeerVersion, TLVStreamLength, TLVByteStream, 0, nullptr);
-
-        WX_RETURN_NTSTATUS_IF_NOT_NT_SUCCESS_MSG(
-            Wifi::ConvertNDISSTATUSoNTSTATUS(
-                ParseWdiTaskSetRadioState(
-                    TLVStreamLength,
-                    TLVByteStream,
-                    &deviceContext->TlvContext,
-                    &RadioStateParams)),
-            "Set Radio State parsing failed");
-
-        // No WX_RETURN here since we need to cleanup.
-        ntStatus = deviceContext->wifiHAL->WifiIhvSetRadioState(
-            Device, RadioStateParams, pWdiHeader, &deviceContext->TlvContext);
-        if(!NT_SUCCESS(ntStatus)) 
-        {
-            WFCError("WifiIhvSetRadioState failed: %!STATUS!\n", ntStatus);
-        }
-
-        CleanupParsedWdiTaskSetRadioState(&RadioStateParams);
-    }
-    break;
-    case WDI_TASK_SCAN:
-    {
-        ULONG scanDataLength = static_cast<ULONG>(InBufferLen - sizeof(WDI_MESSAGE_HEADER));
-        auto pScanData = buffer + sizeof(WDI_MESSAGE_HEADER);
-
-        WDI_SCAN_PARAMETERS scanParams = {};
-
-        WX_RETURN_NTSTATUS_IF_NOT_NT_SUCCESS_MSG(
-            Wifi::ConvertNDISSTATUSoNTSTATUS(
-                ParseWdiTaskScan(
-                    scanDataLength,
-                    pScanData,
-                    &deviceContext->TlvContext, 
-                    &scanParams)),
-                "Scan parameters parsing failed");
-
-        // No WX_RETURN here since we need to cleanup.
-        ntStatus = deviceContext->wifiHAL->WifiIhvScan(
-            Device, scanParams, pWdiHeader, &deviceContext->TlvContext);
-        if (!NT_SUCCESS(ntStatus))
-        {
-            WFCError("WifiIhvSetRadioState failed: %!STATUS!\n", ntStatus);
-        }
-
-        CleanupParsedWdiTaskScan(&scanParams);
-    }
-    break;
-
-    case WDI_TASK_DOT11_RESET:
-    {
-        WDI_TASK_DOT11_RESET_PARAMETERS resetParameters;
-
-        ULONG resetDataLength = static_cast<ULONG>(InBufferLen - sizeof(WDI_MESSAGE_HEADER));
-        auto pResetData = buffer + sizeof(WDI_MESSAGE_HEADER);
-
-        WX_RETURN_NTSTATUS_IF_NOT_NT_SUCCESS_MSG(
-            Wifi::ConvertNDISSTATUSoNTSTATUS(ParseWdiTaskDot11Reset(resetDataLength, pResetData, &deviceContext->TlvContext, &resetParameters)),
-                "DOT11 Reset parameters parsing failed");
-
-        // No WX_RETURN here since we need to cleanup.
-        ntStatus = deviceContext->wifiHAL->WifiIhvReset(
-            Device, resetParameters, pWdiHeader, &deviceContext->TlvContext);
-        if (!NT_SUCCESS(ntStatus))
-        {
-            WFCError("WifiIhvReset failed: %!STATUS!\n", ntStatus);
-        }
-
-        CleanupParsedWdiTaskDot11Reset(&resetParameters);
-    }
-    break;
-
-    case WDI_TASK_CONNECT:
-    {
-        ULONG connectDataLength = static_cast<ULONG>(InBufferLen - sizeof(WDI_MESSAGE_HEADER));
-        auto pConnectData = buffer + sizeof(WDI_MESSAGE_HEADER);
-
-        DumpMessageTlvByteStream(WDI_TASK_CONNECT, TRUE, deviceContext->TlvContext.PeerVersion, connectDataLength, pConnectData, 0, nullptr);
-
-        WDI_TASK_CONNECT_PARAMETERS connectParams{};
-
-
-        // We search for the BSSID pattern to figure out what AP we are trying to connect to
-        WX_RETURN_NTSTATUS_IF_NOT_NT_SUCCESS_MSG(
-            Wifi::ConvertNDISSTATUSoNTSTATUS(ParseWdiTaskConnect(connectDataLength, pConnectData, &deviceContext->TlvContext, &connectParams)),
-                "Connect parameters parsing failed");
-        
-        // No WX_RETURN here since we need to cleanup.
-        ntStatus = deviceContext->wifiHAL->WifiIhvConnect(
-            Device, connectParams, pWdiHeader, &deviceContext->TlvContext);
-        if (!NT_SUCCESS(ntStatus))
-        {
-            WFCError("WifiIhvConnect failed: %!STATUS!\n", ntStatus);
-        }
-
-        CleanupParsedWdiTaskConnect(&connectParams);
-    }
-    break;
-
-    case WDI_TASK_DISCONNECT:
-    {
-        deviceContext->wifiHAL->WifiIhvPerformDisassociation(Device, pWdiHeader, WDI_ASSOC_STATUS_DISASSOCIATED_BY_HOST);
-    }
-    break;
-
-    case WDI_SET_PRIVACY_EXEMPTION_LIST:
-    case WDI_SET_DEFAULT_KEY_ID:
-    case WDI_SET_ADD_CIPHER_KEYS:
-    case WDI_SET_DELETE_CIPHER_KEYS:
-    case WDI_SET_RECEIVE_PACKET_FILTER:
-    case WDI_SET_CONNECTION_QUALITY:
-    case WDI_SET_ADAPTER_CONFIGURATION:
-        // Doing Nothing special
-        break;
-
-    case WDI_SET_SAE_AUTH_PARAMS:
-    {
-        WDI_SET_SAE_AUTH_PARAMS_COMMAND setSAEAuthParams;
-        
-        WX_RETURN_NTSTATUS_IF_NOT_NT_SUCCESS_MSG(
-            Wifi::ConvertNDISSTATUSoNTSTATUS(ParseWdiSetSaeAuthParams(
-            InBufferLen - sizeof(WDI_MESSAGE_HEADER), (PUINT8)buffer + sizeof(WDI_MESSAGE_HEADER), &deviceContext->TlvContext, &setSAEAuthParams)),
-            "Failed to parse WDI_SET_SAE_AUTH_PARAMS");
-
-        // No WX_RETURN here since we need to cleanup.
-        ntStatus = deviceContext->wifiHAL->WifiIhvSetSaeAuthParams(
-            Device, setSAEAuthParams, pWdiHeader, &deviceContext->TlvContext);
-        if (!NT_SUCCESS(ntStatus))
-        {
-            WFCError("WifiIhvSetSaeAuthParams failed: %!STATUS!\n", ntStatus);
-        }
-
-        CleanupParsedWdiSetSaeAuthParams(&setSAEAuthParams);
-    }
-    break;
-
-    default:
-        ntStatus = STATUS_NOT_SUPPORTED;
-        break;
-    }
-
-    // Set the output length
-    if (pBytesWritten)
+    WifiIhvSendIndicationToOs(Device, pWdiHeader, MessageId, 0, STATUS_SUCCESS, pTlvData, TlvDataSize); //TransactionId required to be 0 for unsolicited indications.
+}
+_Use_decl_annotations_
+void WifiIhvNotifyM3Completion(WIFIREQUEST Request, NTSTATUS WifiRequestM3Status, UINT BytesWritten)
+{
+    if(!NT_SUCCESS(WifiRequestM3Status))
     {
-        *pBytesWritten = BytesWritten;
-    }
-    return ntStatus;
+        WFCError("WifiRequest M3 failed: %!STATUS!, BytesWritten: %d\n", WifiRequestM3Status, BytesWritten);
+    }   
+    // Report the M3 status back to OS
+    // OS expects M3 before the M4
+    WifiRequestComplete(Request, WifiRequestM3Status, BytesWritten);
 }
+
+_Use_decl_annotations_
+void WifiIhvSendM4IndicationToOs(WDFDEVICE Device, UINT16 WifiRequestMessageId, const PWDI_MESSAGE_HEADER pWdiHeader, NTSTATUS WifiRequestM4Status)
+{
+    WifiIhvSendIndicationToOs(Device, pWdiHeader, WifiRequestMessageId, pWdiHeader->TransactionId, WifiRequestM4Status, nullptr, 0);
+}
\ No newline at end of file
diff --git a/network/wlan/WIFICX/drivercode/wifirequest.h b/network/wlan/WIFICX/drivercode/wifirequest.h
index 601fec7d..c3317eb3 100644
--- a/network/wlan/WIFICX/drivercode/wifirequest.h
+++ b/network/wlan/WIFICX/drivercode/wifirequest.h
@@ -7,10 +7,20 @@ EVT_WIFI_DEVICE_SEND_COMMAND EvtWifiDeviceSendCommand;
 static NTSTATUS ProcessWifiRequest(
     _In_ WDFDEVICE Device, _In_ UINT16 MessageId, _In_ void* Buffer, _In_ UINT InBufferLen, _In_ UINT OutBufferLen, _In_ UINT* pBytesWritten);
 
-void WifiIhvSendIndicationToOs(
+void WifiIhvSendUnsolicitedIndicationToOs(
     _In_ WDFDEVICE Device,
     _In_ PWDI_MESSAGE_HEADER pWdiHeader,
     _In_ UINT16 MessageId,
-    _In_ UINT32 TransactionId,
     _In_opt_bytecount_(TlvDataSize) PUCHAR pTlvData,
     _In_ UINT32 TlvDataSize);
+
+void WifiIhvNotifyM3Completion(
+    _In_ WIFIREQUEST Request,
+    _In_ NTSTATUS WifiRequestM3Status,
+    _In_ UINT BytesWritten);
+
+void WifiIhvSendM4IndicationToOs(
+    _In_ WDFDEVICE Device,
+    _In_ UINT16 WifiRequestMessageId,
+    _In_ const PWDI_MESSAGE_HEADER pOriginalWdiHeader,
+    _In_ NTSTATUS WifiRequestM4Status);
\ No newline at end of file
diff --git a/network/wlan/WIFICX/drivercode/wifitransition.cpp b/network/wlan/WIFICX/drivercode/wifitransition.cpp
new file mode 100644
index 00000000..4e55f39c
--- /dev/null
+++ b/network/wlan/WIFICX/drivercode/wifitransition.cpp
@@ -0,0 +1,25 @@
+#include "wifitransition.h"
+
+// Runtime dispatcher switches on MessageId and invokes the matching compile-time runner.
+NTSTATUS RunTransitionByMessage(TransitionContext& ctx, UINT16 messageId)
+{
+    switch (messageId)
+    {
+    case WDI_TASK_SET_RADIO_STATE:
+        return RunTransition<WDI_TASK_SET_RADIO_STATE>(ctx);
+    case WDI_TASK_SCAN:
+        return RunTransition<WDI_TASK_SCAN>(ctx);
+    case WDI_TASK_DOT11_RESET:
+        return RunTransition<WDI_TASK_DOT11_RESET>(ctx);
+    case WDI_TASK_CONNECT:
+        return RunTransition<WDI_TASK_CONNECT>(ctx);
+    case WDI_TASK_DISCONNECT:
+        return RunTransition<WDI_TASK_DISCONNECT>(ctx);
+    case WDI_SET_SAE_AUTH_PARAMS:
+        return RunTransition<WDI_SET_SAE_AUTH_PARAMS>(ctx);
+    default:
+        UINT bytesWritten = sizeof(WDI_MESSAGE_HEADER);
+        WifiRequestComplete(ctx.WifiRequest, STATUS_NOT_SUPPORTED, bytesWritten);
+        return STATUS_NOT_SUPPORTED;
+    }
+}
\ No newline at end of file
diff --git a/network/wlan/WIFICX/drivercode/wifitransition.h b/network/wlan/WIFICX/drivercode/wifitransition.h
new file mode 100644
index 00000000..d57e73cb
--- /dev/null
+++ b/network/wlan/WIFICX/drivercode/wifitransition.h
@@ -0,0 +1,332 @@
+#pragma once
+#include "precomp.h"
+#include "wifiHAL.h"
+#include "wifirequest.h" // WifiIhvSendM4IndicationToOs, DumpMessageTlvByteStream, Wifi::ConvertNDISSTATUSToNTSTATUS
+
+// Generic execution context passed between steps and parse/handle
+struct TransitionContext
+{
+    WDFDEVICE Device;
+    PWIFI_IHV_DEVICE_CONTEXT DevCtx;
+    WIFIREQUEST WifiRequest;
+    PWDI_MESSAGE_HEADER Header;
+    void* RawBuffer;
+    UINT InLen;
+    UINT OutLen;
+};
+
+// Common TLV parsing helper: computes TLV span, optional dump, calls a parser, converts status.
+template<typename Param, typename ParserFn>
+NTSTATUS ParseTlvCommon(TransitionContext& ctx,
+    UINT16 messageId,
+    ParserFn parser,
+    Param& outParams,
+    bool dumpStream = true)
+{
+    if (ctx.InLen < sizeof(WDI_MESSAGE_HEADER))
+    {
+        return STATUS_INVALID_PARAMETER;
+    }
+
+    auto* tlvBytes = static_cast<UCHAR*>(ctx.RawBuffer) + sizeof(WDI_MESSAGE_HEADER);
+    auto tlvLen = static_cast<ULONG>(ctx.InLen - sizeof(WDI_MESSAGE_HEADER));
+
+    if (dumpStream)
+    {
+        DumpMessageTlvByteStream(
+            messageId,
+            TRUE,
+            ctx.DevCtx->TlvContext.PeerVersion,
+            tlvLen,
+            tlvBytes,
+            0,
+            nullptr);
+    }
+
+    auto ndisStatus = parser(tlvLen, tlvBytes, &ctx.DevCtx->TlvContext, &outParams);
+    return Wifi::ConvertNDISSTATUSToNTSTATUS(ndisStatus);
+}
+
+// Execute two step callables.
+// StepM3Fn signature: NTSTATUS (TransitionContext&, Param&, UINT& bytesWritten)
+// StepM4Fn signature: NTSTATUS (TransitionContext&, Param&)
+// Always calls WifiRequestComplete after StepM3 with the bytesWritten produced by StepM3.
+// Skips StepM4 if StepM3 failed
+template<typename Param, typename StepM3Fn, typename StepM4Fn>
+NTSTATUS ExecuteSteps(TransitionContext& ctx, Param& p, StepM3Fn stepM3, StepM4Fn stepM4)
+{
+    UINT bytesWritten = sizeof(WDI_MESSAGE_HEADER); // default minimum
+    NTSTATUS m3Status = stepM3(ctx, p, bytesWritten);
+    // Report the M3 status back to OS
+    // OS expects M3 before the M4
+    WifiIhvNotifyM3Completion(ctx.WifiRequest, m3Status, bytesWritten);
+    if (!NT_SUCCESS(m3Status))
+    {
+        return m3Status;
+    }
+    return stepM4(ctx, p);
+}
+// Primary traits template (specialize per MessageId)
+template<UINT16 MsgId>
+struct TransitionTraits;
+
+// -------- WDI_TASK_SET_RADIO_STATE --------
+template<>
+struct TransitionTraits<WDI_TASK_SET_RADIO_STATE>
+{
+    using ParamType = WDI_SET_RADIO_STATE_PARAMETERS;
+    enum : UINT16 { CompleteIndication = WDI_INDICATION_SET_RADIO_STATE_COMPLETE };
+
+    NTSTATUS Parse(TransitionContext& ctx, ParamType& p)
+    {
+        return ParseTlvCommon(ctx,
+            WDI_TASK_SET_RADIO_STATE,
+            [](ULONG len, UINT8* bytes, PCTLV_CONTEXT tlvCtx, ParamType* out)
+            {
+                return ParseWdiTaskSetRadioState(len, bytes, tlvCtx, out);
+            },
+            p,
+            true);
+    }
+    void Cleanup(ParamType& p) { CleanupParsedWdiTaskSetRadioState(&p); }
+
+    NTSTATUS Handle(TransitionContext& ctx, ParamType& p)
+    {
+        return ExecuteSteps(ctx, p,
+            // STEPM3
+            [](TransitionContext&, ParamType&, UINT& bytesWritten) 
+            { 
+                bytesWritten = sizeof(WDI_MESSAGE_HEADER);
+                return STATUS_SUCCESS; 
+            },
+            // STEPM4
+            [](TransitionContext& c, ParamType& par)
+            { 
+                return c.DevCtx->wifiHAL->WifiIhvSetRadioState(par, c.Header);
+            }
+            );
+    }
+    bool ShouldSendComplete(NTSTATUS) const { return true; }
+};
+
+// -------- WDI_TASK_SCAN --------
+template<>
+struct TransitionTraits<WDI_TASK_SCAN>
+{
+    using ParamType = WDI_SCAN_PARAMETERS;
+    enum : UINT16 { CompleteIndication = WDI_INDICATION_SCAN_COMPLETE };
+
+    NTSTATUS Parse(TransitionContext& ctx, ParamType& p)
+    {
+        return ParseTlvCommon(ctx,
+            WDI_TASK_SCAN,
+            [](ULONG len, UINT8* bytes, PCTLV_CONTEXT tlvCtx, ParamType* out)
+            {
+                return ParseWdiTaskScan(len, bytes, tlvCtx, out);
+            },
+            p,
+            true);
+    }
+    void Cleanup(ParamType& p){ CleanupParsedWdiTaskScan(&p); }
+
+    NTSTATUS Handle(TransitionContext& ctx, ParamType& p)
+    {
+        return ExecuteSteps(ctx, p,
+            // STEPM3
+            [](TransitionContext&, ParamType&, UINT& bytesWritten)
+            {
+                bytesWritten = sizeof(WDI_MESSAGE_HEADER);
+                return STATUS_SUCCESS;
+            },
+            // STEPM4
+            [](TransitionContext& c, ParamType& par)
+            {
+                return c.DevCtx->wifiHAL->WifiIhvScan(par, c.Header);
+            }
+        );
+    }
+    bool ShouldSendComplete(NTSTATUS) const { return true; }
+};
+
+// -------- WDI_TASK_DOT11_RESET --------
+template<>
+struct TransitionTraits<WDI_TASK_DOT11_RESET>
+{
+    using ParamType = WDI_TASK_DOT11_RESET_PARAMETERS;
+    enum : UINT16 { CompleteIndication = WDI_INDICATION_DOT11_RESET_COMPLETE };
+
+    NTSTATUS Parse(TransitionContext& ctx, ParamType& p)
+    {
+        return ParseTlvCommon(ctx,
+            WDI_TASK_DOT11_RESET,
+            [](ULONG len, UINT8* bytes, PCTLV_CONTEXT tlvCtx, ParamType* out)
+            {
+                return ParseWdiTaskDot11Reset(len, bytes, tlvCtx, out);
+            },
+            p,
+            false);
+    }
+    void Cleanup(ParamType& p){ CleanupParsedWdiTaskDot11Reset(&p); }
+
+    NTSTATUS Handle(TransitionContext& ctx, ParamType& p)
+    {
+        return ExecuteSteps(ctx, p,
+            // STEPM3
+            [](TransitionContext&, ParamType&, UINT& bytesWritten)
+            {
+                bytesWritten = sizeof(WDI_MESSAGE_HEADER);
+                return STATUS_SUCCESS;
+            },
+            // STEPM4
+            [](TransitionContext& c, ParamType& par)
+            {
+                return c.DevCtx->wifiHAL->WifiIhvReset(par, c.Header);
+            });
+    }
+    bool ShouldSendComplete(NTSTATUS) const { return true; }
+};
+
+// -------- WDI_TASK_CONNECT --------
+template<>
+struct TransitionTraits<WDI_TASK_CONNECT>
+{
+    using ParamType = WDI_TASK_CONNECT_PARAMETERS;
+    enum : UINT16 { CompleteIndication = WDI_INDICATION_CONNECT_COMPLETE };
+
+    NTSTATUS Parse(TransitionContext& ctx, ParamType& p)
+    {
+        return ParseTlvCommon(ctx,
+            WDI_TASK_CONNECT,
+            [](ULONG len, UINT8* bytes, PCTLV_CONTEXT tlvCtx, ParamType* out)
+            {
+                return ParseWdiTaskConnect(len, bytes, tlvCtx, out);
+            },
+            p,
+            true);
+    }
+    void Cleanup(ParamType& p){ CleanupParsedWdiTaskConnect(&p); }
+
+    NTSTATUS Handle(TransitionContext& ctx, ParamType& p)
+    {
+        return ExecuteSteps(ctx, p,
+            // STEPM3
+            [](TransitionContext&, ParamType&, UINT& bytesWritten)
+            {
+                bytesWritten = sizeof(WDI_MESSAGE_HEADER);
+                return STATUS_SUCCESS;
+            },
+            // STEPM4
+            [](TransitionContext& c, ParamType& par)
+            {
+                return c.DevCtx->wifiHAL->WifiIhvConnect(par, c.Header);
+            });
+    }
+    bool ShouldSendComplete(NTSTATUS s) const { return s != STATUS_PENDING; }
+};
+
+// -------- WDI_TASK_DISCONNECT --------
+template<>
+struct TransitionTraits<WDI_TASK_DISCONNECT>
+{
+    struct ParamType {};
+    enum : UINT16 { CompleteIndication = WDI_INDICATION_DISCONNECT_COMPLETE };
+
+    NTSTATUS Parse(TransitionContext&, ParamType&){ return STATUS_SUCCESS; }
+    void Cleanup(ParamType&){}
+
+    NTSTATUS Handle(TransitionContext& ctx, ParamType& p)
+    {
+        return ExecuteSteps(ctx, p,
+            // STEPM3
+            [](TransitionContext&, ParamType&, UINT& bytesWritten)
+            {
+                bytesWritten = sizeof(WDI_MESSAGE_HEADER);
+                return STATUS_SUCCESS;
+            },
+            // STEPM4
+            [](TransitionContext& c, ParamType&)
+            {
+                c.DevCtx->wifiHAL->WifiIhvPerformDisassociation(c.Header, WDI_ASSOC_STATUS_DISASSOCIATED_BY_HOST);
+                return STATUS_SUCCESS;
+            });
+    }
+    bool ShouldSendComplete(NTSTATUS) const { return true; }
+};
+
+// -------- WDI_SET_SAE_AUTH_PARAMS --------
+template<>
+struct TransitionTraits<WDI_SET_SAE_AUTH_PARAMS>
+{
+    using ParamType = WDI_SET_SAE_AUTH_PARAMS_COMMAND;
+    enum : UINT16 { CompleteIndication = 0 };
+
+    NTSTATUS Parse(TransitionContext& ctx, ParamType& p)
+    {
+        return ParseTlvCommon(ctx,
+            WDI_SET_SAE_AUTH_PARAMS,
+            [](ULONG len, UINT8* bytes, PCTLV_CONTEXT tlvCtx, ParamType* out)
+            {
+                return ParseWdiSetSaeAuthParams(len, bytes, tlvCtx, out);
+            },
+            p,
+            /*dumpStream*/ false);
+    }
+
+    void Cleanup(ParamType& p)
+    {
+        CleanupParsedWdiSetSaeAuthParams(&p);
+    }
+
+    NTSTATUS Handle(TransitionContext& ctx, ParamType& p)
+    {
+        return ExecuteSteps(ctx, p,
+            // Step M3: invoke HAL handler, report bytes written (just header)
+            [](TransitionContext& c, ParamType& par, UINT& bytesWritten)
+            {
+                bytesWritten = sizeof(WDI_MESSAGE_HEADER);
+                return c.DevCtx->wifiHAL->WifiIhvSetSaeAuthParams(par, c.Header);
+            },
+            // Step M4: no-op (no completion indication for this OID)
+            [](TransitionContext&, ParamType&) { return STATUS_SUCCESS; });
+    }
+
+    bool ShouldSendComplete(NTSTATUS) const
+    {
+        // Do not send M4 indication for WDI_SET_SAE_AUTH_PARAMS (original code omitted it).
+        return false;
+    }
+};
+
+// -------- Generic runner (compile-time) --------
+template<UINT16 MsgId>
+NTSTATUS RunTransition(TransitionContext& ctx)
+{
+    TransitionTraits<MsgId> traits;
+    typename TransitionTraits<MsgId>::ParamType params{};
+    NTSTATUS parseStatus = traits.Parse(ctx, params);
+    if (!NT_SUCCESS(parseStatus))
+    {
+        traits.Cleanup(params);
+        // Report Failed M3 to OS
+        // Note: No M4 indication on parse failure
+        WifiIhvNotifyM3Completion(ctx.WifiRequest, parseStatus, 0);
+        return parseStatus;
+    }
+
+    NTSTATUS m4Status = traits.Handle(ctx, params);
+
+    if (traits.ShouldSendComplete(m4Status))
+    {
+        WifiIhvSendM4IndicationToOs(
+            ctx.Device,
+            TransitionTraits<MsgId>::CompleteIndication,
+            ctx.Header,
+            m4Status);
+    }
+
+    traits.Cleanup(params);
+    return m4Status;
+}
+
+// -------- Runtime dispatcher (decl) --------
+NTSTATUS RunTransitionByMessage(TransitionContext& ctx, UINT16 messageId);
\ No newline at end of file
diff --git a/network/wlan/WIFICX/km/wificxsampleclientkm.vcxproj b/network/wlan/WIFICX/km/wificxsampleclientkm.vcxproj
index aea719ed..49510981 100644
--- a/network/wlan/WIFICX/km/wificxsampleclientkm.vcxproj
+++ b/network/wlan/WIFICX/km/wificxsampleclientkm.vcxproj
@@ -193,6 +193,7 @@
     <ClCompile Include="..\drivercode\driver.cpp" />
     <ClCompile Include="..\drivercode\memorymanagement.cpp" />
     <ClCompile Include="..\drivercode\wifihal.cpp" />
+    <ClCompile Include="..\drivercode\wifitransition.cpp" />
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="..\drivercode\adapter.h" />
@@ -205,6 +206,7 @@
     <ClInclude Include="..\drivercode\precomp.h" />
     <ClInclude Include="..\drivercode\trace.h" />
     <ClInclude Include="..\drivercode\wifihal.h" />
+    <ClInclude Include="..\drivercode\wifitransition.h" />
   </ItemGroup>
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.targets" />
   <ImportGroup Label="ExtensionTargets">
diff --git a/network/wlan/WIFICX/km/wificxsampleclientkm.vcxproj.filters b/network/wlan/WIFICX/km/wificxsampleclientkm.vcxproj.filters
index d1dec69c..2930b24d 100644
--- a/network/wlan/WIFICX/km/wificxsampleclientkm.vcxproj.filters
+++ b/network/wlan/WIFICX/km/wificxsampleclientkm.vcxproj.filters
@@ -17,15 +17,6 @@
       <UniqueIdentifier>{8E41214B-6785-4CFE-B992-037D68949A14}</UniqueIdentifier>
       <Extensions>inf;inv;inx;mof;mc;</Extensions>
     </Filter>
-    <Filter Include="Source Files\controlpath">
-      <UniqueIdentifier>{eed063b0-f52a-4d8e-b03f-1f616b73061c}</UniqueIdentifier>
-    </Filter>
-    <Filter Include="Source Files\datapath">
-      <UniqueIdentifier>{56f07a12-b9fb-4ff6-a9b5-e4af54eed417}</UniqueIdentifier>
-    </Filter>
-    <Filter Include="Source Files\firmware">
-      <UniqueIdentifier>{58bc1adc-c940-478a-898b-59d217c27f36}</UniqueIdentifier>
-    </Filter>
   </ItemGroup>
   <ItemGroup>
     <Inf Include="wificxsampleclientkm.inf">
@@ -33,51 +24,60 @@
     </Inf>
   </ItemGroup>
   <ItemGroup>
-    <ClCompile Include="..\common\adapter.cpp">
+    <ClCompile Include="..\drivercode\adapter.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
-    <ClCompile Include="..\common\device.cpp">
+    <ClCompile Include="..\drivercode\wifirequest.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
-    <ClCompile Include="..\common\driver.cpp">
+    <ClCompile Include="..\drivercode\device.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
-    <ClCompile Include="..\common\new.cpp">
+    <ClCompile Include="..\drivercode\driver.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
-    <ClCompile Include="..\common\power.cpp">
+    <ClCompile Include="..\drivercode\memorymanagement.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
-    <ClCompile Include="..\common\controlpath\wifirequest.cpp">
-      <Filter>Source Files\controlpath</Filter>
-    </ClCompile>
-    <ClCompile Include="..\common\datapath\rxqueue.cpp">
-      <Filter>Source Files\datapath</Filter>
+    <ClCompile Include="..\drivercode\wifihal.cpp">
+      <Filter>Source Files</Filter>
     </ClCompile>
-    <ClCompile Include="..\common\datapath\txqueue.cpp">
-      <Filter>Source Files\datapath</Filter>
+    <ClCompile Include="..\drivercode\wifitransition.cpp">
+      <Filter>Source Files</Filter>
     </ClCompile>
   </ItemGroup>
   <ItemGroup>
-    <ClInclude Include="..\common\datapath\netringiterator.h">
+    <ClInclude Include="..\drivercode\adapter.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
+    <ClInclude Include="..\drivercode\device.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
+    <ClInclude Include="..\drivercode\driver.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
+    <ClInclude Include="..\drivercode\sharedtypes.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
+    <ClInclude Include="..\drivercode\umkmfusion.h">
       <Filter>Header Files</Filter>
     </ClInclude>
-    <ClInclude Include="..\common\datapath\rxqueue.h">
+    <ClInclude Include="..\drivercode\wifihaltestdata.h">
       <Filter>Header Files</Filter>
     </ClInclude>
-    <ClInclude Include="..\common\datapath\txqueue.h">
+    <ClInclude Include="..\drivercode\wifirequest.h">
       <Filter>Header Files</Filter>
     </ClInclude>
-    <ClInclude Include="..\common\firmware\testdata.h">
+    <ClInclude Include="..\drivercode\precomp.h">
       <Filter>Header Files</Filter>
     </ClInclude>
-    <ClInclude Include="..\common\controlpath\wifirequest.h">
+    <ClInclude Include="..\drivercode\trace.h">
       <Filter>Header Files</Filter>
     </ClInclude>
-    <ClInclude Include="..\common\precomp.h">
+    <ClInclude Include="..\drivercode\wifihal.h">
       <Filter>Header Files</Filter>
     </ClInclude>
-    <ClInclude Include="..\common\trace.h">
+    <ClInclude Include="..\drivercode\wifitransition.h">
       <Filter>Header Files</Filter>
     </ClInclude>
   </ItemGroup>
diff --git a/network/wlan/WIFICX/um/wificxsampleclientum.vcxproj b/network/wlan/WIFICX/um/wificxsampleclientum.vcxproj
index 70383f39..715859e3 100644
--- a/network/wlan/WIFICX/um/wificxsampleclientum.vcxproj
+++ b/network/wlan/WIFICX/um/wificxsampleclientum.vcxproj
@@ -205,6 +205,7 @@
     <ClInclude Include="..\drivercode\wifihal.h" />
     <ClInclude Include="..\drivercode\wifirequest.h" />
     <ClInclude Include="..\drivercode\umkmfusion.h" />
+    <ClInclude Include="..\drivercode\wifitransition.h" />
   </ItemGroup>
   <ItemGroup>
     <ClCompile Include="..\drivercode\adapter.cpp" />
@@ -213,6 +214,7 @@
     <ClCompile Include="..\drivercode\memorymanagement.cpp" />
     <ClCompile Include="..\drivercode\wifihal.cpp" />
     <ClCompile Include="..\drivercode\wifirequest.cpp" />
+    <ClCompile Include="..\drivercode\wifitransition.cpp" />
   </ItemGroup>
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.targets" />
   <ImportGroup Label="ExtensionTargets">
diff --git a/network/wlan/WIFICX/um/wificxsampleclientum.vcxproj.filters b/network/wlan/WIFICX/um/wificxsampleclientum.vcxproj.filters
index 099c33f8..a10fb659 100644
--- a/network/wlan/WIFICX/um/wificxsampleclientum.vcxproj.filters
+++ b/network/wlan/WIFICX/um/wificxsampleclientum.vcxproj.filters
@@ -54,6 +54,9 @@
     <ClInclude Include="..\drivercode\wifihal.h">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="..\drivercode\wifitransition.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
   </ItemGroup>
   <ItemGroup>
     <ClCompile Include="..\drivercode\adapter.cpp">
@@ -74,5 +77,8 @@
     <ClCompile Include="..\drivercode\wifihal.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="..\drivercode\wifitransition.cpp">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
 </Project>
\ No newline at end of file