NASHelperNGAP_RewriteIdentities.cpp

#include "NASHelperNGAP.hpp"

#include <algorithm>
#include <cstddef>
#include <vector>

namespace {

bool fillRange(std::vector<std::uint8_t>& pdu,
               const std::uint8_t* begin,
               std::size_t len,
               std::uint8_t fill,
               std::size_t& totalWritten) {
    if (begin == nullptr || len == 0U) {
        return false;
    }
    const std::uint8_t* const base = pdu.data();
    const std::uint8_t* const end = base + pdu.size();
    if (begin < base || begin + len > end) {
        return false;
    }
    const std::size_t off = static_cast<std::size_t>(begin - base);
    std::fill(pdu.begin() + static_cast<std::ptrdiff_t>(off),
              pdu.begin() + static_cast<std::ptrdiff_t>(off + len), fill);
    totalWritten += len;
    return true;
}

/// Fills a 5GS mobile identity **value** (TS 24.501 §9.11.3.4): the octet after the length has the **upper**
/// nibble as the first identity digit(s) and the **lower** nibble for type of identity (bits 1–3) and odd/even
/// (bit 4). Preserves @c (firstOctet & 0x0F); sets the upper nibble from @p fill (same as @c fillRange).
/// Remaining value octets are filled entirely with @p fill. @return false only if @p valueBegin is null, length
/// is zero, or the span is out of range.
bool fillFiveGsMobileIdentityValuePreservingTypeOctet(std::vector<std::uint8_t>& pdu,
                                                       const std::uint8_t* valueBegin,
                                                       std::size_t valueLen,
                                                       std::uint8_t fill,
                                                       std::size_t& totalWritten) {
    if (valueBegin == nullptr || valueLen == 0U) {
        return false;
    }
    const std::uint8_t* const base = pdu.data();
    const std::uint8_t* const end = base + pdu.size();
    if (valueBegin < base || valueBegin + valueLen > end) {
        return false;
    }
    const std::size_t firstOff = static_cast<std::size_t>(valueBegin - base);
    const std::uint8_t orig0 = pdu[firstOff];
    const std::uint8_t new0 = static_cast<std::uint8_t>((orig0 & 0x0FU) | (fill & 0xF0U));
    if (new0 != orig0) {
        pdu[firstOff] = new0;
        totalWritten += 1U;
    }
    if (valueLen <= 1U) {
        return true;
    }
    return fillRange(pdu, valueBegin + 1U, valueLen - 1U, fill, totalWritten);
}

void rewriteSlice(std::vector<std::uint8_t>& pdu,
                  std::size_t off,
                  std::size_t len,
                  std::uint8_t fill,
                  NASHelperNGAP& secForProtectedPdu,
                  bool innerCipheredPlainLayout,
                  std::size_t& totalWritten);

void rewritePlain5gmmFromPrefix(std::vector<std::uint8_t>& pdu,
                                const NASHelperNGAP::PlainNasPrefix& plain,
                                std::uint8_t fill,
                                std::size_t& totalWritten);

// --- Registration request optional walk (must match NASHelperNGAP_RegistrationRequest.cpp) ---

bool registrationRequestOptionalIsTlvE(std::uint8_t iei) {
    switch (iei) {
        case 0x77:
        case 0x70:
        case 0x74:
        case 0x7B:
        case 0x71:
        case 0x72:
            return true;
        default:
            return false;
    }
}

bool registrationRequestOptionalIsTlv(std::uint8_t iei) {
    switch (iei) {
        case 0x10:
        case 0x2E:
        case 0x2F:
        case 0x52:
        case 0x17:
        case 0x40:
        case 0x50:
        case 0x2B:
        case 0x25:
        case 0x18:
        case 0x51:
        case 0x53:
        case 0x41:
        case 0x42:
        case 0x60:
        case 0x6E:
        case 0x6A:
        case 0x67:
        case 0x35:
        case 0x48:
        case 0x1A:
        case 0x30:
        case 0x29:
        case 0x28:
        case 0x32:
        case 0x16:
        case 0x2A:
        case 0x3B:
        case 0x3C:
        case 0x3F:
        case 0x56:
        case 0x64:
            return true;
        default:
            return false;
    }
}

bool registrationRequestOptionalTv1HighNibble(std::uint8_t firstOctet) {
    const unsigned h = firstOctet >> 4U;
    return h >= 8U && h <= 12U;
}

// --- Registration accept optional walk (must match NASHelperNGAP_RegistrationAccept.cpp) ---

bool registrationAcceptOptionalIsTlvE(std::uint8_t iei) {
    switch (iei) {
        case 0x77:
        case 0x72:
        case 0x79:
        case 0x7A:
        case 0x73:
        case 0x78:
        case 0x76:
        case 0x74:
        case 0x75:
        case 0x7B:
        case 0x70:
        case 0x71:
        case 0x7C:
        case 0x7D:
            return true;
        default:
            return false;
    }
}

bool registrationAcceptOptionalIsTlv(std::uint8_t iei) {
    switch (iei) {
        case 0x4A:
        case 0x54:
        case 0x15:
        case 0x11:
        case 0x31:
        case 0x21:
        case 0x50:
        case 0x26:
        case 0x27:
        case 0x5E:
        case 0x5D:
        case 0x16:
        case 0x34:
        case 0x51:
        case 0x60:
        case 0x6E:
        case 0x6C:
        case 0x6B:
        case 0x6A:
        case 0x67:
        case 0x39:
        case 0x1B:
        case 0x1C:
        case 0x29:
        case 0x68:
        case 0x33:
        case 0x35:
        case 0x14:
        case 0x2C:
        case 0x13:
        case 0x1D:
        case 0x1E:
        case 0x3D:
        case 0x32:
        case 0x4B:
        case 0x4C:
        case 0x4F:
        case 0x5B:
        case 0x3C:
        case 0x5C:
        case 0x61:
        case 0x63:
        case 0x64:
            return true;
        default:
            return false;
    }
}

bool registrationAcceptOptionalTv1HighNibble(std::uint8_t firstOctet) {
    const unsigned h = firstOctet >> 4U;
    if (h >= 8U && h <= 11U) {
        return true;
    }
    return h == 13U || h == 14U;
}

bool securityModeCompleteOptionalIsTlvE(std::uint8_t iei) {
    switch (iei) {
        case 0x77:
        case 0x71:
        case 0x78:
            return true;
        default:
            return false;
    }
}

void rewritePlain5gmmFromPrefix(std::vector<std::uint8_t>& pdu,
                                const NASHelperNGAP::PlainNasPrefix& plain,
                                std::uint8_t fill,
                                std::size_t& totalWritten) {
    if (!NASHelperNGAP::is5gsMobilityManagementNas(plain.extended_protocol_discriminator)) {
        return;
    }
    if (NASHelperNGAP::collectMobileIdentifierLabels(plain).empty()) {
        return;
    }

    const std::uint8_t* const ie = plain.information_elements;
    const std::size_t n = plain.information_element_octet_count;

    switch (plain.message_type) {
        case NASHelperNGAP::k5gmmMessageTypeRegistrationRequest: {
            NASHelperNGAP::RegistrationRequestInformationElementsParse reg{};
            if (!NASHelperNGAP::tryParseRegistrationRequestInformationElements(ie, n, reg)) {
                break;
            }
            fillFiveGsMobileIdentityValuePreservingTypeOctet(
                pdu, reg.mandatory.five_gs_mobile_identity_value,
                reg.mandatory.five_gs_mobile_identity_value_octet_count, fill, totalWritten);
            std::size_t i = reg.mandatory.mandatory_total_octets;
            while (i < n) {
                const std::uint8_t b = ie[i];
                if (registrationRequestOptionalIsTlvE(b)) {
                    if (i + 3U > n) {
                        break;
                    }
                    const std::uint16_t vlen = static_cast<std::uint16_t>(
                        (static_cast<unsigned>(ie[i + 1]) << 8U) | static_cast<unsigned>(ie[i + 2]));
                    const std::size_t chunk = 3U + static_cast<std::size_t>(vlen);
                    if (i + chunk > n || i + chunk < i) {
                        break;
                    }
                    const std::uint8_t* vptr = ie + i + 3U;
                    if (b == 0x77U) {
                        fillFiveGsMobileIdentityValuePreservingTypeOctet(pdu, vptr, static_cast<std::size_t>(vlen),
                                                                         fill, totalWritten);
                    } else if (b == NASHelperNGAP::kIeiNasMessageContainer) {
                        const std::size_t voff = static_cast<std::size_t>(vptr - pdu.data());
                        NASHelperNGAP emptyNested{};
                        rewriteSlice(pdu, voff, static_cast<std::size_t>(vlen), fill, emptyNested, false,
                                     totalWritten);
                    }
                    i += chunk;
                    continue;
                }
                if (registrationRequestOptionalIsTlv(b)) {
                    if (i + 2U > n) {
                        break;
                    }
                    const std::uint8_t l = ie[i + 1];
                    const std::size_t chunk = 2U + static_cast<std::size_t>(l);
                    if (i + chunk > n || i + chunk < i) {
                        break;
                    }
                    i += chunk;
                    continue;
                }
                if (registrationRequestOptionalTv1HighNibble(b)) {
                    i += 1U;
                    continue;
                }
                break;
            }
            break;
        }
        case NASHelperNGAP::k5gmmMessageTypeRegistrationAccept: {
            NASHelperNGAP::RegistrationAcceptInformationElementsParse acc{};
            if (!NASHelperNGAP::tryParseRegistrationAcceptInformationElements(ie, n, acc)) {
                break;
            }
            std::size_t i = acc.mandatory.mandatory_total_octets;
            while (i < n) {
                const std::uint8_t b = ie[i];
                if (registrationAcceptOptionalIsTlvE(b)) {
                    if (i + 3U > n) {
                        break;
                    }
                    const std::uint16_t vlen = static_cast<std::uint16_t>(
                        (static_cast<unsigned>(ie[i + 1]) << 8U) | static_cast<unsigned>(ie[i + 2]));
                    const std::size_t chunk = 3U + static_cast<std::size_t>(vlen);
                    if (i + chunk > n || i + chunk < i) {
                        break;
                    }
                    if (b == 0x77U) {
                        fillFiveGsMobileIdentityValuePreservingTypeOctet(pdu, ie + i + 3U,
                                                                         static_cast<std::size_t>(vlen), fill,
                                                                         totalWritten);
                    }
                    i += chunk;
                    continue;
                }
                if (registrationAcceptOptionalIsTlv(b)) {
                    if (i + 2U > n) {
                        break;
                    }
                    const std::uint8_t l = ie[i + 1];
                    const std::size_t chunk = 2U + static_cast<std::size_t>(l);
                    if (i + chunk > n || i + chunk < i) {
                        break;
                    }
                    i += chunk;
                    continue;
                }
                if (registrationAcceptOptionalTv1HighNibble(b)) {
                    i += 1U;
                    continue;
                }
                break;
            }
            break;
        }
        case NASHelperNGAP::k5gmmMessageTypeDeregistrationRequestUeOriginating: {
            NASHelperNGAP::DeregistrationRequestUeOriginatingInformationElementsParse d{};
            if (!NASHelperNGAP::tryParseDeregistrationRequestUeOriginatingInformationElements(ie, n, d)) {
                break;
            }
            fillFiveGsMobileIdentityValuePreservingTypeOctet(
                pdu, d.mandatory.five_gs_mobile_identity_value,
                d.mandatory.five_gs_mobile_identity_value_octet_count, fill, totalWritten);
            std::size_t i = d.mandatory.mandatory_total_octets;
            while (i < n) {
                const std::uint8_t b = ie[i];
                if (b == 0x71U) {
                    if (i + 3U > n) {
                        break;
                    }
                    const std::uint16_t vlen = static_cast<std::uint16_t>(
                        (static_cast<unsigned>(ie[i + 1]) << 8U) | static_cast<unsigned>(ie[i + 2]));
                    const std::size_t chunk = 3U + static_cast<std::size_t>(vlen);
                    if (i + chunk > n || i + chunk < i) {
                        break;
                    }
                    const std::uint8_t* vptr = ie + i + 3U;
                    const std::size_t voff = static_cast<std::size_t>(vptr - pdu.data());
                    NASHelperNGAP emptyNested{};
                    rewriteSlice(pdu, voff, static_cast<std::size_t>(vlen), fill, emptyNested, false, totalWritten);
                    i += chunk;
                    continue;
                }
                if (b == 0x3CU) {
                    if (i + 2U > n) {
                        break;
                    }
                    const std::uint8_t l = ie[i + 1];
                    const std::size_t chunk = 2U + static_cast<std::size_t>(l);
                    if (i + chunk > n || i + chunk < i) {
                        break;
                    }
                    i += chunk;
                    continue;
                }
                break;
            }
            break;
        }
        case NASHelperNGAP::k5gmmMessageTypeServiceRequest: {
            NASHelperNGAP::ServiceRequestInformationElementsParse sr{};
            if (!NASHelperNGAP::tryParseServiceRequestInformationElements(ie, n, sr)) {
                break;
            }
            fillFiveGsMobileIdentityValuePreservingTypeOctet(
                pdu, sr.mandatory.five_g_s_tmsi_value, sr.mandatory.five_g_s_tmsi_value_octet_count, fill,
                totalWritten);
            break;
        }
        case NASHelperNGAP::k5gmmMessageTypeIdentityResponse: {
            NASHelperNGAP::IdentityResponseInformationElementsParse idr{};
            if (!NASHelperNGAP::tryParseIdentityResponseInformationElements(ie, n, idr)) {
                break;
            }
            fillFiveGsMobileIdentityValuePreservingTypeOctet(pdu, idr.mobile_identity_value,
                                                             idr.mobile_identity_value_octet_count, fill,
                                                             totalWritten);
            break;
        }
        case NASHelperNGAP::k5gmmMessageTypeSecurityModeComplete: {
            NASHelperNGAP::SecurityModeCompleteInformationElementsParse smc{};
            if (!NASHelperNGAP::tryParseSecurityModeCompleteInformationElements(ie, n, smc)) {
                break;
            }
            std::size_t i = 0U;
            while (i < n) {
                const std::uint8_t b = ie[i];
                if (!securityModeCompleteOptionalIsTlvE(b)) {
                    break;
                }
                if (i + 3U > n) {
                    break;
                }
                const std::uint16_t vlen = static_cast<std::uint16_t>(
                    (static_cast<unsigned>(ie[i + 1]) << 8U) | static_cast<unsigned>(ie[i + 2]));
                const std::size_t chunk = 3U + static_cast<std::size_t>(vlen);
                if (i + chunk > n || i + chunk < i) {
                    break;
                }
                const std::uint8_t* vptr = ie + i + 3U;
                if (b == 0x77U || b == 0x78U) {
                    fillFiveGsMobileIdentityValuePreservingTypeOctet(pdu, vptr, static_cast<std::size_t>(vlen), fill,
                                                                     totalWritten);
                } else if (b == NASHelperNGAP::kIeiNasMessageContainer) {
                    const std::size_t voff = static_cast<std::size_t>(vptr - pdu.data());
                    NASHelperNGAP emptyNested{};
                    rewriteSlice(pdu, voff, static_cast<std::size_t>(vlen), fill, emptyNested, false, totalWritten);
                }
                i += chunk;
            }
            break;
        }
        default:
            break;
    }
}

void rewriteSlice(std::vector<std::uint8_t>& pdu,
                  std::size_t off,
                  std::size_t len,
                  std::uint8_t fill,
                  NASHelperNGAP& secForProtectedPdu,
                  bool innerCipheredPlainLayout,
                  std::size_t& totalWritten) {
    if (off > pdu.size() || len > pdu.size() - off) {
        return;
    }
    std::uint8_t* const slice = pdu.data() + off;

    NASHelperNGAP::PlainNasPrefix plain{};
    if (NASHelperNGAP::tryParsePlainNasPrefix(slice, len, plain, innerCipheredPlainLayout)) {
        rewritePlain5gmmFromPrefix(pdu, plain, fill, totalWritten);
        return;
    }

    NASHelperNGAP::SecurityProtectedNasEnvelope sec{};
    if (secForProtectedPdu.tryParseSecurityProtectedNasEnvelope(slice, len, sec)) {
        const std::size_t innerOff = static_cast<std::size_t>(sec.inner_nas_octets - pdu.data());
        NASHelperNGAP empty{};
        rewriteSlice(pdu, innerOff, sec.inner_nas_octet_count, fill, empty, true, totalWritten);
        return;
    }

    NASHelperNGAP emptyOuter{};
    if (emptyOuter.tryParseSecurityProtectedNasEnvelope(slice, len, sec)) {
        const std::size_t innerOff = static_cast<std::size_t>(sec.inner_nas_octets - pdu.data());
        rewriteSlice(pdu, innerOff, sec.inner_nas_octet_count, fill, emptyOuter, true, totalWritten);
    }
}

}  // namespace

