//  Copyright (c) 2020-present, Facebook, Inc.  All rights reserved.
//  This source code is licensed under both the GPLv2 (found in the
//  COPYING file in the root directory) and Apache 2.0 License
//  (found in the LICENSE.Apache file in the root directory).
//
// This file contains classes containing fields to protect individual entries.
// The classes are named "ProtectionInfo<suffix>", where <suffix> indicates the
// combination of fields that are covered. Each field has a single letter
// abbreviation as follows.
//
// K = key
// V = value
// O = optype aka value type
// S = seqno
// C = CF ID
//
// Then, for example, a class that protects an entry consisting of key, value,
// optype, and CF ID (i.e., a `WriteBatch` entry) would be named
// `ProtectionInfoKVOC`.
//
// The `ProtectionInfo.*` classes are templated on the integer type used to hold
// the XOR of hashes for each field. Only unsigned integer types are supported,
// and the maximum supported integer width is 64 bits. When the integer type is
// narrower than the hash values, we lop off the most significant bits to make
// them fit.
//
// The `ProtectionInfo.*` classes are all intended to be non-persistent. We do
// not currently make the byte order consistent for integer fields before
// hashing them, so the resulting values are endianness-dependent.

#pragma once

#include <type_traits>

#include "db/dbformat.h"
#include "rocksdb/types.h"
#include "util/hash.h"

namespace ROCKSDB_NAMESPACE {

template <typename T>
class ProtectionInfo;
template <typename T>
class ProtectionInfoKVO;
template <typename T>
class ProtectionInfoKVOC;
template <typename T>
class ProtectionInfoKVOS;
template <typename T>
class ProtectionInfoKV;

// Aliases for 64-bit protection infos.
using ProtectionInfo64 = ProtectionInfo<uint64_t>;
using ProtectionInfoKVO64 = ProtectionInfoKVO<uint64_t>;
using ProtectionInfoKVOC64 = ProtectionInfoKVOC<uint64_t>;
using ProtectionInfoKVOS64 = ProtectionInfoKVOS<uint64_t>;

template <typename T>
class ProtectionInfo {
 public:
  ProtectionInfo() = default;

  Status GetStatus() const;
  ProtectionInfoKVO<T> ProtectKVO(const Slice& key, const Slice& value,
                                  ValueType op_type) const;
  ProtectionInfoKVO<T> ProtectKVO(const SliceParts& key,
                                  const SliceParts& value,
                                  ValueType op_type) const;
  ProtectionInfoKV<T> ProtectKV(const Slice& key, const Slice& value) const;

 private:
  friend class ProtectionInfoKVO<T>;
  friend class ProtectionInfoKVOS<T>;
  friend class ProtectionInfoKVOC<T>;
  friend class ProtectionInfoKV<T>;

  // Each field is hashed with an independent value so we can catch fields being
  // swapped. Per the `NPHash64()` docs, using consecutive seeds is a pitfall,
  // and we should instead vary our seeds by a large odd number. This value by
  // which we increment (0xD28AAD72F49BD50B) was taken from
  // `head -c8 /dev/urandom | hexdump`, run repeatedly until it yielded an odd
  // number. The values are computed manually since the Windows C++ compiler
  // complains about the overflow when adding constants.
  static const uint64_t kSeedK = 0;
  static const uint64_t kSeedV = 0xD28AAD72F49BD50B;
  static const uint64_t kSeedO = 0xA5155AE5E937AA16;
  static const uint64_t kSeedS = 0x77A00858DDD37F21;
  static const uint64_t kSeedC = 0x4A2AB5CBD26F542C;

  ProtectionInfo(T val) : val_(val) {
    static_assert(sizeof(ProtectionInfo<T>) == sizeof(T), "");
  }

  T GetVal() const { return val_; }
  void SetVal(T val) { val_ = val; }

