have the M3 amd M4 general types

3 files changed, 268 insertions, 242 deletions

diff --git a/network/wlan/WIFICX/drivercode/wifihal.cpp b/network/wlan/WIFICX/drivercode/wifihal.cpp
index edeb8fe0..bbfa04e0 100644
--- a/network/wlan/WIFICX/drivercode/wifihal.cpp
+++ b/network/wlan/WIFICX/drivercode/wifihal.cpp
@@ -6,6 +6,24 @@
 #include "WifiHal.h"
 #include "WifiHal.tmh"
 
+WifiHAL::WifiHAL(WDFDEVICE Device)
+{
+    m_Device = Device;
+    m_TlvContext = &WifiGetIhvDeviceContext(Device)->TlvContext;
+}
+
+NTSTATUS WifiHAL::WifiIhvIsDeviceReadyForRequest()
+{
+    return ((m_Device != WDF_NO_HANDLE) // Make sure device is initialized (since this is hardware abstraction layer, IHV can replace with firmware state)
+        && (WifiGetIhvDeviceContext(m_Device)->primaryStaAdapter != WDF_NO_HANDLE) // In WIFICX, the logic sits on top of primary STA adapter, make sure it is initialized
+        && (m_CurrentRadioState == TRUE) ? STATUS_SUCCESS : STATUS_DEVICE_NOT_READY); // Device is ready only if radio is ON (this may be replaced by other state in IHV implementation)
+}
+
+NTSTATUS WifiHAL::WifiIhvGetPendingTransitionStatus()
+{
+    return m_LastConnectTransactionId ==0 ? STATUS_SUCCESS : STATUS_PENDING;
+}
+
 NTSTATUS WifiHAL::WifiIhvSetDeviceCapabilities()
 {
     WIFI_DEVICE_CAPABILITIES deviceCaps = {};
@@ -432,7 +450,7 @@ NTSTATUS WifiHAL::WifiIhvSetDeviceCapabilities()
     return STATUS_SUCCESS;
 }
 
