From 27076a2fdf2c4adce7bf34cf62ad8990a439872b Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Mon, 2 May 2022 17:11:42 +0200
Subject: [PATCH] wifi: Use a single SCANNING state for StaWifiMac

---
 src/wifi/model/sta-wifi-mac.cc | 22 +++++++---------------
 src/wifi/model/sta-wifi-mac.h  | 34 ++++++++++++----------------------
 2 files changed, 19 insertions(+), 37 deletions(-)

diff --git a/src/wifi/model/sta-wifi-mac.cc b/src/wifi/model/sta-wifi-mac.cc
index c1b7122aa..eee40bf50 100644
--- a/src/wifi/model/sta-wifi-mac.cc
+++ b/src/wifi/model/sta-wifi-mac.cc
@@ -126,7 +126,7 @@ StaWifiMac::SetActiveProbing (bool enable)
 {
   NS_LOG_FUNCTION (this << enable);
   m_activeProbing = enable;
-  if (m_state == WAIT_PROBE_RESP || m_state == WAIT_BEACON)
+  if (m_state == SCANNING)
     {
       NS_LOG_DEBUG ("STA is still scanning, reset scanning process");
       StartScanning ();
@@ -304,16 +304,9 @@ StaWifiMac::TryToEnsureAssociated (void)
     case ASSOCIATED:
       return;
       break;
-    case WAIT_PROBE_RESP:
-      /* we have sent a probe request earlier so we
-         do not need to re-send a probe request immediately.
-         We just need to wait until probe-request-timeout
-         or until we get a probe response
-       */
-      break;
-    case WAIT_BEACON:
-      /* we have initiated passive scanning, continue to wait
-         and gather beacons
+    case SCANNING:
+      /* we have initiated active or passive scanning, continue to wait
+         and gather beacons or probe responses until the scanning timeout
        */
       break;
     case UNASSOCIATED:
@@ -355,7 +348,6 @@ StaWifiMac::StartScanning (void)
     }
   if (GetActiveProbing ())
     {
-      SetState (WAIT_PROBE_RESP);
       SendProbeRequest ();
       m_probeRequestEvent = Simulator::Schedule (m_probeRequestTimeout,
                                                  &StaWifiMac::ScanningTimeout,
@@ -363,11 +355,11 @@ StaWifiMac::StartScanning (void)
     }
   else
     {
-      SetState (WAIT_BEACON);
       m_waitBeaconEvent = Simulator::Schedule (m_waitBeaconTimeout,
                                                &StaWifiMac::ScanningTimeout,
                                                this);
     }
+  SetState (SCANNING);
 }
 
 void
@@ -658,7 +650,7 @@ StaWifiMac::Receive (Ptr<WifiMacQueueItem> mpdu, uint8_t linkId)
           RestartBeaconWatchdog (delay);
           UpdateApInfoFromBeacon (beacon, hdr->GetAddr2 (), hdr->GetAddr3 (), linkId);
         }
-      if (goodBeacon && m_state == WAIT_BEACON)
+      if (goodBeacon && m_state == SCANNING)
         {
           NS_LOG_DEBUG ("Beacon received while scanning from " << hdr->GetAddr2 ());
           SnrTag snrTag;
@@ -677,7 +669,7 @@ StaWifiMac::Receive (Ptr<WifiMacQueueItem> mpdu, uint8_t linkId)
     }
   else if (hdr->IsProbeResp ())
     {
-      if (m_state == WAIT_PROBE_RESP)
+      if (m_state == SCANNING)
         {
           NS_LOG_DEBUG ("Probe response received while scanning from " << hdr->GetAddr2 ());
           MgtProbeResponseHeader probeResp;
diff --git a/src/wifi/model/sta-wifi-mac.h b/src/wifi/model/sta-wifi-mac.h
index fa017d92b..d42c4e4ff 100644
--- a/src/wifi/model/sta-wifi-mac.h
+++ b/src/wifi/model/sta-wifi-mac.h
@@ -64,32 +64,23 @@ struct ApInfo
    ---------       --------------                           |      /      -----------
       |              |   /------------------------------\   |     /
       \              v   v                              |   v    /
-       \    ----------------     ---------------     -----------------------------
-        \-> | Unassociated | --> | Wait Beacon | --> | Wait Association Response |
-            ----------------     ---------------     -----------------------------
-                  \                  ^     ^ |              ^    ^ |
-                   \                 |     | |              |    | |
-                    \                v      -               /     -
-                     \    -----------------------          /
-                      \-> | Wait Probe Response | --------/
-                          -----------------------
-                                  ^ |
-                                  | |
-                                   -
+       \    ----------------     ------------     -----------------------------
+        \-> | Unassociated | --> | Scanning | --> | Wait Association Response |
+            ----------------     ------------     -----------------------------
+                                      ^ |                  ^ |
+                                      | |                  | |
+                                       -                    -
    \endverbatim
  *
  * Notes:
  * 1. The state 'Start' is not included in #MacState and only used
  *    for illustration purpose.
  * 2. The Unassociated state is a transient state before STA starts the
- *    scanning procedure which moves it into either Wait Beacon or Wait
- *    Probe Response, based on whether passive or active scanning is
- *    selected.
- * 3. In Wait Beacon and Wait Probe Response, STA is gathering beacon or
- *    probe response packets from APs, resulted in a list of candidate AP.
- *    After the respective timeout, it then tries to associate to the best
- *    AP (i.e., best SNR). STA will switch between the two states and
- *    restart the scanning procedure if SetActiveProbing() called.
+ *    scanning procedure which moves it into the Scanning state.
+ * 3. In Scanning, STA is gathering beacon or probe response frames from APs,
+ *    resulted in a list of candidate AP. After the timeout, it then tries to
+ *    associate to the best AP, which is indicated by the Association Manager.
+ *    STA will restart the scanning procedure if SetActiveProbing() called.
  * 4. In the case when AP responded to STA's association request with a
  *    refusal, STA will try to associate to the next best AP until the list
  *    of candidate AP is exhausted which sends STA to Refused state.
@@ -161,8 +152,7 @@ private:
   enum MacState
   {
     ASSOCIATED,
-    WAIT_BEACON,
-    WAIT_PROBE_RESP,
+    SCANNING,
     WAIT_ASSOC_RESP,
     UNASSOCIATED,
     REFUSED