  void Encode(uint8_t len, char* dst) const {
    assert(sizeof(val_) >= len);
    switch (len) {
      case 1:
        dst[0] = static_cast<uint8_t>(val_);
        break;
      case 2:
        EncodeFixed16(dst, static_cast<uint16_t>(val_));
        break;
      case 4:
        EncodeFixed32(dst, static_cast<uint32_t>(val_));
        break;
      case 8:
        EncodeFixed64(dst, static_cast<uint64_t>(val_));
        break;
      default:
        assert(false);
    }
  }

  bool Verify(uint8_t len, const char* checksum_ptr) const {
    assert(sizeof(val_) >= len);
    switch (len) {
      case 1:
        return static_cast<uint8_t>(checksum_ptr[0]) ==
               static_cast<uint8_t>(val_);
      case 2:
        return DecodeFixed16(checksum_ptr) == static_cast<uint16_t>(val_);
      case 4:
        return DecodeFixed32(checksum_ptr) == static_cast<uint32_t>(val_);
      case 8:
        return DecodeFixed64(checksum_ptr) == static_cast<uint64_t>(val_);
      default:
        assert(false);
        return false;
    }
  }

  T val_ = 0;
};

template <typename T>
class ProtectionInfoKVO {
 public:
  ProtectionInfoKVO() = default;

  ProtectionInfo<T> StripKVO(const Slice& key, const Slice& value,
                             ValueType op_type) const;
  ProtectionInfo<T> StripKVO(const SliceParts& key, const SliceParts& value,
                             ValueType op_type) const;

  ProtectionInfoKVOC<T> ProtectC(ColumnFamilyId column_family_id) const;
  ProtectionInfoKVOS<T> ProtectS(SequenceNumber sequence_number) const;

  void UpdateK(const Slice& old_key, const Slice& new_key);
  void UpdateK(const SliceParts& old_key, const SliceParts& new_key);
  void UpdateV(const Slice& old_value, const Slice& new_value);
  void UpdateV(const SliceParts& old_value, const SliceParts& new_value);
  void UpdateO(ValueType old_op_type, ValueType new_op_type);

  // Encode this protection info into `len` bytes and stores them in `dst`.
  void Encode(uint8_t len, char* dst) const { info_.Encode(len, dst); }
  // Verify this protection info against the protection info encoded by Encode()
  // at the first `len` bytes of `checksum_ptr`.
  // Returns true iff the verification is successful.
  bool Verify(uint8_t len, const char* checksum_ptr) const {
    return info_.Verify(len, checksum_ptr);
  }

 private:
  friend class ProtectionInfo<T>;
  friend class ProtectionInfoKVOS<T>;
  friend class ProtectionInfoKVOC<T>;

  explicit ProtectionInfoKVO(T val) : info_(val) {
    static_assert(sizeof(ProtectionInfoKVO<T>) == sizeof(T), "");
  }

  T GetVal() const { return info_.GetVal(); }
  void SetVal(T val) { info_.SetVal(val); }

  ProtectionInfo<T> info_;
};

template <typename T>
class ProtectionInfoKVOC {
 public:
  ProtectionInfoKVOC() = default;

  ProtectionInfoKVO<T> StripC(ColumnFamilyId column_family_id) const;

  void UpdateK(const Slice& old_key, const Slice& new_key) {
    kvo_.UpdateK(old_key, new_key);
  }
  void UpdateK(const SliceParts& old_key, const SliceParts& new_key) {
    kvo_.UpdateK(old_key, new_key);
  }
  void UpdateV(const Slice& old_value, const Slice& new_value) {
    kvo_.UpdateV(old_value, new_value);
  }
  void UpdateV(const SliceParts& old_value, const SliceParts& new_value) {
    kvo_.UpdateV(old_value, new_value);
  }
  void UpdateO(ValueType old_op_type, ValueType new_op_type) {
    kvo_.UpdateO(old_op_type, new_op_type);
  }
  void UpdateC(ColumnFamilyId old_column_family_id,
               ColumnFamilyId new_column_family_id);

