#pragma once

#ifndef SEC_LOCKED_BUFFER_HPP
  #define SEC_LOCKED_BUFFER_HPP

  #include <cstdlib>
  #include <cstring>
  #include <stdexcept>
  #include <sys/mman.h>
  #include <unistd.h>

namespace sec {

template <typename T, std::size_t N> class locked_buffer {
public:
  locked_buffer() {
    this->size_ = N * sizeof(T);
    this->ptr_ = static_cast<T *>(std::malloc(this->size_));
    if (!this->ptr_)
      throw std::bad_alloc();
    if (mlock(this->ptr_, this->size_) != 0) {
      std::free(this->ptr_);
      throw std::runtime_error("mlock failed");
    }
  }

  ~locked_buffer() noexcept {
    cleanse();
    if (this->ptr_) {
      munlock(this->ptr_, this->size_);
      std::free(this->ptr_);
    }
  }

  locked_buffer(const locked_buffer &) = delete;
  locked_buffer &operator=(const locked_buffer &) = delete;

  locked_buffer(locked_buffer &&other) noexcept
      : ptr_(other.ptr_), size_(other.size_) {
    other.ptr_ = nullptr;
    other.size_ = 0;
  }

  locked_buffer &operator=(locked_buffer &&other) noexcept {
    if (this != &other) {
      cleanse();
      if (this->ptr_) {
        munlock(this->ptr_, this->size_);
        std::free(this->ptr_);
      }
      this->ptr_ = other.ptr_;
      this->size_ = other.size_;
      other.ptr_ = nullptr;
      other.size_ = 0;
    }
    return *this;
  }

  T *data() noexcept { return this->ptr_; }
  const T *data() const noexcept { return this->ptr_; }
  std::size_t size() const noexcept { return N; }

  T &operator[](std::size_t i) { return this->ptr_[i]; }
  const T &operator[](std::size_t i) const { return this->ptr_[i]; }

  T *begin() noexcept { return this->ptr_; }
  T *end() noexcept { return this->ptr_ + N; }
  const T *begin() const noexcept { return this->ptr_; }
  const T *end() const noexcept { return this->ptr_ + N; }

private:
  void cleanse() noexcept {
    if (this->ptr_) {
      std::memset(this->ptr_, 0, this->size_);
    }
  }

  T *ptr_{nullptr};
  std::size_t size_{0};
};

} // namespace sec

#endif