-NTSTATUS WifiHAL::WifiIhvReset(const WDI_TASK_DOT11_RESET_PARAMETERS& ResetParameters, const PWDI_MESSAGE_HEADER)
+NTSTATUS WifiHAL::WifiIhvReset(const WDI_TASK_DOT11_RESET_PARAMETERS& ResetParameters, const PWDI_MESSAGE_HEADER, UINT)
 {
     if (0 == ResetParameters.Optional.ResetMACAddress_IsPresent)
     {
@@ -458,7 +476,7 @@ NTSTATUS WifiHAL::WifiIhvReset(const WDI_TASK_DOT11_RESET_PARAMETERS& ResetParam
 }
 
 _Use_decl_annotations_
-NTSTATUS WifiHAL::WifiIhvSetRadioState(const WDI_SET_RADIO_STATE_PARAMETERS& RadioState, const PWDI_MESSAGE_HEADER pWdiHeader)
+NTSTATUS WifiHAL::WifiIhvSetRadioState(const WDI_SET_RADIO_STATE_PARAMETERS& RadioState, const PWDI_MESSAGE_HEADER pWdiHeader, UINT)
 {
     if (RadioState.SoftwareRadioState != m_CurrentRadioState)
     {
@@ -483,7 +501,7 @@ NTSTATUS WifiHAL::WifiIhvSetRadioState(const WDI_SET_RADIO_STATE_PARAMETERS& Rad
 }
 
 _Use_decl_annotations_
-NTSTATUS WifiHAL::WifiIhvScan(WDI_SCAN_PARAMETERS& ScanParameters, const PWDI_MESSAGE_HEADER pWdiHeader)
+NTSTATUS WifiHAL::WifiIhvScan(const WDI_SCAN_PARAMETERS& ScanParameters, const PWDI_MESSAGE_HEADER pWdiHeader, UINT)
 {
     for (UINT8 connectEntry = 1; connectEntry < ConnectEntryId_MAX; connectEntry++)
     {
@@ -530,7 +548,7 @@ NTSTATUS WifiHAL::WifiIhvScan(WDI_SCAN_PARAMETERS& ScanParameters, const PWDI_ME
     return STATUS_SUCCESS;
 }
 
-NTSTATUS WifiHAL::WifiIhvConnect(WDI_TASK_CONNECT_PARAMETERS& ConnectParameters, const PWDI_MESSAGE_HEADER pWdiHeader)
+NTSTATUS WifiHAL::WifiIhvConnect(const WDI_TASK_CONNECT_PARAMETERS& ConnectParameters, const PWDI_MESSAGE_HEADER pWdiHeader, UINT)
 {
     NT_ASSERT(m_LastConnectEntryId == 0);
     if (m_LastConnectEntryId != 0) // Not Disconnected State
@@ -550,7 +568,11 @@ NTSTATUS WifiHAL::WifiIhvConnect(WDI_TASK_CONNECT_PARAMETERS& ConnectParameters,
     //
     if (WDI_AUTH_ALGO_WPA3_SAE == m_LastAuthAlgo)
     {
-        return STATUS_PENDING;
+        m_LastConnectTransactionId = pWdiHeader->TransactionId;
+    }
+    else 
+    {
+        m_LastConnectTransactionId = 0;
     }
     return STATUS_SUCCESS;
 }
@@ -618,9 +640,9 @@ NTSTATUS WifiHAL::WifiIhvSendLinkStateIndication(_In_ PWDI_MESSAGE_HEADER pWdiHe
 
 _Use_decl_annotations_
 NTSTATUS WifiHAL::WifiIhvPerformAssociation(
-    struct ArrayOfElements<WDI_CONNECT_BSS_ENTRY_CONTAINER>* pPreferredBSSEntryList,
-    struct ArrayOfElements<WDI_AUTH_ALGORITHM>* pAuthenticationAlgorithms,
-    PWDI_MESSAGE_HEADER pWdiHeader)
+    const struct ArrayOfElements<WDI_CONNECT_BSS_ENTRY_CONTAINER>* pPreferredBSSEntryList,
+    const struct ArrayOfElements<WDI_AUTH_ALGORITHM>* pAuthenticationAlgorithms,
+    const PWDI_MESSAGE_HEADER pWdiHeader)
 {
 
     ULONG bssIndex = 0;
@@ -740,7 +762,7 @@ NTSTATUS WifiHAL::WifiIhvPerformAssociation(
     return ntStatus;
 }
 
-NTSTATUS WifiHAL::WifiIhvSetSaeAuthParams(const WDI_SET_SAE_AUTH_PARAMS_COMMAND& setSAEAuthParams, const PWDI_MESSAGE_HEADER pWdiHeader)
+NTSTATUS WifiHAL::WifiIhvSetSaeAuthParams(const WDI_SET_SAE_AUTH_PARAMS_COMMAND& setSAEAuthParams, const PWDI_MESSAGE_HEADER pWdiHeader, UINT)
 {
     //Since this is DIRECT OID, need to check the m_LastConnectTransactionId match
     if (pWdiHeader->TransactionId != m_LastConnectTransactionId) 
@@ -837,8 +859,9 @@ NTSTATUS WifiHAL::WifiIhvSetSaeAuthParams(const WDI_SET_SAE_AUTH_PARAMS_COMMAND&
         WX_RETURN_NTSTATUS_IF_NOT_NT_SUCCESS_MSG(WifiIhvSendLinkStateIndication(pWdiHeader, 1),
             "Failed WifiIhvSendLinkStateIndication");
 
-        WifiIhvSendM4IndicationToOs(
-            m_Device, WDI_INDICATION_CONNECT_COMPLETE, pWdiHeader, STATUS_SUCCESS);
+        // Complete the transaction to let the M4 hanlder WifiIhvGetPendingTransitionStatus know that association is complete
+        m_LastConnectTransactionId = 0;
+
 #ifdef WIFI_IHV_HANDSHAKE
         //
         // Receive M1 frame of 4-way handshake
@@ -850,6 +873,7 @@ NTSTATUS WifiHAL::WifiIhvSetSaeAuthParams(const WDI_SET_SAE_AUTH_PARAMS_COMMAND&
         //
         RecieveDatapathFrame(0x33, sizeof(pucM3SaeFrame), pucM3SaeFrame);
 #endif
+        return STATUS_SUCCESS;
     }
 
     RtlCopyMemory(SAEAuthParamsNeeded.BssId.Address, &g_ConnectEntries[m_LastConnectEntryId].pMacAddress, sizeof(DOT11_MAC_ADDRESS));
@@ -871,7 +895,7 @@ NTSTATUS WifiHAL::WifiIhvSetSaeAuthParams(const WDI_SET_SAE_AUTH_PARAMS_COMMAND&
 }
 
 _Use_decl_annotations_
-void WifiHAL::WifiIhvPerformDisassociation(_In_ PWDI_MESSAGE_HEADER pWdiHeader, _In_ WDI_ASSOC_STATUS DisassocStatus)
+NTSTATUS WifiHAL::WifiIhvDisconnect(const WDI_TASK_DISCONNECT_PARAMETERS&, const PWDI_MESSAGE_HEADER pWdiHeader, UINT)
 {
     UCHAR s_TLV_Disassociation[] =
     {
@@ -888,6 +912,7 @@ void WifiHAL::WifiIhvPerformDisassociation(_In_ PWDI_MESSAGE_HEADER pWdiHeader,
 
             0x00
     };
+    WDI_ASSOC_STATUS DisassocStatus = WDI_ASSOC_STATUS_DISASSOCIATED_BY_HOST;
 
     RtlCopyMemory(&s_TLV_Disassociation[4], &m_ConnectedPeer, sizeof(DOT11_MAC_ADDRESS));
     RtlCopyMemory(&s_TLV_Disassociation[10], &DisassocStatus, sizeof(ULONG));
@@ -900,10 +925,6 @@ void WifiHAL::WifiIhvPerformDisassociation(_In_ PWDI_MESSAGE_HEADER pWdiHeader,
 #ifdef WIFI_IHV_NETV
     DeleteDatapathPeer(0x33);
 #endif
-}
 
-WifiHAL::WifiHAL(WDFDEVICE Device)
-{
-    m_Device = Device;
-    m_TlvContext = &WifiGetIhvDeviceContext(Device)->TlvContext;
+    return STATUS_SUCCESS;
 }
diff --git a/network/wlan/WIFICX/drivercode/wifihal.h b/network/wlan/WIFICX/drivercode/wifihal.h
index d2d770d9..844a9ccb 100644
--- a/network/wlan/WIFICX/drivercode/wifihal.h
+++ b/network/wlan/WIFICX/drivercode/wifihal.h
@@ -6,17 +6,25 @@ class WifiHAL
 {
 public:
     WifiHAL(_In_ WDFDEVICE Device);
+
+    // Wifi request M3 working condition verification function
+    NTSTATUS WifiIhvIsDeviceReadyForRequest();
+
+    // Wifi request pending transition status check function
+    // for example, during SAE authentication, the connect request is pending until SAE exchange is complete
+    NTSTATUS WifiIhvGetPendingTransitionStatus();
+
     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);
+    NTSTATUS WifiIhvReset(_In_ const WDI_TASK_DOT11_RESET_PARAMETERS& ResetParameters, _In_ const PWDI_MESSAGE_HEADER pWdiHeader, _In_ UINT BytesWritten);
+    NTSTATUS WifiIhvSetRadioState(_In_ const WDI_SET_RADIO_STATE_PARAMETERS& RadioState, _In_ const PWDI_MESSAGE_HEADER pWdiHeader, _In_ UINT BytesWritten);
+    NTSTATUS WifiIhvScan(_In_ const WDI_SCAN_PARAMETERS& ScanParameters, _In_ const PWDI_MESSAGE_HEADER pWdiHeader, _In_ UINT BytesWritten);
+    NTSTATUS WifiIhvConnect(_In_ const WDI_TASK_CONNECT_PARAMETERS& ConnectParameters, _In_ const PWDI_MESSAGE_HEADER pWdiHeader, _In_ UINT BytesWritten);
+    NTSTATUS WifiIhvSetSaeAuthParams(_In_ const WDI_SET_SAE_AUTH_PARAMS_COMMAND& setSAEAuthParams, _In_ const PWDI_MESSAGE_HEADER pWdiHeader, _In_ UINT BytesWritten);
+    NTSTATUS WifiIhvDisconnect(_In_ const WDI_TASK_DISCONNECT_PARAMETERS& disconnectParameters, _In_ const PWDI_MESSAGE_HEADER pWdiHeader, _In_ UINT BytesWritten);
 private:
-    NTSTATUS WifiIhvSendLinkStateIndication(_In_ PWDI_MESSAGE_HEADER pWdiHeader, ULONG numLinks);
-    
+    NTSTATUS WifiIhvPerformAssociation(_In_ const struct ArrayOfElements<WDI_CONNECT_BSS_ENTRY_CONTAINER>* pPreferredBSSEntryList, _In_ const struct ArrayOfElements<WDI_AUTH_ALGORITHM>* pAuthenticationAlgorithms, _In_ const PWDI_MESSAGE_HEADER pWdiHeader);
+    NTSTATUS WifiIhvSendLinkStateIndication(_In_ const PWDI_MESSAGE_HEADER pWdiHeader, ULONG numLinks);
+
     WDFDEVICE m_Device{};
     PCTLV_CONTEXT m_TlvContext{};
     
@@ -27,8 +35,8 @@ private:
     DOT11_MAC_ADDRESS m_ConnectedPeer{};
 
     const WDI_MAC_ADDRESS m_LocalLinkAddresses[2] = {
-    {0x11, 0x01, 0x02, 0x03, 0x04, 0x21},
-    {0x11, 0x01, 0x02, 0x03, 0x04, 0x22},
+        {0x11, 0x01, 0x02, 0x03, 0x04, 0x21},
+        {0x11, 0x01, 0x02, 0x03, 0x04, 0x22},
     };
 
     const ULONG m_SupportedBands = (WDI_BAND_ID_2400 | WDI_BAND_ID_5000 | WDI_BAND_ID_6000);
diff --git a/network/wlan/WIFICX/drivercode/wifitransition.cpp b/network/wlan/WIFICX/drivercode/wifitransition.cpp
index 9ac857bc..eccf9dc3 100644
--- a/network/wlan/WIFICX/drivercode/wifitransition.cpp
+++ b/network/wlan/WIFICX/drivercode/wifitransition.cpp
@@ -32,255 +32,132 @@ NTSTATUS ParseTlvCommon(TransitionContext& ctx,
     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>
+// --- Generic pure-type traits template (add before existing specializations) ---
+// WIFIREQUEST always needs M3 notification, so TPreM3Fn is mandatory.
+// and WIFICX expectes the M3 then M4 order, so we always execute M3 then M4.
+// using template parameters to configure parsing, cleanup, M3/M4 steps.
+// to make sure that all transitions have consistent implementations.
+template<
+    UINT16 TMsgId,
+    typename TParam,
+    UINT16 TCompleteIndication,
+    bool TDumpTlvStream,
+    NDIS_STATUS (*TParseFn)(ULONG, const UINT8*, PCTLV_CONTEXT, TParam*),
+    void (*TCleanupFn)(TParam*),
+    NTSTATUS (WifiHAL::*TPreM3Fn)(),                                                                   // mandatory pre-M3 hook
+    NTSTATUS (WifiHAL::*THalM3Fn)(const TParam&, const PWDI_MESSAGE_HEADER, UINT BytesWriten),         // optional HAL M3 (may be nullptr)
+    NTSTATUS (WifiHAL::*TPreM4Fn)(),                                                                   // optional pre-M4 hook (may be nullptr)
+    NTSTATUS (WifiHAL::*THalM4Fn)(const PWDI_MESSAGE_HEADER)                                           // optional HAL M4 (may be nullptr)
+>
+struct GenericTransitionTraits
 {
-    using ParamType = WDI_SET_RADIO_STATE_PARAMETERS;
-    enum : UINT16 { CompleteIndication = WDI_INDICATION_SET_RADIO_STATE_COMPLETE };
+    using ParamType = TParam;
+    enum : UINT16 { CompleteIndication = TCompleteIndication };
 
     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); }
+        if (ctx.InLen < sizeof(WDI_MESSAGE_HEADER))
+        {
+            return STATUS_INVALID_PARAMETER;
+        }
 
-    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; }
-};
+        auto* tlvBytes = static_cast<UCHAR*>(ctx.RawBuffer) + sizeof(WDI_MESSAGE_HEADER);
+        auto tlvLen = static_cast<ULONG>(ctx.InLen - sizeof(WDI_MESSAGE_HEADER));
 
-// -------- WDI_TASK_SCAN --------
-template<>
-struct TransitionTraits<WDI_TASK_SCAN>
-{
-    using ParamType = WDI_SCAN_PARAMETERS;
-    enum : UINT16 { CompleteIndication = WDI_INDICATION_SCAN_COMPLETE };
+        if (TDumpTlvStream)
+        {
+            DumpMessageTlvByteStream(
+                TMsgId,
+                TRUE,
+                ctx.DevCtx->TlvContext.PeerVersion,
+                tlvLen,
+                tlvBytes,
+                0,
+                nullptr);
+        }
 
-    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);
+        auto ndisStatus = TParseFn(tlvLen, tlvBytes, &ctx.DevCtx->TlvContext, &p);
+        return Wifi::ConvertNDISSTATUSToNTSTATUS(ndisStatus);
     }
-    void Cleanup(ParamType& p) { CleanupParsedWdiTaskScan(&p); }
 
-    NTSTATUS Handle(TransitionContext& ctx, ParamType& p)
+    void Cleanup(ParamType& p) { TCleanupFn(&p); }
+
+    // Make static so pointer matches ExecuteSteps expected callable type (no implicit this)
+    static NTSTATUS StepM3(TransitionContext& c, ParamType& p, UINT& bytesWritten)
     {
-        return ExecuteSteps(ctx, p,
-            // STEPM3
-            [](TransitionContext&, ParamType&, UINT& bytesWritten)
-            {
-                bytesWritten = sizeof(WDI_MESSAGE_HEADER);
-                return STATUS_SUCCESS;
-            },
-            // STEPM4
-            [](TransitionContext& c, ParamType& par)
+        bytesWritten = sizeof(WDI_MESSAGE_HEADER);
+        ASSERT(TPreM3Fn);
+
+        if (TPreM3Fn)
+        {
+            // Call member function pointer on WifiHAL instance
+            NTSTATUS preStatus = (c.DevCtx->wifiHAL->*TPreM3Fn)();
+            if (!NT_SUCCESS(preStatus))
             {
-                return c.DevCtx->wifiHAL->WifiIhvScan(par, c.Header);
+                return preStatus;
             }
-        );
-    }
-    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 };
+        if (THalM3Fn)
+        {
+            // Pass required third argument (BytesWriten) to HAL M3 function
+            return (c.DevCtx->wifiHAL->*THalM3Fn)(p, c.Header, bytesWritten);
+        }
 
-    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);
+        return STATUS_SUCCESS;
     }
-    void Cleanup(ParamType& p) { CleanupParsedWdiTaskDot11Reset(&p); }
 
-    NTSTATUS Handle(TransitionContext& ctx, ParamType& p)
+    static NTSTATUS StepM4(TransitionContext& c, ParamType&)
     {
-        return ExecuteSteps(ctx, p,
-            // STEPM3
-            [](TransitionContext&, ParamType&, UINT& bytesWritten)
+        if (TPreM4Fn)
+        {
+            NTSTATUS preStatus = (c.DevCtx->wifiHAL->*TPreM4Fn)();
+            if (!NT_SUCCESS(preStatus))
             {
-                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; }
-};
+                return preStatus;
+            }
+        }
 
-// -------- WDI_TASK_CONNECT --------
-template<>
-struct TransitionTraits<WDI_TASK_CONNECT>
-{
-    using ParamType = WDI_TASK_CONNECT_PARAMETERS;
-    enum : UINT16 { CompleteIndication = WDI_INDICATION_CONNECT_COMPLETE };
+        if (THalM4Fn)
+        {
+            return (c.DevCtx->wifiHAL->*THalM4Fn)(c.Header);
+        }
 
-    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);
+        return (TPreM4Fn == nullptr && THalM4Fn == nullptr) ? STATUS_PENDING : STATUS_SUCCESS;
     }
-    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);
-            });
+        return ExecuteSteps(ctx, p, &GenericTransitionTraits::StepM3, &GenericTransitionTraits::StepM4);
     }
-    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; }
+    bool ShouldSendComplete(NTSTATUS s) const { return s != STATUS_PENDING; }
 };
 
