From 3e6f00d76b622256b2caf2ba1d1a944412f9009c Mon Sep 17 00:00:00 2001
From: Jonathan Hui <jonhui@nestlabs.com>
Date: Mon, 1 Aug 2016 09:35:51 -0700
Subject: [PATCH] Accept MLE messages that have MAC security enabled. (#311)

---
 src/core/thread/mle_router.cpp | 80 +++++++++++++++++++++-------------
 1 file changed, 50 insertions(+), 30 deletions(-)

diff --git a/src/core/thread/mle_router.cpp b/src/core/thread/mle_router.cpp
index 8a9d6acaf..cdb43ff72 100644
--- a/src/core/thread/mle_router.cpp
+++ b/src/core/thread/mle_router.cpp
@@ -2132,25 +2132,35 @@ Neighbor *MleRouter::GetNeighbor(uint16_t aAddress)
         ExitNow();
     }
 
-    if (mDeviceState == kDeviceStateChild && (rval = Mle::GetNeighbor(aAddress)) != NULL)
+    switch (mDeviceState)
     {
-        ExitNow();
-    }
+    case kDeviceStateDisabled:
+        break;
 
-    for (int i = 0; i < kMaxChildren; i++)
-    {
-        if (mChildren[i].mState == Neighbor::kStateValid && mChildren[i].mValid.mRloc16 == aAddress)
-        {
-            ExitNow(rval = &mChildren[i]);
-        }
-    }
+    case kDeviceStateDetached:
+    case kDeviceStateChild:
+        rval = Mle::GetNeighbor(aAddress);
+        break;
 
-    for (int i = 0; i < kMaxRouterId; i++)
-    {
-        if (mRouters[i].mState == Neighbor::kStateValid && mRouters[i].mValid.mRloc16 == aAddress)
+    case kDeviceStateRouter:
+    case kDeviceStateLeader:
+        for (int i = 0; i < kMaxChildren; i++)
         {
-            ExitNow(rval = &mRouters[i]);
+            if (mChildren[i].mState == Neighbor::kStateValid && mChildren[i].mValid.mRloc16 == aAddress)
+            {
+                ExitNow(rval = &mChildren[i]);
+            }
         }
+
+        for (int i = 0; i < kMaxRouterId; i++)
+        {
+            if (mRouters[i].mState == Neighbor::kStateValid && mRouters[i].mValid.mRloc16 == aAddress)
+            {
+                ExitNow(rval = &mRouters[i]);
+            }
+        }
+
+        break;
     }
 
 exit:
@@ -2161,27 +2171,37 @@ Neighbor *MleRouter::GetNeighbor(const Mac::ExtAddress &aAddress)
 {
     Neighbor *rval = NULL;
 
-    if (mDeviceState == kDeviceStateChild && (rval = Mle::GetNeighbor(aAddress)) != NULL)
+    switch (mDeviceState)
     {
-        ExitNow();
-    }
+    case kDeviceStateDisabled:
+        break;
 
-    for (int i = 0; i < kMaxChildren; i++)
-    {
-        if (mChildren[i].mState == Neighbor::kStateValid &&
-            memcmp(&mChildren[i].mMacAddr, &aAddress, sizeof(mChildren[i].mMacAddr)) == 0)
-        {
-            ExitNow(rval = &mChildren[i]);
-        }
-    }
+    case kDeviceStateDetached:
+    case kDeviceStateChild:
+        rval = Mle::GetNeighbor(aAddress);
+        break;
 
-    for (int i = 0; i < kMaxRouterId; i++)
-    {
-        if (mRouters[i].mState == Neighbor::kStateValid &&
-            memcmp(&mRouters[i].mMacAddr, &aAddress, sizeof(mRouters[i].mMacAddr)) == 0)
+    case kDeviceStateRouter:
+    case kDeviceStateLeader:
+        for (int i = 0; i < kMaxChildren; i++)
         {
-            ExitNow(rval = &mRouters[i]);
+            if (mChildren[i].mState == Neighbor::kStateValid &&
+                memcmp(&mChildren[i].mMacAddr, &aAddress, sizeof(mChildren[i].mMacAddr)) == 0)
+            {
+                ExitNow(rval = &mChildren[i]);
+            }
         }
+
+        for (int i = 0; i < kMaxRouterId; i++)
+        {
+            if (mRouters[i].mState == Neighbor::kStateValid &&
+                memcmp(&mRouters[i].mMacAddr, &aAddress, sizeof(mRouters[i].mMacAddr)) == 0)
+            {
+                ExitNow(rval = &mRouters[i]);
+            }
+        }
+
+        break;
     }
 
 exit: