From eeb24229f0823c8b6c037a00a5391b6a220fd53f Mon Sep 17 00:00:00 2001
From: Abtin Keshavarzian <abtink@google.com>
Date: Thu, 4 Jan 2018 09:07:16 -0800
Subject: [PATCH] [mesh-forwarder] use extended address for data poll during
 parent switch attempt (#2449)

This commit changes the logic for preparing a data poll (802.15.4 MAC
Data Request) frame and selecting short/extended address. It ensures
that during a parent switch attempt on a sleepy-end-device, the data
polls use the extended address as the source MAC address.
---
 src/core/thread/mesh_forwarder.cpp | 59 ++++++++++++++++--------------
 src/core/thread/mesh_forwarder.hpp |  1 +
 2 files changed, 32 insertions(+), 28 deletions(-)

diff --git a/src/core/thread/mesh_forwarder.cpp b/src/core/thread/mesh_forwarder.cpp
index 801bc519d..6eeee7260 100644
--- a/src/core/thread/mesh_forwarder.cpp
+++ b/src/core/thread/mesh_forwarder.cpp
@@ -615,35 +615,8 @@ Message *MeshForwarder::GetDirectTransmission(void)
             break;
 
         case Message::kTypeMacDataPoll:
-        {
-            ThreadNetif &netif = GetNetif();
-            Neighbor *parent = netif.GetMle().GetParentCandidate();
-
-            if ((parent != NULL) && (parent->IsStateValidOrRestoring()))
-            {
-                mMacSource.mShortAddress = netif.GetMac().GetShortAddress();
-
-                if (mMacSource.mShortAddress != Mac::kShortAddrInvalid)
-                {
-                    mMacSource.mLength = sizeof(mMacSource.mShortAddress);
-                    mMacDest.mLength = sizeof(mMacDest.mShortAddress);
-                    mMacDest.mShortAddress = parent->GetRloc16();
-                }
-                else
-                {
-                    mMacSource.mLength = sizeof(mMacSource.mExtAddress);
-                    mMacSource.mExtAddress = netif.GetMac().GetExtAddress();
-                    mMacDest.mLength = sizeof(mMacDest.mExtAddress);
-                    mMacDest.mExtAddress = parent->GetExtAddress();
-                }
-            }
-            else
-            {
-                error = OT_ERROR_DROP;
-            }
-
+            error = PrepareDataPoll();
             break;
-        }
 
         case Message::kTypeSupervision:
             error = OT_ERROR_DROP;
@@ -778,6 +751,36 @@ void MeshForwarder::PrepareIndirectTransmission(Message &aMessage, const Child &
     }
 }
 
+otError MeshForwarder::PrepareDataPoll(void)
+{
+    otError error = OT_ERROR_NONE;
+    ThreadNetif &netif = GetNetif();
+    Neighbor *parent = netif.GetMle().GetParentCandidate();
+    uint16_t shortAddress;
+
+    VerifyOrExit((parent != NULL) && parent->IsStateValidOrRestoring(), error = OT_ERROR_DROP);
+
+    shortAddress = netif.GetMac().GetShortAddress();
+
+    if ((shortAddress == Mac::kShortAddrInvalid) || (parent != netif.GetMle().GetParent()))
+    {
+        mMacSource.mLength = sizeof(mMacSource.mExtAddress);
+        mMacSource.mExtAddress = netif.GetMac().GetExtAddress();
+        mMacDest.mLength = sizeof(mMacDest.mExtAddress);
+        mMacDest.mExtAddress = parent->GetExtAddress();
+    }
+    else
+    {
+        mMacSource.mLength = sizeof(mMacSource.mShortAddress);
+        mMacSource.mShortAddress = shortAddress;
+        mMacDest.mLength = sizeof(mMacDest.mShortAddress);
+        mMacDest.mShortAddress = parent->GetRloc16();
+    }
+
+exit:
+    return error;
+}
+
 otError MeshForwarder::UpdateMeshRoute(Message &aMessage)
 {
     ThreadNetif &netif = GetNetif();
diff --git a/src/core/thread/mesh_forwarder.hpp b/src/core/thread/mesh_forwarder.hpp
index e0b37822b..e08b290cc 100644
--- a/src/core/thread/mesh_forwarder.hpp
+++ b/src/core/thread/mesh_forwarder.hpp
@@ -275,6 +275,7 @@ private:
     Message *GetIndirectTransmission(Child &aChild);
     otError PrepareDiscoverRequest(void);
     void PrepareIndirectTransmission(Message &aMessage, const Child &aChild);
+    otError PrepareDataPoll(void);
     void HandleMesh(uint8_t *aFrame, uint8_t aPayloadLength, const Mac::Address &aMacSource,
                     const otThreadLinkInfo &aLinkInfo);
     void HandleFragment(uint8_t *aFrame, uint8_t aPayloadLength,