-// -------- WDI_SET_SAE_AUTH_PARAMS --------
-template<>
-struct TransitionTraits<WDI_SET_SAE_AUTH_PARAMS>
+// 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)
 {
-    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
+    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))
     {
-        // Do not send M4 indication for WDI_SET_SAE_AUTH_PARAMS (original code omitted it).
-        return false;
+        return m3Status;
     }
-};
+    return stepM4(ctx, p);
+}
 
 // -------- Generic runner (compile-time) --------
 template<UINT16 MsgId>
@@ -313,6 +190,126 @@ NTSTATUS RunTransition(TransitionContext& ctx)
     return m4Status;
 }
 
+//// -------- SCENARIO: [Connect with a SAE WI-FI7 network --------
+///     Demo: Handle WDI_TASK_CONNECT + WDI_SET_SAE_AUTH_PARAMS then WDI_TASK_DISCONNECT
+///     Scope:
+///            -WifiRequest WDI_TASK_CONNECT & WDI_TASK_DISCONNECT are both WIFICX task commands, which is a two step M3/M4 transition
+///            -The direct WifiRequest WDI_SET_SAE_AUTH_PARAMS, which is a single step transition but
+///             is logically part of the connect scenario. since WDI_SET_SAE_AUTH_PARAMS is WIFICX property command,
+///             it only has M3 step, no M4 step.
+///            - The WifiCx unsolicited indication e.g. WDI_INDICATION_SAE_AUTH_PARAMS_NEEDED is sent from the HAL during the connect process,
+///     Notes:
+///            - M3 and M4 status mainly used for WifiCx to track progress of the transition. e.g. the hung detection and trigger recovery.
+///            - The actual scenario result is reported through unsolicited indication.
+/// 
+
+// -------- WDI_TASK_CONNECT --------
+template<>
+struct TransitionTraits<WDI_TASK_CONNECT>
+    : GenericTransitionTraits <
+    WDI_TASK_CONNECT,
+    WDI_TASK_CONNECT_PARAMETERS,
+    WDI_INDICATION_CONNECT_COMPLETE,
+    true, // dump TLV stream? (was true in original)
+    ParseWdiTaskConnect,
+    CleanupParsedWdiTaskConnect,
+    &WifiHAL::WifiIhvIsDeviceReadyForRequest, // pre-M3
+    &WifiHAL::WifiIhvConnect, // HAL M3
+    &WifiHAL::WifiIhvGetPendingTransitionStatus, // pre-M4
+    nullptr
+    >
+{
+};
+
+// --- WDI_SET_SAE_AUTH_PARAMS ---
+template<>
+struct TransitionTraits<WDI_SET_SAE_AUTH_PARAMS>
+    : GenericTransitionTraits<
+        WDI_SET_SAE_AUTH_PARAMS,
+        WDI_SET_SAE_AUTH_PARAMS_COMMAND,
+        WDI_INDICATION_CONNECT_COMPLETE,
+        false, // dump TLV stream? (was false in original)
+        ParseWdiSetSaeAuthParams,
+        CleanupParsedWdiSetSaeAuthParams,
+    &WifiHAL::WifiIhvIsDeviceReadyForRequest, // pre-M3
+    &WifiHAL::WifiIhvSetSaeAuthParams, // HAL M3
+    &WifiHAL::WifiIhvGetPendingTransitionStatus, // pre-M4
+    nullptr// HAL M4
+    >
+{};
+
+// -------- WDI_TASK_DISCONNECT --------
+template<>
+struct TransitionTraits<WDI_TASK_DISCONNECT>
+    : GenericTransitionTraits<
+    WDI_TASK_DISCONNECT,
+    WDI_TASK_DISCONNECT_PARAMETERS,
+    WDI_INDICATION_DISCONNECT_COMPLETE,
+    false, // dump TLV stream? (was false in original)
+    ParseWdiTaskDisconnect,
+    CleanupParsedWdiTaskDisconnect,
+    &WifiHAL::WifiIhvIsDeviceReadyForRequest, // pre-M3
+    &WifiHAL::WifiIhvDisconnect, // HAL M3
+    &WifiHAL::WifiIhvGetPendingTransitionStatus, // pre-M4
+    nullptr // HAL M4
+    >
+{
+};
+/// ----- End of scenario [Connect with a SAE WI-FI7 network]-----
+
+// -------- WDI_TASK_DOT11_RESET --------
+template<>
+struct TransitionTraits<WDI_TASK_DOT11_RESET>
+    : GenericTransitionTraits <
+    WDI_TASK_DOT11_RESET,
+    WDI_TASK_DOT11_RESET_PARAMETERS,
+    WDI_INDICATION_DOT11_RESET_COMPLETE,
+    false, // dump TLV stream? (was false in original)
+    ParseWdiTaskDot11Reset,
+    CleanupParsedWdiTaskDot11Reset,
+    &WifiHAL::WifiIhvIsDeviceReadyForRequest, // pre-M3
+    &WifiHAL::WifiIhvReset, // HAL M3
+    &WifiHAL::WifiIhvGetPendingTransitionStatus, // pre-M4
+    nullptr  // HAL M4
+    >
+{
+};
+
+// -------- WDI_TASK_SCAN --------
+template<>
+struct TransitionTraits<WDI_TASK_SCAN>
+    : GenericTransitionTraits <
+    WDI_TASK_SCAN,
+    WDI_SCAN_PARAMETERS,
+    WDI_INDICATION_SCAN_COMPLETE,
+    true, // dump TLV stream? (was true in original)
+    ParseWdiTaskScan,
+    CleanupParsedWdiTaskScan,
+    &WifiHAL::WifiIhvIsDeviceReadyForRequest, // pre-M3
+    &WifiHAL::WifiIhvScan, // HAL M3
+    &WifiHAL::WifiIhvGetPendingTransitionStatus, // pre-M4
+    nullptr  // HAL M4
+    >
+{
+};
+
+// -------- WDI_TASK_SET_RADIO_STATE --------
+template<>
+struct TransitionTraits<WDI_TASK_SET_RADIO_STATE>
+    : GenericTransitionTraits <
+    WDI_TASK_SET_RADIO_STATE,
+    WDI_SET_RADIO_STATE_PARAMETERS,
+    WDI_INDICATION_SET_RADIO_STATE_COMPLETE,
+    true, // dump TLV stream? (was true in original)
+    ParseWdiTaskSetRadioState,
+    CleanupParsedWdiTaskSetRadioState,
+    &WifiHAL::WifiIhvIsDeviceReadyForRequest, // pre-M3
+    &WifiHAL::WifiIhvSetRadioState, // HAL M3
+    &WifiHAL::WifiIhvGetPendingTransitionStatus, // pre-M4
+    nullptr  // HAL M4
+    >
+{
+};
 
 // Runtime dispatcher switches on MessageId and invokes the matching compile-time runner.
 NTSTATUS RunTransitionByMessage(TransitionContext& ctx, UINT16 messageId)