  void Encode(uint8_t len, char* dst) const { kvo_.Encode(len, dst); }
  bool Verify(uint8_t len, const char* checksum_ptr) const {
    return kvo_.Verify(len, checksum_ptr);
  }

 private:
  friend class ProtectionInfoKVO<T>;

  explicit ProtectionInfoKVOC(T val) : kvo_(val) {
    static_assert(sizeof(ProtectionInfoKVOC<T>) == sizeof(T), "");
  }

  T GetVal() const { return kvo_.GetVal(); }
  void SetVal(T val) { kvo_.SetVal(val); }

  ProtectionInfoKVO<T> kvo_;
};

template <typename T>
class ProtectionInfoKVOS {
 public:
  ProtectionInfoKVOS() = default;

  ProtectionInfoKVO<T> StripS(SequenceNumber sequence_number) const;

  void UpdateK(const Slice& old_key, const Slice& new_key) {
    kvo_.UpdateK(old_key, new_key);
  }
  void UpdateK(const SliceParts& old_key, const SliceParts& new_key) {
    kvo_.UpdateK(old_key, new_key);
  }
  void UpdateV(const Slice& old_value, const Slice& new_value) {
    kvo_.UpdateV(old_value, new_value);
  }
  void UpdateV(const SliceParts& old_value, const SliceParts& new_value) {
    kvo_.UpdateV(old_value, new_value);
  }
  void UpdateO(ValueType old_op_type, ValueType new_op_type) {
    kvo_.UpdateO(old_op_type, new_op_type);
  }
  void UpdateS(SequenceNumber old_sequence_number,
               SequenceNumber new_sequence_number);

  void Encode(uint8_t len, char* dst) const { kvo_.Encode(len, dst); }
  bool Verify(uint8_t len, const char* checksum_ptr) const {
    return kvo_.Verify(len, checksum_ptr);
  }

 private:
  friend class ProtectionInfoKVO<T>;

  explicit ProtectionInfoKVOS(T val) : kvo_(val) {
    static_assert(sizeof(ProtectionInfoKVOS<T>) == sizeof(T), "");
  }

  T GetVal() const { return kvo_.GetVal(); }
  void SetVal(T val) { kvo_.SetVal(val); }

  ProtectionInfoKVO<T> kvo_;
};

template <typename T>
class ProtectionInfoKV {
 public:
  ProtectionInfoKV() = default;

  void Encode(uint8_t len, char* dst) const { info_.Encode(len, dst); }
  bool Verify(uint8_t len, const char* checksum_ptr) const {
    return info_.Verify(len, checksum_ptr);
  }

 private:
  friend class ProtectionInfo<T>;

  explicit ProtectionInfoKV(T val) : info_(val) {
    static_assert(sizeof(ProtectionInfoKV<T>) == sizeof(T));
  }

