1  #pragma once
  2  
  3  #include <cstddef>
  4  #include <cstdint>
  5  #include <cstring>
  6  #include <type_traits>
  7  
  8  #include <biscuit/assert.hpp>
  9  
 10  namespace biscuit {
 11  
 12  /**
 13   * An arbitrarily sized buffer that code is written into.
 14   *
 15   * Also contains other member functions for manipulating
 16   * the data within the code buffer.
 17   */
 18  class CodeBuffer {
 19  public:
 20      // Default capacity of 4KB.
 21      static constexpr size_t default_capacity = 4096;
 22  
 23      /**
 24       * Constructor
 25       *
 26       * @param capacity The initial capacity of the code buffer in bytes.
 27       */
 28      explicit CodeBuffer(size_t capacity = default_capacity);
 29  
 30      /**
 31       * Constructor
 32       *
 33       * @param buffer   A non-null pointer to an allocated buffer of size `capacity`.
 34       * @param capacity The capacity of the memory pointed to by `buffer`.
 35       *
 36       * @pre The given memory buffer must not be null.
 37       * @pre The given memory buffer must be at minimum `capacity` bytes in size.
 38       *
 39       * @note The caller is responsible for managing the lifetime of the given memory.
 40       *       CodeBuffer will *not* free the memory once it goes out of scope.
 41       */
 42      explicit CodeBuffer(uint8_t* buffer, size_t capacity);
 43  
 44      // Copy constructor and assignment is deleted in order to prevent unintentional memory leaks.
 45      CodeBuffer(const CodeBuffer&) = delete;
 46      CodeBuffer& operator=(const CodeBuffer&) = delete;
 47  
 48      // Move constructing or moving the buffer in general is allowed, as it's a transfer of control.
 49      CodeBuffer(CodeBuffer&& other) noexcept;
 50      CodeBuffer& operator=(CodeBuffer&& other) noexcept;
 51  
 52      /**
 53       * Destructor
 54       *
 55       * If a custom memory buffer is not given to the code buffer,
 56       * then the code buffer will automatically free any memory
 57       * it had allocated in order to be able to emit code.
 58       */
 59      ~CodeBuffer() noexcept;
 60  
 61      /// Returns whether or not the memory is managed by the code buffer.
 62      [[nodiscard]] bool IsManaged() const noexcept { return m_is_managed; }
 63  
 64      /// Retrieves the current cursor position within the buffer.
 65      [[nodiscard]] ptrdiff_t GetCursorOffset() const noexcept {
 66          return m_cursor - m_buffer;
 67      }
 68  
 69      /// Retrieves the current address of the cursor within the buffer.
 70      [[nodiscard]] uintptr_t GetCursorAddress() const noexcept {
 71          return GetOffsetAddress(GetCursorOffset());
 72      }
 73  
 74      /// Retrieves the cursor pointer
 75      [[nodiscard]] uint8_t* GetCursorPointer() noexcept {
 76          return GetOffsetPointer(GetCursorOffset());
 77      }
 78  
 79      /// Retrieves the cursor pointer
 80      [[nodiscard]] const uint8_t* GetCursorPointer() const noexcept {
 81          return GetOffsetPointer(GetCursorOffset());
 82      }
 83  
 84      /// Retrieves the address of an arbitrary offset within the buffer.
 85      [[nodiscard]] uintptr_t GetOffsetAddress(ptrdiff_t offset) const noexcept {
 86          return reinterpret_cast<uintptr_t>(GetOffsetPointer(offset));
 87      }
 88  
 89      /// Retrieves the pointer to an arbitrary location within the buffer.
 90      [[nodiscard]] uint8_t* GetOffsetPointer(ptrdiff_t offset) noexcept {
 91          BISCUIT_ASSERT(offset >= 0 && offset <= GetCursorOffset());
 92          return m_buffer + offset;
 93      }
 94  
 95      /// Retrieves the pointer to an arbitrary location within the buffer.
 96      [[nodiscard]] const uint8_t* GetOffsetPointer(ptrdiff_t offset) const noexcept {
 97          BISCUIT_ASSERT(offset >= 0 && offset <= GetCursorOffset());
 98          return m_buffer + offset;
 99      }
100  
101      /**
102       * Allows rewinding of the code buffer cursor.
103       *
104       * @param offset The offset to rewind the cursor by.
105       *
106       * @note If no offset is provided, then this function rewinds the
107       *       cursor to the beginning of the buffer.
108       *
109       * @note The offset may not be larger than the current cursor offset
110       *       and may not be less than the current buffer starting address.
111       */
112      void RewindCursor(ptrdiff_t offset = 0) noexcept {
113          auto* rewound = m_buffer + offset;
114          BISCUIT_ASSERT(m_buffer <= rewound && rewound <= m_cursor);
115          m_cursor = rewound;
116      }
117  
118      /**
119       * Whether or not the underlying buffer has enough room for the
120       * given number of bytes.
121       *
122       * @param num_bytes The number of bytes to store in the buffer.
123       */
124      [[nodiscard]] bool HasSpaceFor(size_t num_bytes) const noexcept {
125          return GetRemainingBytes() >= num_bytes;
126      }
127  
128      /// Returns the size of the data written to the buffer in bytes.
129      [[nodiscard]] size_t GetSizeInBytes() const noexcept {
130          EnsureBufferRange();
131          return static_cast<size_t>(m_cursor - m_buffer);
132      }
133  
134      /// Returns the total number of remaining bytes in the buffer.
135      [[nodiscard]] size_t GetRemainingBytes() const noexcept {
136          EnsureBufferRange();
137          return static_cast<size_t>((m_buffer + m_capacity) - m_cursor);
138      }
139  
140      /**
141       * Grows the underlying memory of the code buffer
142       *
143       * @param new_capacity The new capacity of the code buffer in bytes.
144       *
145       * @pre The underlying memory of the code buffer *must* be managed
146       *      by the code buffer itself. Attempts to grow the buffer
147       *      with memory that is not managed by it will result in
148       *      an assertion being hit.
149       *
150       * @note Calling this with a new capacity that is less than or equal
151       *       to the current capacity of the buffer will result in
152       *       this function doing nothing.
153       */
154      void Grow(size_t new_capacity);
155  
156      /**
157       * Emits a given value into the code buffer.
158       *
159       * @param value The value to emit into the code buffer.
160       * @tparam T    A trivially-copyable type.
161       */
162      template <typename T>
163      void Emit(T value) noexcept {
164          static_assert(std::is_trivially_copyable_v<T>,
165                        "It's undefined behavior to memcpy a non-trivially-copyable type.");
166          BISCUIT_ASSERT(HasSpaceFor(sizeof(T)));
167  
168          std::memcpy(m_cursor, &value, sizeof(T));
169          m_cursor += sizeof(T);
170      }
171  
172      /// Emits a 16-bit value into the code buffer.
173      void Emit16(uint32_t value) noexcept {
174          Emit(static_cast<uint16_t>(value));
175      }
176  
177      /// Emits a 32-bit value into the code buffer.
178      void Emit32(uint32_t value) noexcept {
179          Emit(value);
180      }
181  
182      /**
183       * Sets the internal code buffer to be executable.
184       *
185       * @note This will make the contained region of memory non-writable
186       *       to satisfy operating under W^X contexts. To make the
187       *       region writable again, use SetWritable().
188       */
189      void SetExecutable();
190  
191      /**
192       * Sets the internal code buffer to be writable
193       *
194       * @note This will make the contained region of memory non-executable
195       *       to satisfy operating under W^X contexts. To make the region
196       *       executable again, use SetExecutable().
197       */
198      void SetWritable();
199  
200  private:
201      void EnsureBufferRange() const noexcept {
202          BISCUIT_ASSERT(m_cursor >= m_buffer && m_cursor <= m_buffer + m_capacity);
203      }
204  
205      uint8_t* m_buffer = nullptr;
206      uint8_t* m_cursor = nullptr;
207      size_t m_capacity = 0;
208      bool m_is_managed = false;
209  };
210  
211  } // namespace biscuit