/* This Source Code Form is subject to the terms of the Mozilla Public

 * License, v. 2.0. If a copy of the MPL was not distributed with this

 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include "mutators.h"

#include <algorithm>

#include <cstddef>

#include <cstdint>

#include <cstring>

#include <random>

#include <vector>

#include "tls_parser.h"

// TLS record header: ContentType(1) + ProtocolVersion(2) + length(2) = 5 bytes.

// DTLS adds epoch(2) + sequence_number(6) before the length field.

// See RFC 8446 Section 5.1 (TLS) and RFC 9147 Section 4 (DTLS).

constexpr size_t kHeaderBytes = 5 + EXTRA_HEADER_BYTES;

constexpr size_t kLengthOffset = 3 + EXTRA_HEADER_BYTES;

// RFC 8446 Section 5.1: change_cipher_spec(20), alert(21),

// handshake(22), application_data(23). Also heartbeat(24, RFC 6520),

// tls12_cid(25, RFC 9146), and ACK(26, RFC 9147 DTLS 1.3).

constexpr uint8_t kContentTypes[] = {20, 21, 22, 23, 24, 25, 26};

struct Record {

  uint8_t* data;

  size_t size;

  size_t remaining;

  uint8_t contentType() { return data[0]; }

};

static std::vector<Record> ParseRecords(uint8_t* data, size_t size) {

  std::vector<Record> records;

  nss_test::TlsParser parser(data, size);

  while (parser.remaining()) {

    size_t offset = parser.consumed();

    if (!parser.Skip(kLengthOffset)) {

      break;

    }

    // Read the 2-byte length field, then skip that many payload bytes.

    if (!parser.SkipVariable(2)) {

      break;

    }

    records.push_back(

        {data + offset, parser.consumed() - offset, parser.remaining()});

  }

  return records;

}

static size_t DropRecord(uint8_t* data, size_t size, size_t maxSize,

                         std::mt19937& rng) {

  std::vector<Record> records = ParseRecords(data, size);

  if (records.empty()) {

    return 0;

  }

  std::uniform_int_distribution<size_t> dist(0, records.size() - 1);

  Record& record = records.at(dist(rng));

  memmove(record.data, record.data + record.size, record.remaining);

  return size - record.size;

}

static size_t SwapRecords(uint8_t* data, size_t size, size_t maxSize,

                          std::mt19937& rng) {

  std::vector<Record> records = ParseRecords(data, size);

  if (records.size() < 2) {

    return 0;

  }

  std::uniform_int_distribution<size_t> dist(0, records.size() - 1);

  size_t i = dist(rng);

  size_t j;

  do {

    j = dist(rng);

  } while (j == i);

  // Read from buf (snapshot), write to data. Do not remove the copy;

  // without it, in-place reordering of variable-size records aliases.

  std::vector<uint8_t> buf(data, data + size);

  std::swap(records.at(i), records.at(j));

  uint8_t* dest = data;

  for (Record& rec : records) {

    memcpy(dest, buf.data() + (rec.data - data), rec.size);

    dest += rec.size;

  }

  return size;

}

// Copy a random record and insert the duplicate before a random record.

static size_t DuplicateRecord(uint8_t* data, size_t size, size_t maxSize,

                              std::mt19937& rng) {

  std::vector<Record> records = ParseRecords(data, size);

  if (records.empty()) {

    return 0;

  }

  std::uniform_int_distribution<size_t> dist(0, records.size() - 1);

  Record& record = records.at(dist(rng));

  if (size + record.size > maxSize) {

    return 0;

  }

  Record& other = records.at(dist(rng));

  memmove(other.data + record.size, other.data, other.size + other.remaining);

  uint8_t* src =

      (record.data < other.data) ? record.data : record.data + record.size;

  memcpy(other.data, src, record.size);

  return size + record.size;

}

// Shorten a random record's payload to a random length.

static size_t TruncateRecord(uint8_t* data, size_t size, size_t maxSize,

                             std::mt19937& rng) {

  std::vector<Record> records = ParseRecords(data, size);

  if (records.empty()) {

    return 0;

  }

  std::uniform_int_distribution<size_t> dist(0, records.size() - 1);

  Record& record = records.at(dist(rng));

  if (record.size <= kHeaderBytes) {

    return 0;

  }

  std::uniform_int_distribution<size_t> dist2(kHeaderBytes, record.size - 1);

  size_t length = dist2(rng);

  size_t payloadLength = length - kHeaderBytes;

  record.data[kLengthOffset] = (payloadLength >> 8) & 0xff;

  record.data[kLengthOffset + 1] = payloadLength & 0xff;

  memmove(record.data + length, record.data + record.size, record.remaining);

  return size + length - record.size;

}

// Split a random record into two records of the same ContentType.

// RFC 8446 Section 5.1: "Multiple fragments of a Handshake message may be

// coalesced into a single TLSPlaintext record" -- the inverse is also valid.

// For DTLS the second record inherits the first's epoch/seqno; real DTLS

// fragments at the handshake layer, but duplicate seqnos exercise anti-replay.

static size_t FragmentRecord(uint8_t* data, size_t size, size_t maxSize,

                             std::mt19937& rng) {

  if (size + kHeaderBytes > maxSize) {

    return 0;

  }

  std::vector<Record> records = ParseRecords(data, size);

  if (records.empty()) {

    return 0;

  }

  std::uniform_int_distribution<size_t> dist(0, records.size() - 1);

  Record& record = records.at(dist(rng));

  size_t contentLength = record.size - kHeaderBytes;

  if (contentLength < 2) {

    return 0;

  }

  std::uniform_int_distribution<size_t> dist2(1, contentLength - 1);

  size_t firstLength = dist2(rng);

  size_t secondLength = contentLength - firstLength;

  // Update first record's length.

  record.data[kLengthOffset] = (firstLength >> 8) & 0xff;

  record.data[kLengthOffset + 1] = firstLength & 0xff;

  // Make room for the second record's header at the split point.

  uint8_t* second = record.data + kHeaderBytes + firstLength;

  memmove(second + kHeaderBytes, second, record.remaining + secondLength);

  // Write second header: copy type/version from first, set new length.

  memcpy(second, record.data, kLengthOffset);

  second[kLengthOffset] = (secondLength >> 8) & 0xff;

  second[kLengthOffset + 1] = secondLength & 0xff;

  return size + kHeaderBytes;

}

// Insert a zero-length record with a random content type before a random

// record. RFC 8446 Section 5.4: "Application Data records may contain a

// zero-length TLSInnerPlaintext.content".

static size_t InjectEmptyRecord(uint8_t* data, size_t size, size_t maxSize,

                                std::mt19937& rng) {

  if (size + kHeaderBytes > maxSize) {

    return 0;

  }

  std::vector<Record> records = ParseRecords(data, size);

  if (records.empty()) {

    return 0;

  }

  std::uniform_int_distribution<size_t> dist(0, records.size() - 1);

  Record& record = records.at(dist(rng));

  uint8_t* dest = record.data;

  memmove(dest + kHeaderBytes, dest, record.size + record.remaining);

  // Copy ProtocolVersion (+ epoch/seqno for DTLS) from the adjacent

  // record, pick a random ContentType, and set length to zero.

  memcpy(dest + 1, dest + kHeaderBytes + 1, kLengthOffset - 1);

  std::uniform_int_distribution<size_t> dist2(0, sizeof(kContentTypes) - 1);

  dest[0] = kContentTypes[dist2(rng)];

  dest[kLengthOffset] = 0;

  dest[kLengthOffset + 1] = 0;

  return size + kHeaderBytes;

}

// Merge two adjacent records that share a ContentType into one.

// Inverse of FragmentRecord. RFC 8446 Section 5.1: "Handshake messages may

// be fragmented over several records".

static size_t MergeRecords(uint8_t* data, size_t size, size_t maxSize,

                           std::mt19937& rng) {

  std::vector<Record> records = ParseRecords(data, size);

  if (records.size() < 2) {

    return 0;

  }

  std::vector<size_t> candidates;

  for (size_t i = 0; i + 1 < records.size(); i++) {

    if (records[i].contentType() == records[i + 1].contentType()) {

      candidates.push_back(i);

    }

  }

  if (candidates.empty()) {

    return 0;

  }

  std::uniform_int_distribution<size_t> dist(0, candidates.size() - 1);

  size_t idx = candidates.at(dist(rng));

  Record& first = records.at(idx);

  Record& second = records.at(idx + 1);

  size_t combinedPayload =

      (first.size - kHeaderBytes) + (second.size - kHeaderBytes);

  // RFC 8446 Section 5.1: TLSPlaintext.length MUST NOT exceed 2^14.

  // We use the uint16 max here to also exercise the overflow path.

  if (combinedPayload > 0xffff) {

    return 0;

  }

  first.data[kLengthOffset] = (combinedPayload >> 8) & 0xff;

  first.data[kLengthOffset + 1] = combinedPayload & 0xff;

  uint8_t* secondHeader = second.data;

  memmove(secondHeader, secondHeader + kHeaderBytes,

          second.size - kHeaderBytes + second.remaining);

  return size - kHeaderBytes;

}

// Extend a random record's payload with random bytes.

// Inverse of TruncateRecord. Exercises overlong-message handling and

// "trailing data after message" code paths.

static size_t ExtendRecord(uint8_t* data, size_t size, size_t maxSize,

                           std::mt19937& rng) {

  std::vector<Record> records = ParseRecords(data, size);

  if (records.empty()) {

    return 0;

  }

  std::uniform_int_distribution<size_t> dist(0, records.size() - 1);

  Record& record = records.at(dist(rng));

  size_t available = maxSize - size;

  if (available == 0) {

    return 0;

  }

  size_t currentPayload = record.size - kHeaderBytes;

  if (currentPayload >= 0xffff) {

    return 0;

  }

  size_t maxExtend = std::min(available, (size_t)0xffff - currentPayload);

  std::uniform_int_distribution<size_t> dist2(1, maxExtend);

  size_t extendBy = dist2(rng);

  // Make room after the record's current payload.

  uint8_t* insertPoint = record.data + record.size;

  memmove(insertPoint + extendBy, insertPoint, record.remaining);

  // Fill extension with random bytes.

  std::uniform_int_distribution<uint8_t> dist3(0, 255);

  for (size_t i = 0; i < extendBy; ++i) {

    insertPoint[i] = dist3(rng);

  }

  // Update the length field.

  size_t newPayload = currentPayload + extendBy;

  record.data[kLengthOffset] = (newPayload >> 8) & 0xff;

  record.data[kLengthOffset + 1] = newPayload & 0xff;

  return size + extendBy;

}

using Mutator = size_t (*)(uint8_t*, size_t, size_t, std::mt19937&);

constexpr Mutator kMutators[] = {

    DropRecord,     SwapRecords,       DuplicateRecord, TruncateRecord,

    FragmentRecord, InjectEmptyRecord, MergeRecords,    ExtendRecord,

};

namespace TlsMutators {

extern "C" size_t LLVMFuzzerMutate(uint8_t* data, size_t size, size_t maxSize);

// Pick a random TLS-aware mutator or fall back to libFuzzer's default.

size_t CustomMutator(uint8_t* data, size_t size, size_t maxSize,

                     unsigned int seed) {

  std::mt19937 rng(seed);

  std::bernoulli_distribution coin;

  if (coin(rng)) {

    std::uniform_int_distribution<size_t> dist(

        0, (sizeof(kMutators) / sizeof(kMutators[0])) - 1);

    return kMutators[dist(rng)](data, size, maxSize, rng);

  }

  return LLVMFuzzerMutate(data, size, maxSize);

}

// Merge records from two transcripts, shuffle, and write to `out`.

// const_cast is safe: records are only read via memcpy to `out`.

size_t CustomCrossOver(const uint8_t* data1, size_t size1, const uint8_t* data2,

                       size_t size2, uint8_t* out, size_t maxOutSize,

                       unsigned int seed) {

  std::vector<Record> records1 =

      ParseRecords(const_cast<uint8_t*>(data1), size1);

  if (records1.empty()) {

    return 0;

  }

  std::vector<Record> records2 =

      ParseRecords(const_cast<uint8_t*>(data2), size2);

  if (records2.empty()) {

    return 0;

  }

  // Append `records2` to the back of `records1`.

  std::move(records2.begin(), records2.end(), std::back_inserter(records1));

  std::mt19937 rng(seed);

  std::shuffle(records1.begin(), records1.end(), rng);

  size_t total = 0;

  for (Record& record : records1) {

    size_t length = record.size;

    if (total + length > maxOutSize) {

      break;

    }

    memcpy(out + total, record.data, length);

    total += length;

  }

  return total;

}

}  // namespace TlsMutators