  ProtectionInfo<T> info_;
};

template <typename T>
Status ProtectionInfo<T>::GetStatus() const {
  if (val_ != 0) {
    return Status::Corruption("ProtectionInfo mismatch");
  }
  return Status::OK();
}

template <typename T>
ProtectionInfoKVO<T> ProtectionInfo<T>::ProtectKVO(const Slice& key,
                                                   const Slice& value,
                                                   ValueType op_type) const {
  T val = GetVal();
  val = val ^ static_cast<T>(GetSliceNPHash64(key, ProtectionInfo<T>::kSeedK));
  val =
      val ^ static_cast<T>(GetSliceNPHash64(value, ProtectionInfo<T>::kSeedV));
  val = val ^
        static_cast<T>(NPHash64(reinterpret_cast<char*>(&op_type),
                                sizeof(op_type), ProtectionInfo<T>::kSeedO));
  return ProtectionInfoKVO<T>(val);
}

template <typename T>
ProtectionInfoKVO<T> ProtectionInfo<T>::ProtectKVO(const SliceParts& key,
                                                   const SliceParts& value,
                                                   ValueType op_type) const {
  T val = GetVal();
  val = val ^
        static_cast<T>(GetSlicePartsNPHash64(key, ProtectionInfo<T>::kSeedK));
  val = val ^
        static_cast<T>(GetSlicePartsNPHash64(value, ProtectionInfo<T>::kSeedV));
  val = val ^
        static_cast<T>(NPHash64(reinterpret_cast<char*>(&op_type),
                                sizeof(op_type), ProtectionInfo<T>::kSeedO));
  return ProtectionInfoKVO<T>(val);
}

template <typename T>
ProtectionInfoKV<T> ProtectionInfo<T>::ProtectKV(const Slice& key,
                                                 const Slice& value) const {
  T val = GetVal();
  val = val ^ static_cast<T>(GetSliceNPHash64(key, ProtectionInfo<T>::kSeedK));
  val =
      val ^ static_cast<T>(GetSliceNPHash64(value, ProtectionInfo<T>::kSeedV));
  return ProtectionInfoKV<T>(val);
}

template <typename T>
void ProtectionInfoKVO<T>::UpdateK(const Slice& old_key, const Slice& new_key) {
  T val = GetVal();
  val = val ^
        static_cast<T>(GetSliceNPHash64(old_key, ProtectionInfo<T>::kSeedK));
  val = val ^
        static_cast<T>(GetSliceNPHash64(new_key, ProtectionInfo<T>::kSeedK));
  SetVal(val);
}

template <typename T>
void ProtectionInfoKVO<T>::UpdateK(const SliceParts& old_key,
                                   const SliceParts& new_key) {
  T val = GetVal();
  val = val ^ static_cast<T>(
                  GetSlicePartsNPHash64(old_key, ProtectionInfo<T>::kSeedK));
  val = val ^ static_cast<T>(
                  GetSlicePartsNPHash64(new_key, ProtectionInfo<T>::kSeedK));
  SetVal(val);
}

template <typename T>
void ProtectionInfoKVO<T>::UpdateV(const Slice& old_value,
                                   const Slice& new_value) {
  T val = GetVal();
  val = val ^
        static_cast<T>(GetSliceNPHash64(old_value, ProtectionInfo<T>::kSeedV));
  val = val ^
        static_cast<T>(GetSliceNPHash64(new_value, ProtectionInfo<T>::kSeedV));
  SetVal(val);
}

template <typename T>
void ProtectionInfoKVO<T>::UpdateV(const SliceParts& old_value,
                                   const SliceParts& new_value) {
  T val = GetVal();
  val = val ^ static_cast<T>(
                  GetSlicePartsNPHash64(old_value, ProtectionInfo<T>::kSeedV));
  val = val ^ static_cast<T>(
                  GetSlicePartsNPHash64(new_value, ProtectionInfo<T>::kSeedV));
  SetVal(val);
}

template <typename T>
void ProtectionInfoKVO<T>::UpdateO(ValueType old_op_type,
                                   ValueType new_op_type) {
  T val = GetVal();
  val = val ^ static_cast<T>(NPHash64(reinterpret_cast<char*>(&old_op_type),
                                      sizeof(old_op_type),
                                      ProtectionInfo<T>::kSeedO));
  val = val ^ static_cast<T>(NPHash64(reinterpret_cast<char*>(&new_op_type),
                                      sizeof(new_op_type),
                                      ProtectionInfo<T>::kSeedO));
  SetVal(val);
}

template <typename T>
ProtectionInfo<T> ProtectionInfoKVO<T>::StripKVO(const Slice& key,
                                                 const Slice& value,
                                                 ValueType op_type) const {
  T val = GetVal();
  val = val ^ static_cast<T>(GetSliceNPHash64(key, ProtectionInfo<T>::kSeedK));
  val =
      val ^ static_cast<T>(GetSliceNPHash64(value, ProtectionInfo<T>::kSeedV));
  val = val ^
        static_cast<T>(NPHash64(reinterpret_cast<char*>(&op_type),
                                sizeof(op_type), ProtectionInfo<T>::kSeedO));
  return ProtectionInfo<T>(val);
}

template <typename T>
ProtectionInfo<T> ProtectionInfoKVO<T>::StripKVO(const SliceParts& key,
                                                 const SliceParts& value,
                                                 ValueType op_type) const {
  T val = GetVal();
  val = val ^
        static_cast<T>(GetSlicePartsNPHash64(key, ProtectionInfo<T>::kSeedK));
  val = val ^
        static_cast<T>(GetSlicePartsNPHash64(value, ProtectionInfo<T>::kSeedV));
  val = val ^
        static_cast<T>(NPHash64(reinterpret_cast<char*>(&op_type),
                                sizeof(op_type), ProtectionInfo<T>::kSeedO));
  return ProtectionInfo<T>(val);
}

template <typename T>
ProtectionInfoKVOC<T> ProtectionInfoKVO<T>::ProtectC(
    ColumnFamilyId column_family_id) const {
  T val = GetVal();
  val = val ^ static_cast<T>(NPHash64(
                  reinterpret_cast<char*>(&column_family_id),
                  sizeof(column_family_id), ProtectionInfo<T>::kSeedC));
  return ProtectionInfoKVOC<T>(val);
}

template <typename T>
ProtectionInfoKVO<T> ProtectionInfoKVOC<T>::StripC(
    ColumnFamilyId column_family_id) const {
  T val = GetVal();
  val = val ^ static_cast<T>(NPHash64(
                  reinterpret_cast<char*>(&column_family_id),
                  sizeof(column_family_id), ProtectionInfo<T>::kSeedC));
  return ProtectionInfoKVO<T>(val);
}

template <typename T>
void ProtectionInfoKVOC<T>::UpdateC(ColumnFamilyId old_column_family_id,
                                    ColumnFamilyId new_column_family_id) {
  T val = GetVal();
  val = val ^ static_cast<T>(NPHash64(
                  reinterpret_cast<char*>(&old_column_family_id),
                  sizeof(old_column_family_id), ProtectionInfo<T>::kSeedC));
  val = val ^ static_cast<T>(NPHash64(
                  reinterpret_cast<char*>(&new_column_family_id),
                  sizeof(new_column_family_id), ProtectionInfo<T>::kSeedC));
  SetVal(val);
}

template <typename T>
ProtectionInfoKVOS<T> ProtectionInfoKVO<T>::ProtectS(
    SequenceNumber sequence_number) const {
  T val = GetVal();
  val = val ^ static_cast<T>(NPHash64(reinterpret_cast<char*>(&sequence_number),
                                      sizeof(sequence_number),
                                      ProtectionInfo<T>::kSeedS));
  return ProtectionInfoKVOS<T>(val);
}

template <typename T>
ProtectionInfoKVO<T> ProtectionInfoKVOS<T>::StripS(
    SequenceNumber sequence_number) const {
  T val = GetVal();
  val = val ^ static_cast<T>(NPHash64(reinterpret_cast<char*>(&sequence_number),
                                      sizeof(sequence_number),
                                      ProtectionInfo<T>::kSeedS));
  return ProtectionInfoKVO<T>(val);
}

template <typename T>
void ProtectionInfoKVOS<T>::UpdateS(SequenceNumber old_sequence_number,
                                    SequenceNumber new_sequence_number) {
  T val = GetVal();
  val = val ^ static_cast<T>(NPHash64(
                  reinterpret_cast<char*>(&old_sequence_number),
                  sizeof(old_sequence_number), ProtectionInfo<T>::kSeedS));
  val = val ^ static_cast<T>(NPHash64(
                  reinterpret_cast<char*>(&new_sequence_number),
                  sizeof(new_sequence_number), ProtectionInfo<T>::kSeedS));
  SetVal(val);
}
}  // namespace ROCKSDB_NAMESPACE