[sub-mac] ensure AES key is set before restoring transmit security on retx (#13378)

This commit introduces `SubMac::ReprocessSecurityForRetx()` to handle
re-processing transmit security for frames with Header IEs during
retransmission.

When transmit security is delegated to the radio platform, the radio is
not required to set or preserve the AES key on `TxFrame`. While it may be
expected that a radio platform implementation might do this, it is not
guaranteed. This raised a risk in the previous implementation where calling
`RestoreTransmitSecurity()` to decrypt the frame back to plaintext could
potentially use invalid key material.

This commit resolves this issue by determining and setting the AES key on
`TxFrame` from its key index before restoring transmit security.

src/core/mac/sub_mac.cpp

@@ -577,12 +577,7 @@ void SubMac::HandleTransmitDone(TxFrame &aFrame, RxFrame *aAckFrame, Error aErro
         aFrame.SetIsARetransmission(true);
 
 #if OPENTHREAD_CONFIG_MAC_HEADER_IE_SUPPORT && OPENTHREAD_CONFIG_MAC_SOFTWARE_RETX_SECURITY_ENABLE
-        if (aFrame.GetSecurityEnabled() && aFrame.HasAnyHeaderIe())
-        {
-            aFrame.RestoreTransmitSecurity(GetExtAddress());
-        }
-
-        ProcessTransmitSecurity();
+        ReprocessSecurityForRetx(aFrame);
 #endif
 
 #if OPENTHREAD_CONFIG_MAC_ADD_DELAY_ON_NO_ACK_ERROR_BEFORE_RETRY
@@ -621,6 +616,41 @@ exit:
     return;
 }
 
+#if OPENTHREAD_CONFIG_MAC_HEADER_IE_SUPPORT && OPENTHREAD_CONFIG_MAC_SOFTWARE_RETX_SECURITY_ENABLE
+
+void SubMac::ReprocessSecurityForRetx(TxFrame &aFrame)
+{
+    // Re-processes transmit security on a frame being retransmitted if
+    // it contains Header IEs. The frame is first restored back to
+    // plaintext and then re-encrypted with a new frame counter value.
+
+    VerifyOrExit(aFrame.GetSecurityEnabled() && aFrame.HasAnyHeaderIe());
+
+    // When transmit security is handled by `SubMac`, the AES key is already set
+    // on `aFrame`. However, when transmit security is delegated to the radio
+    // platform, the radio is not required to set or preserve the AES key on
+    // `aFrame`. To ensure `RestoreTransmitSecurity()` can properly decrypt the
+    // frame back to plaintext, we determine and set the key on `aFrame` using
+    // its key index.
+
+    if (!ShouldHandleTransmitSecurity())
+    {
+        uint8_t keyIndex;
+
+        SuccessOrExit(aFrame.GetKeyIndex(keyIndex));
+        aFrame.SetAesKey(mKeyTrio.SelectKey(keyIndex));
+    }
+
+    aFrame.RestoreTransmitSecurity(GetExtAddress());
+
+    ProcessTransmitSecurity();
+
+exit:
+    return;
+}
+
+#endif // OPENTHREAD_CONFIG_MAC_HEADER_IE_SUPPORT && OPENTHREAD_CONFIG_MAC_SOFTWARE_RETX_SECURITY_ENABLE
+
 void SubMac::SignalFrameCounterUsedOnTxDone(const TxFrame &aFrame)
 {
     uint8_t  keyIdMode;

src/core/mac/sub_mac.hpp

@@ -585,6 +585,7 @@ private:
     bool ShouldHandleTransitionToSleep(void) const;
 
     void ProcessTransmitSecurity(void);
+    void ReprocessSecurityForRetx(TxFrame &aFrame);
     void SignalFrameCounterUsed(uint32_t aFrameCounter, uint8_t aKeyIndex);
     void StartCsmaBackoff(void);
     void StartTimerForBackoff(uint8_t aBackoffExponent);