std::size_t NASHelperNGAP::rewriteKnownMobileIdentifiers(std::vector<std::uint8_t>& pdu,
                                                         NASHelperNGAP& securityContextForOuterProtected,
                                                         const NasIdentifierRewriteOptions& options) {
    if (pdu.empty()) {
        return 0U;
    }

    std::size_t totalWritten = 0U;
    PlainNasPrefix plain{};
    if (tryParsePlainNasPrefix(pdu.data(), pdu.size(), plain, false)) {
        rewritePlain5gmmFromPrefix(pdu, plain, options.fill_byte, totalWritten);
        return totalWritten;
    }

    SecurityProtectedNasEnvelope sec{};
    if (securityContextForOuterProtected.tryParseSecurityProtectedNasEnvelope(pdu.data(), pdu.size(), sec)) {
        const std::size_t innerOff = static_cast<std::size_t>(sec.inner_nas_octets - pdu.data());
        NASHelperNGAP empty{};
        rewriteSlice(pdu, innerOff, sec.inner_nas_octet_count, options.fill_byte, empty, true, totalWritten);
        return totalWritten;
    }

    NASHelperNGAP emptyTop{};
    if (emptyTop.tryParseSecurityProtectedNasEnvelope(pdu.data(), pdu.size(), sec)) {
        const std::size_t innerOff = static_cast<std::size_t>(sec.inner_nas_octets - pdu.data());
        rewriteSlice(pdu, innerOff, sec.inner_nas_octet_count, options.fill_byte, emptyTop, true, totalWritten);
        return totalWritten;
    }

    return 0U;
}