From bbec8132e4aed8b16df7a1c977d72f351a025291 Mon Sep 17 00:00:00 2001
From: rongli <rongli@nestlabs.com>
Date: Wed, 25 Jan 2017 11:06:26 +0800
Subject: [PATCH] enable mac security for fragmented mle message (#1202)

---
 src/core/common/message.cpp        | 15 +++++++++++++++
 src/core/common/message.hpp        | 10 ++++++++++
 src/core/thread/mesh_forwarder.cpp | 18 ++++++++++++++++--
 src/core/thread/mle.cpp            |  1 +
 4 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/src/core/common/message.cpp b/src/core/common/message.cpp
index f7a80d3b6..1f2232957 100644
--- a/src/core/common/message.cpp
+++ b/src/core/common/message.cpp
@@ -372,6 +372,21 @@ void Message::SetSubType(uint8_t aSubType)
     mInfo.mSubType = aSubType;
 }
 
+bool Message::IsSubTypeMle(void) const
+{
+    bool rval = false;
+
+    if (mInfo.mSubType == kSubTypeMleAnnounce ||
+        mInfo.mSubType == kSubTypeMleDiscoverRequest ||
+        mInfo.mSubType == kSubTypeMleDiscoverResponse ||
+        mInfo.mSubType == kSubTypeMleGeneral)
+    {
+        rval = true;
+    }
+
+    return rval;
+}
+
 uint8_t Message::GetPriority(void) const
 {
     return mInfo.mPriority;
diff --git a/src/core/common/message.hpp b/src/core/common/message.hpp
index 3af12b25d..f6dc879bf 100644
--- a/src/core/common/message.hpp
+++ b/src/core/common/message.hpp
@@ -215,6 +215,7 @@ public:
         kSubTypeMleDiscoverResponse = 3,  ///< MLE Discover Response
         kSubTypeJoinerEntrust       = 4,  ///< Joiner Entrust
         kSubTypeMplRetransmission   = 5,  ///< MPL next retranmission message
+        kSubTypeMleGeneral          = 6,  ///< General MLE
     };
 
     enum
@@ -320,6 +321,15 @@ public:
      */
     void SetSubType(uint8_t aSubType);
 
+    /**
+     * This method returns whether or not the message is of MLE subtype.
+     *
+     * @retval TRUE   If message is of MLE subtype.
+     * @retval FLASE  If message is not of MLE subtype.
+     *
+     */
+    bool IsSubTypeMle(void) const;
+
     /**
      * This method returns the message priority level.
      *
diff --git a/src/core/thread/mesh_forwarder.cpp b/src/core/thread/mesh_forwarder.cpp
index 02c955ece..a82d7dd74 100644
--- a/src/core/thread/mesh_forwarder.cpp
+++ b/src/core/thread/mesh_forwarder.cpp
@@ -989,7 +989,11 @@ ThreadError MeshForwarder::HandleFrameRequest(Mac::Frame &aFrame)
             aFrame.SetChannel(mScanChannel);
         }
 
-        SendFragment(*mSendMessage, aFrame);
+        if (SendFragment(*mSendMessage, aFrame) == kThreadError_NotCapable)
+        {
+            SendFragment(*mSendMessage, aFrame);
+        }
+
         assert(aFrame.GetLength() != 7);
         break;
 
@@ -1115,6 +1119,7 @@ ThreadError MeshForwarder::SendFragment(Message &aMessage, Mac::Frame &aFrame)
     uint16_t fragmentLength;
     uint16_t dstpan;
     uint8_t secCtl = Mac::Frame::kSecNone;
+    ThreadError error = kThreadError_None;
 
     if (mAddMeshHeader)
     {
@@ -1254,6 +1259,13 @@ ThreadError MeshForwarder::SendFragment(Message &aMessage, Mac::Frame &aFrame)
 
         if (payloadLength > fragmentLength)
         {
+            if ((!aMessage.IsLinkSecurityEnabled()) && aMessage.IsSubTypeMle())
+            {
+                aMessage.SetLinkSecurityEnabled(true);
+                aMessage.SetOffset(0);
+                ExitNow(error = kThreadError_NotCapable);
+            }
+
             // write Fragment header
             if (aMessage.GetDatagramTag() == 0)
             {
@@ -1322,7 +1334,9 @@ ThreadError MeshForwarder::SendFragment(Message &aMessage, Mac::Frame &aFrame)
         aFrame.SetFramePending(true);
     }
 
-    return kThreadError_None;
+exit:
+
+    return error;
 }
 
 ThreadError MeshForwarder::SendEmptyFrame(Mac::Frame &aFrame)
diff --git a/src/core/thread/mle.cpp b/src/core/thread/mle.cpp
index e0fc7091b..f85a6fa29 100644
--- a/src/core/thread/mle.cpp
+++ b/src/core/thread/mle.cpp
@@ -882,6 +882,7 @@ Message *Mle::NewMleMessage(void)
     message = mSocket.NewMessage(0);
     VerifyOrExit(message != NULL, ;);
 
+    message->SetSubType(Message::kSubTypeMleGeneral);
     message->SetLinkSecurityEnabled(false);
     message->SetPriority(kMleMessagePriority);