Use std::span (backported) and std::byte to make our buffers safer

11 files changed, 189 insertions, 163 deletions

diff --git a/src/audio/CMakeLists.txt b/src/audio/CMakeLists.txt
index d5564f13..6fb2634e 100644
--- a/src/audio/CMakeLists.txt
+++ b/src/audio/CMakeLists.txt
@@ -2,6 +2,6 @@ idf_component_register(
   SRCS "audio_decoder.cpp" "audio_task.cpp" "chunk.cpp" "fatfs_audio_input.cpp"
   "stream_info.cpp" "stream_message.cpp"
   INCLUDE_DIRS "include"
-  REQUIRES "codecs" "drivers" "cbor" "result" "tasks")
+  REQUIRES "codecs" "drivers" "cbor" "result" "tasks" "span" "psram_allocator")
 
 target_compile_options(${COMPONENT_LIB} PRIVATE ${EXTRA_WARNINGS})
diff --git a/src/audio/audio_decoder.cpp b/src/audio/audio_decoder.cpp
index 48637f08..31cfb50e 100644
--- a/src/audio/audio_decoder.cpp
+++ b/src/audio/audio_decoder.cpp
@@ -22,13 +22,14 @@ namespace audio {
 AudioDecoder::AudioDecoder()
     : IAudioElement(),
       stream_info_({}),
-      chunk_buffer_(static_cast<uint8_t*>(
-          heap_caps_malloc(kMaxChunkSize, MALLOC_CAP_SPIRAM)))
+      raw_chunk_buffer_(static_cast<std::byte*>(
+          heap_caps_malloc(kMaxChunkSize, MALLOC_CAP_SPIRAM))),
+      chunk_buffer_(raw_chunk_buffer_, kMaxChunkSize)
 
 {}
 
 AudioDecoder::~AudioDecoder() {
-  free(chunk_buffer_);
+  free(raw_chunk_buffer_);
 }
 
 auto AudioDecoder::SetInputBuffer(MessageBufferHandle_t* buffer) -> void {
@@ -61,21 +62,21 @@ auto AudioDecoder::ProcessStreamInfo(StreamInfo& info)
   return {};
 }
 
-auto AudioDecoder::ProcessChunk(uint8_t* data, std::size_t length)
+auto AudioDecoder::ProcessChunk(cpp::span<std::byte>& chunk)
     -> cpp::result<size_t, AudioProcessingError> {
   if (current_codec_ == nullptr) {
     // Should never happen, but fail explicitly anyway.
     return cpp::fail(UNSUPPORTED_STREAM);
   }
 
-  current_codec_->SetInput(data, length);
+  current_codec_->SetInput(chunk);
 
   bool has_samples_to_send = false;
   bool needs_more_input = false;
   std::optional<codecs::ICodec::ProcessingError> error = std::nullopt;
   WriteChunksToStream(
-      output_buffer_, chunk_buffer_, kMaxChunkSize,
-      [&](uint8_t* buf, size_t len) -> std::size_t {
+      output_buffer_, chunk_buffer_,
+      [&](cpp::span<std::byte> buffer) -> std::size_t {
         std::size_t bytes_written = 0;
         // Continue filling up the output buffer so long as we have samples
         // leftover, or are able to synthesize more samples from the input.
@@ -92,7 +93,7 @@ auto AudioDecoder::ProcessChunk(uint8_t* data, std::size_t length)
             }
           } else {
             auto result = current_codec_->WriteOutputSamples(
-                buf + bytes_written, len - bytes_written);
+                buffer.last(buffer.size() - bytes_written));
             bytes_written += result.first;
             has_samples_to_send = !result.second;
           }
diff --git a/src/audio/audio_task.cpp b/src/audio/audio_task.cpp
index 9ffcab74..112f8f34 100644
--- a/src/audio/audio_task.cpp
+++ b/src/audio/audio_task.cpp
@@ -9,6 +9,7 @@
 #include "freertos/portmacro.h"
 #include "freertos/queue.h"
 #include "freertos/stream_buffer.h"
+#include "span.hpp"
 
 #include "audio_element.hpp"
 #include "chunk.hpp"
@@ -46,8 +47,8 @@ void AudioTaskMain(void* args) {
       bool has_received_message = false;
       if (element->InputBuffer() != nullptr) {
         ChunkReadResult chunk_res = chunk_reader.ReadChunkFromStream(
-            [&](uint8_t* data, std::size_t length) -> std::optional<size_t> {
-              process_res = element->ProcessChunk(data, length);
+            [&](cpp::span<std::byte> data) -> std::optional<size_t> {
+              process_res = element->ProcessChunk(data);
               if (process_res.has_value()) {
                 return process_res.value();
               } else {
@@ -65,11 +66,10 @@ void AudioTaskMain(void* args) {
       }
 
       if (has_received_message) {
-        auto [buffer, length] = chunk_reader.GetLastMessage();
-        MessageType msg = ReadMessageType(buffer, length);
-        if (msg == TYPE_STREAM_INFO) {
-          auto parse_res =
-              ReadMessage<StreamInfo>(&StreamInfo::Parse, buffer, length);
+        auto message = chunk_reader.GetLastMessage();
+        MessageType type = ReadMessageType(message);
+        if (type == TYPE_STREAM_INFO) {
+          auto parse_res = ReadMessage<StreamInfo>(&StreamInfo::Parse, message);
           if (parse_res.has_error()) {
             break;  // TODO.
           }
diff --git a/src/audio/chunk.cpp b/src/audio/chunk.cpp
index 89e54605..8e27bdef 100644
--- a/src/audio/chunk.cpp
+++ b/src/audio/chunk.cpp
@@ -1,11 +1,13 @@
 #include "chunk.hpp"
 
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
 #include <optional>
 
 #include "cbor.h"
+#include "psram_allocator.h"
 
 #include "stream_message.hpp"
 
@@ -17,60 +19,34 @@ const std::size_t kMaxChunkSize = 512;
 // TODO: tune
 static const std::size_t kWorkingBufferSize = kMaxChunkSize * 1.5;
 
-/*
- * The amount of space to allocate for the first chunk's header. After the first
- * chunk, we have a more concrete idea of the header's size and can allocate
- * space for future headers more compactly.
- */
-// TODO: measure how big headers tend to be to pick a better value.
-static const std::size_t kInitialHeaderSize = 32;
-
 auto WriteChunksToStream(MessageBufferHandle_t* stream,
-                         uint8_t* working_buffer,
-                         size_t working_buffer_length,
-                         std::function<size_t(uint8_t*, size_t)> callback,
+                         cpp::span<std::byte> working_buffer,
+                         std::function<size_t(cpp::span<std::byte>)> callback,
                          TickType_t max_wait) -> ChunkWriteResult {
-  size_t header_size = kInitialHeaderSize;
   while (1) {
-    // First, ask the callback for some data to write.
-    size_t chunk_size = callback(working_buffer + header_size,
-                                 working_buffer_length - header_size);
+    // First, write out our chunk header so we know how much space to give to
+    // the callback.
+    auto header_size = WriteTypeOnlyMessage(TYPE_CHUNK_HEADER, working_buffer);
+    if (header_size.has_error()) {
+      return CHUNK_ENCODING_ERROR;
+    }
+
+    // Now we can ask the callback to fill the remaining space.
+    size_t chunk_size = std::invoke(
+        callback,
+        working_buffer.subspan(header_size.value(),
+                               working_buffer.size() - header_size.value()));
 
     if (chunk_size == 0) {
       // They had nothing for us, so bail out.
       return CHUNK_OUT_OF_DATA;
     }
 
-    // Put together a header.
-    CborEncoder arr;
-    cpp::result<size_t, CborError> encoder_res = WriteMessage(
-        TYPE_CHUNK_HEADER,
-        [&](CborEncoder& container) {
-          cbor_encoder_create_array(&container, &arr, 2);
-          cbor_encode_uint(&arr, header_size);
-          cbor_encode_uint(&arr, chunk_size);
-          cbor_encoder_close_container(&container, &arr);
-          return std::nullopt;
-        },
-        working_buffer, working_buffer_length);
-
-    size_t new_header_size = header_size;
-    if (encoder_res.has_error()) {
-      return CHUNK_ENCODING_ERROR;
-    } else {
-      // We can now tune the space to allocate for the header to be closer to
-      // its actual size. We pad this by 2 bytes to allow extra space for the
-      // chunk size and header size fields to each spill over into another byte
-      // each.
-      new_header_size = encoder_res.value() + 2;
-    }
-
     // Try to write to the buffer. Note the return type here will be either 0 or
     // header_size + chunk_size, as MessageBuffer doesn't allow partial writes.
-    size_t actual_write_size = xMessageBufferSend(
-        *stream, working_buffer, header_size + chunk_size, max_wait);
-
-    header_size = new_header_size;
+    size_t actual_write_size =
+        xMessageBufferSend(*stream, working_buffer.data(),
+                           header_size.value() + chunk_size, max_wait);
 
     if (actual_write_size == 0) {
       // We failed to write in time, so bail out. This is techinically data loss
@@ -81,38 +57,39 @@ auto WriteChunksToStream(MessageBufferHandle_t* stream,
   }
 }
 
-ChunkReader::ChunkReader(MessageBufferHandle_t* stream) : stream_(stream) {
-  working_buffer_ = static_cast<uint8_t*>(
-      heap_caps_malloc(kWorkingBufferSize, MALLOC_CAP_SPIRAM));
-};
+ChunkReader::ChunkReader(MessageBufferHandle_t* stream)
+    : stream_(stream),
+      raw_working_buffer_(static_cast<std::byte*>(
+          heap_caps_malloc(kWorkingBufferSize, MALLOC_CAP_SPIRAM))),
+      working_buffer_(raw_working_buffer_, kWorkingBufferSize){};
 
 ChunkReader::~ChunkReader() {
-  free(working_buffer_);
-}
+  free(raw_working_buffer_);
+};
 
 auto ChunkReader::Reset() -> void {
   leftover_bytes_ = 0;
   last_message_size_ = 0;
 }
 
-auto ChunkReader::GetLastMessage() -> std::pair<uint8_t*, size_t> {
-  return std::make_pair(working_buffer_ + leftover_bytes_, last_message_size_);
+auto ChunkReader::GetLastMessage() -> cpp::span<std::byte> {
+  return working_buffer_.subspan(leftover_bytes_, last_message_size_);
 }
 
 auto ChunkReader::ReadChunkFromStream(
-    std::function<std::optional<size_t>(uint8_t*, size_t)> callback,
+    std::function<std::optional<size_t>(cpp::span<std::byte>)> callback,
     TickType_t max_wait) -> ChunkReadResult {
   // First, wait for a message to arrive over the buffer.
   last_message_size_ =
-      xMessageBufferReceive(*stream_, working_buffer_ + leftover_bytes_,
-                            kWorkingBufferSize - leftover_bytes_, max_wait);
+      xMessageBufferReceive(*stream_, raw_working_buffer_ + leftover_bytes_,
+                            working_buffer_.size() - leftover_bytes_, max_wait);
 
   if (last_message_size_ == 0) {
     return CHUNK_READ_TIMEOUT;
   }
 
-  MessageType type =
-      ReadMessageType(working_buffer_ + leftover_bytes_, last_message_size_);
+  cpp::span<std::byte> new_data = GetLastMessage();
+  MessageType type = ReadMessageType(new_data);
 
   if (type != TYPE_CHUNK_HEADER) {
     // This message wasn't for us, so let the caller handle it.
@@ -121,31 +98,30 @@ auto ChunkReader::ReadChunkFromStream(
   }
 
   // Work the size and position of the chunk.
-  size_t header_length = 0, chunk_length = 0;
-
-  // TODO: chunker header type to encapsulate this?
+  auto chunk_data = GetAdditionalData(new_data);
 
   // Now we need to stick the end of the last chunk (if it exists) onto the
   // front of the new chunk. Do it this way around bc we assume the old chunk
   // is shorter, and therefore faster to move.
-  uint8_t* combined_buffer = working_buffer_ + header_length - leftover_bytes_;
-  size_t combined_buffer_size = leftover_bytes_ + chunk_length;
+  cpp::span<std::byte> leftover_data = working_buffer_.first(leftover_bytes_);
+  cpp::span<std::byte> combined_data(chunk_data.data() - leftover_data.size(),
+                                     leftover_data.size() + chunk_data.size());
   if (leftover_bytes_ > 0) {
-    memmove(combined_buffer, working_buffer_, leftover_bytes_);
+    std::copy_backward(leftover_data.begin(), leftover_data.end(),
+                       combined_data.begin());
   }
 
   // Tell the callback about the new data.
-  std::optional<size_t> amount_processed =
-      callback(combined_buffer, combined_buffer_size);
+  std::optional<size_t> amount_processed = std::invoke(callback, combined_data);
   if (!amount_processed) {
     return CHUNK_PROCESSING_ERROR;
   }
 
   // Prepare for the next iteration.
-  leftover_bytes_ = combined_buffer_size - amount_processed.value();
+  leftover_bytes_ = combined_data.size() - amount_processed.value();
   if (leftover_bytes_ > 0) {
-    memmove(working_buffer_, combined_buffer + amount_processed.value(),
-            leftover_bytes_);
+    std::copy(combined_data.begin() + amount_processed.value(),
+              combined_data.end(), working_buffer_.begin());
     return CHUNK_LEFTOVER_DATA;
   }
 
diff --git a/src/audio/fatfs_audio_input.cpp b/src/audio/fatfs_audio_input.cpp
index d4cbf6db..6777bef2 100644
--- a/src/audio/fatfs_audio_input.cpp
+++ b/src/audio/fatfs_audio_input.cpp
@@ -23,24 +23,28 @@ static const std::size_t kMinFileReadSize = 1024 * 4;
 static const std::size_t kOutputBufferSize = 1024 * 4;
 
 FatfsAudioInput::FatfsAudioInput(std::shared_ptr<drivers::SdStorage> storage)
-    : IAudioElement(), storage_(storage) {
-  file_buffer_ = static_cast<uint8_t*>(
-      heap_caps_malloc(kFileBufferSize, MALLOC_CAP_SPIRAM));
-  file_buffer_read_pos_ = file_buffer_;
-  file_buffer_write_pos_ = file_buffer_;
-  chunk_buffer_ =
-      static_cast<uint8_t*>(heap_caps_malloc(kMaxChunkSize, MALLOC_CAP_SPIRAM));
-
-  output_buffer_memory_ = static_cast<uint8_t*>(
-      heap_caps_malloc(kOutputBufferSize, MALLOC_CAP_SPIRAM));
+    : IAudioElement(),
+      storage_(storage),
+      raw_file_buffer_(static_cast<std::byte*>(
+          heap_caps_malloc(kFileBufferSize, MALLOC_CAP_SPIRAM))),
+      file_buffer_(raw_file_buffer_, kFileBufferSize),
+      file_buffer_read_pos_(file_buffer_.begin()),
+      file_buffer_write_pos_(file_buffer_.begin()),
+      raw_chunk_buffer_(static_cast<std::byte*>(
+          heap_caps_malloc(kMaxChunkSize, MALLOC_CAP_SPIRAM))),
+      chunk_buffer_(raw_chunk_buffer_, kMaxChunkSize),
+      current_file_(),
+      is_file_open_(false),
+      output_buffer_memory_(static_cast<uint8_t*>(
+          heap_caps_malloc(kOutputBufferSize, MALLOC_CAP_SPIRAM))) {
   output_buffer_ = new MessageBufferHandle_t;
   *output_buffer_ = xMessageBufferCreateStatic(
       kOutputBufferSize, output_buffer_memory_, &output_buffer_metadata_);
 }
 
 FatfsAudioInput::~FatfsAudioInput() {
-  free(file_buffer_);
-  free(chunk_buffer_);
+  free(raw_file_buffer_);
+  free(raw_chunk_buffer_);
   vMessageBufferDelete(output_buffer_);
   free(output_buffer_memory_);
   free(output_buffer_);
@@ -64,22 +68,22 @@ auto FatfsAudioInput::ProcessStreamInfo(StreamInfo& info)
 
   is_file_open_ = true;
 
-  auto write_res =
+  auto write_size =
       WriteMessage(TYPE_STREAM_INFO,
                    std::bind(&StreamInfo::Encode, info, std::placeholders::_1),
-                   chunk_buffer_, kMaxChunkSize);
+                   chunk_buffer_);
 
-  if (write_res.has_error()) {
+  if (write_size.has_error()) {
     return cpp::fail(IO_ERROR);
   } else {
-    xMessageBufferSend(output_buffer_, chunk_buffer_, write_res.value(),
+    xMessageBufferSend(output_buffer_, chunk_buffer_.data(), write_size.value(),
                        portMAX_DELAY);
   }
 
   return {};
 }
 
-auto FatfsAudioInput::ProcessChunk(uint8_t* data, std::size_t length)
+auto FatfsAudioInput::ProcessChunk(cpp::span<std::byte>& chunk)
     -> cpp::result<size_t, AudioProcessingError> {
   return cpp::fail(UNSUPPORTED_STREAM);
 }
@@ -93,9 +97,9 @@ auto FatfsAudioInput::GetRingBufferDistance() -> size_t {
   }
   return
       // Read position to end of buffer.
-      (file_buffer_ + kFileBufferSize - file_buffer_read_pos_)
+      (file_buffer_.end() - file_buffer_read_pos_)
       // Start of buffer to write position.
-      + (file_buffer_write_pos_ - file_buffer_);
+      + (file_buffer_write_pos_ - file_buffer_.begin());
 }
 
 auto FatfsAudioInput::ProcessIdle() -> cpp::result<void, AudioProcessingError> {
@@ -103,19 +107,20 @@ auto FatfsAudioInput::ProcessIdle() -> cpp::result<void, AudioProcessingError> {
   // file's contents.
   if (is_file_open_) {
     size_t ringbuf_distance = GetRingBufferDistance();
-    if (kFileBufferSize - ringbuf_distance > kMinFileReadSize) {
+    if (file_buffer_.size() - ringbuf_distance > kMinFileReadSize) {
       size_t read_size;
       if (file_buffer_write_pos_ < file_buffer_read_pos_) {
         // Don't worry about the start of buffer -> read pos size; we can get to
         // it next iteration.
         read_size = file_buffer_read_pos_ - file_buffer_write_pos_;
       } else {
-        read_size = file_buffer_ - file_buffer_write_pos_;
+        read_size = file_buffer_.begin() - file_buffer_write_pos_;
       }
 
       UINT bytes_read = 0;
-      FRESULT result = f_read(&current_file_, file_buffer_write_pos_, read_size,
-                              &bytes_read);
+      FRESULT result =
+          f_read(&current_file_, std::addressof(file_buffer_write_pos_),
+                 read_size, &bytes_read);
       if (result != FR_OK) {
         ESP_LOGE(kTag, "file I/O error %d", result);
         return cpp::fail(IO_ERROR);
@@ -129,17 +134,17 @@ auto FatfsAudioInput::ProcessIdle() -> cpp::result<void, AudioProcessingError> {
       }
 
       file_buffer_write_pos_ += bytes_read;
-      if (file_buffer_write_pos_ == file_buffer_ + kFileBufferSize) {
-        file_buffer_write_pos_ = file_buffer_;
+      if (file_buffer_write_pos_ == file_buffer_.end()) {
+        file_buffer_write_pos_ = file_buffer_.begin();
       }
     }
   }
 
   // Now stream data into the output buffer until it's full.
-  pending_read_pos_ = nullptr;
+  pending_read_pos_ = file_buffer_read_pos_;
   ChunkWriteResult result = WriteChunksToStream(
-      output_buffer_, chunk_buffer_, kMaxChunkSize,
-      [&](uint8_t* b, size_t s) { return SendChunk(b, s); }, kServiceInterval);
+      output_buffer_, chunk_buffer_,
+      [&](cpp::span<std::byte> d) { return SendChunk(d); }, kServiceInterval);
 
   switch (result) {
     case CHUNK_WRITE_TIMEOUT:
@@ -152,28 +157,28 @@ auto FatfsAudioInput::ProcessIdle() -> cpp::result<void, AudioProcessingError> {
   }
 }
 
-auto FatfsAudioInput::SendChunk(uint8_t* buffer, size_t size) -> size_t {
-  if (pending_read_pos_ != nullptr) {
-    file_buffer_read_pos_ = pending_read_pos_;
-  }
+auto FatfsAudioInput::SendChunk(cpp::span<std::byte> dest) -> size_t {
+  file_buffer_read_pos_ = pending_read_pos_;
 
   if (file_buffer_read_pos_ == file_buffer_write_pos_) {
     return 0;
   }
-  std::size_t write_size;
+  std::size_t chunk_size;
   if (file_buffer_read_pos_ > file_buffer_write_pos_) {
-    write_size = file_buffer_ + kFileBufferSize - file_buffer_read_pos_;
+    chunk_size = file_buffer_.end() - file_buffer_read_pos_;
   } else {
-    write_size = file_buffer_write_pos_ - file_buffer_read_pos_;
+    chunk_size = file_buffer_write_pos_ - file_buffer_read_pos_;
   }
-  write_size = std::min(write_size, size);
-  memcpy(buffer, file_buffer_read_pos_, write_size);
+  chunk_size = std::min(chunk_size, dest.size());
+
+  cpp::span<std::byte> source(file_buffer_read_pos_, chunk_size);
+  std::copy(source.begin(), source.end(), dest.begin());
 
-  pending_read_pos_ = file_buffer_read_pos_ + write_size;
-  if (pending_read_pos_ == file_buffer_ + kFileBufferSize) {
-    pending_read_pos_ = file_buffer_;
+  pending_read_pos_ = file_buffer_read_pos_ + chunk_size;
+  if (pending_read_pos_ == file_buffer_.end()) {
+    pending_read_pos_ = file_buffer_.begin();
   }
-  return write_size;
+  return chunk_size;
 }
 
 }  // namespace audio
diff --git a/src/audio/include/audio_decoder.hpp b/src/audio/include/audio_decoder.hpp
index 4d2fd5f3..a32442da 100644
--- a/src/audio/include/audio_decoder.hpp
+++ b/src/audio/include/audio_decoder.hpp
@@ -1,8 +1,10 @@
 #pragma once
 
 #include <cstddef>
+#include <cstdint>
 
 #include "ff.h"
+#include "span.hpp"
 
 #include "audio_element.hpp"
 #include "codec.hpp"
@@ -23,8 +25,8 @@ class AudioDecoder : public IAudioElement {
 
   auto ProcessStreamInfo(StreamInfo& info)
       -> cpp::result<void, AudioProcessingError>;
-  auto ProcessChunk(uint8_t* data, std::size_t length)
-      -> cpp::result<size_t, AudioProcessingError>;
+  auto ProcessChunk(cpp::span<std::byte>& chunk)
+      -> cpp::result<std::size_t, AudioProcessingError>;
   auto ProcessIdle() -> cpp::result<void, AudioProcessingError>;
 
   AudioDecoder(const AudioDecoder&) = delete;
@@ -34,7 +36,8 @@ class AudioDecoder : public IAudioElement {
   std::unique_ptr<codecs::ICodec> current_codec_;
   std::optional<StreamInfo> stream_info_;
 
-  uint8_t* chunk_buffer_;
+  std::byte* raw_chunk_buffer_;
+  cpp::span<std::byte> chunk_buffer_;
 };
 
 }  // namespace audio
diff --git a/src/audio/include/audio_element.hpp b/src/audio/include/audio_element.hpp
index 5b697784..06e47b35 100644
--- a/src/audio/include/audio_element.hpp
+++ b/src/audio/include/audio_element.hpp
@@ -1,7 +1,5 @@
 #pragma once
 
-#include <stdint.h>
-
 #include <cstdint>
 
 #include "freertos/FreeRTOS.h"
@@ -9,6 +7,7 @@
 #include "freertos/message_buffer.h"
 #include "freertos/portmacro.h"
 #include "result.hpp"
+#include "span.hpp"
 
 #include "stream_info.hpp"
 #include "types.hpp"
@@ -77,8 +76,8 @@ class IAudioElement {
    * bytes in this chunk that were actually used; leftover bytes will be
    * prepended to the next call.
    */
-  virtual auto ProcessChunk(uint8_t* data, std::size_t length)
-      -> cpp::result<size_t, AudioProcessingError> = 0;
+  virtual auto ProcessChunk(cpp::span<std::byte>& chunk)
+      -> cpp::result<std::size_t, AudioProcessingError> = 0;
 
   /*
    * Called when there has been no data received over the input buffer for some
diff --git a/src/audio/include/chunk.hpp b/src/audio/include/chunk.hpp
index 365c83d0..0cbe8d5c 100644
--- a/src/audio/include/chunk.hpp
+++ b/src/audio/include/chunk.hpp
@@ -12,6 +12,7 @@
 #include "freertos/portmacro.h"
 #include "freertos/queue.h"
 #include "result.hpp"
+#include "span.hpp"
 
 namespace audio {
 
@@ -37,9 +38,8 @@ enum ChunkWriteResult {
  * more input to read.
  */
 auto WriteChunksToStream(MessageBufferHandle_t* stream,
-                         uint8_t* working_buffer,
-                         size_t working_buffer_length,
-                         std::function<size_t(uint8_t*, size_t)> callback,
+                         cpp::span<std::byte> working_buffer,
+                         std::function<size_t(cpp::span<std::byte>)> callback,
                          TickType_t max_wait) -> ChunkWriteResult;
 
 enum ChunkReadResult {
@@ -64,7 +64,7 @@ class ChunkReader {
 
   auto Reset() -> void;
 
-  auto GetLastMessage() -> std::pair<uint8_t*, size_t>;
+  auto GetLastMessage() -> cpp::span<std::byte>;
 
   /*
    * Reads chunks of data from the given input stream, and invokes the given
@@ -79,12 +79,13 @@ class ChunkReader {
    * will place the message at the start of the working_buffer and then return.
    */
   auto ReadChunkFromStream(
-      std::function<std::optional<size_t>(uint8_t*, size_t)> callback,
+      std::function<std::optional<size_t>(cpp::span<std::byte>)> callback,
       TickType_t max_wait) -> ChunkReadResult;
 
  private:
   MessageBufferHandle_t* stream_;
-  uint8_t* working_buffer_;
+  std::byte* raw_working_buffer_;
+  cpp::span<std::byte> working_buffer_;
 
   std::size_t leftover_bytes_ = 0;
   std::size_t last_message_size_ = 0;
diff --git a/src/audio/include/fatfs_audio_input.hpp b/src/audio/include/fatfs_audio_input.hpp
index c54b32bd..3ca79457 100644
--- a/src/audio/include/fatfs_audio_input.hpp
+++ b/src/audio/include/fatfs_audio_input.hpp
@@ -8,6 +8,7 @@
 
 #include "freertos/message_buffer.h"
 #include "freertos/queue.h"
+#include "span.hpp"
 
 #include "audio_element.hpp"
 #include "storage.hpp"
@@ -21,28 +22,28 @@ class FatfsAudioInput : public IAudioElement {
 
   auto ProcessStreamInfo(StreamInfo& info)
       -> cpp::result<void, AudioProcessingError>;
-  auto ProcessChunk(uint8_t* data, std::size_t length)
-      -> cpp::result<size_t, AudioProcessingError>;
+  auto ProcessChunk(cpp::span<std::byte>& chunk)
+      -> cpp::result<std::size_t, AudioProcessingError> = 0;
   auto ProcessIdle() -> cpp::result<void, AudioProcessingError>;
 
-  auto SendChunk(uint8_t* buffer, size_t size) -> size_t;
+  auto SendChunk(cpp::span<std::byte> dest) -> size_t;
 
  private:
   auto GetRingBufferDistance() -> size_t;
 
   std::shared_ptr<drivers::SdStorage> storage_;
 
-  uint8_t* file_buffer_;
-  uint8_t* file_buffer_read_pos_;
-  uint8_t* pending_read_pos_;
-  uint8_t* file_buffer_write_pos_;
+  std::byte* raw_file_buffer_;
+  cpp::span<std::byte> file_buffer_;
+  cpp::span<std::byte>::iterator file_buffer_read_pos_;
+  cpp::span<std::byte>::iterator pending_read_pos_;
+  cpp::span<std::byte>::iterator file_buffer_write_pos_;
 
-  uint8_t* chunk_buffer_;
+  std::byte* raw_chunk_buffer_;
+  cpp::span<std::byte> chunk_buffer_;
 
   FIL current_file_;
-  bool is_file_open_ = false;
-
-  MessageBufferHandle_t input_buffer_;
+  bool is_file_open_;
 
   uint8_t* output_buffer_memory_;
   StaticMessageBuffer_t output_buffer_metadata_;
diff --git a/src/audio/include/stream_message.hpp b/src/audio/include/stream_message.hpp
index cbd7c733..043f9dc3 100644
--- a/src/audio/include/stream_message.hpp
+++ b/src/audio/include/stream_message.hpp
@@ -1,12 +1,12 @@
 #pragma once
 
-#include <stdint.h>
-
+#include <cstdint>
 #include <functional>
 #include <optional>
 
 #include "cbor.h"
 #include "result.hpp"
+#include "span.hpp"
 
 namespace audio {
 
@@ -20,14 +20,13 @@ enum MessageType {
 };
 
 template <typename Writer>
-auto WriteMessage(MessageType type,
-                  Writer&& writer,
-                  uint8_t* buffer,
-                  size_t length) -> cpp::result<size_t, CborError> {
+auto WriteMessage(MessageType type, Writer&& writer, cpp::span<std::byte> data)
+    -> cpp::result<size_t, CborError> {
   CborEncoder root;
   CborEncoder container;
+  uint8_t* cast_data = reinterpret_cast<uint8_t*>(data.data());
 
-  cbor_encoder_init(&root, buffer, length, kEncoderFlags);
+  cbor_encoder_init(&root, cast_data, data.size(), kEncoderFlags);
   cbor_encoder_create_array(&root, &container, 2);
   cbor_encode_uint(&container, type);
 
@@ -37,17 +36,18 @@ auto WriteMessage(MessageType type,
   }
 
   cbor_encoder_close_container(&root, &container);
-  return cbor_encoder_get_buffer_size(&root, buffer);
+  return cbor_encoder_get_buffer_size(&root, cast_data);
 }
 
 template <typename Result, typename Reader>
-auto ReadMessage(Reader&& reader, uint8_t* buffer, size_t length)
+auto ReadMessage(Reader&& reader, cpp::span<std::byte> data)
     -> cpp::result<Result, CborError> {
   CborParser parser;
   CborValue root;
   CborValue container;
 
-  cbor_parser_init(buffer, length, kDecoderFlags, &parser, &root);
+  cbor_parser_init(reinterpret_cast<uint8_t*>(data.data()), data.size(),
+                   kDecoderFlags, &parser, &root);
   cbor_value_enter_container(&root, &container);
   // Skip the type header
   cbor_value_advance_fixed(&container);
@@ -55,6 +55,9 @@ auto ReadMessage(Reader&& reader, uint8_t* buffer, size_t length)
   return std::invoke(reader, container);
 }
 
-auto ReadMessageType(uint8_t* buffer, size_t length) -> MessageType;
+auto WriteTypeOnlyMessage(MessageType type, cpp::span<std::byte> data)
+    -> cpp::result<size_t, CborError>;
+auto ReadMessageType(cpp::span<std::byte> msg) -> MessageType;
+auto GetAdditionalData(cpp::span<std::byte> msg) -> cpp::span<std::byte>;
 
 }  // namespace audio
diff --git a/src/audio/stream_message.cpp b/src/audio/stream_message.cpp
index 58868ce8..055e7e96 100644
--- a/src/audio/stream_message.cpp
+++ b/src/audio/stream_message.cpp
@@ -3,24 +3,61 @@
 #include <cstdint>
 
 #include "cbor.h"
+#include "esp-idf/components/cbor/tinycbor/src/cbor.h"
+#include "span.hpp"
 
 namespace audio {
 
 const int kEncoderFlags = 0;
 const int kDecoderFlags = 0;
 
-auto ReadMessageType(uint8_t* buffer, size_t length) -> MessageType {
+auto WriteTypeOnlyMessage(MessageType type, cpp::span<std::byte> data)
+    -> cpp::result<size_t, CborError> {
+  CborEncoder root;
+  CborEncoder container;
+  uint8_t* cast_data = reinterpret_cast<uint8_t*>(data.data());
+
+  cbor_encoder_init(&root, cast_data, data.size(), kEncoderFlags);
+  cbor_encode_uint(&container, type);
+
+  return cbor_encoder_get_buffer_size(&root, cast_data);
+}
+
+auto ReadMessageType(cpp::span<std::byte> msg) -> MessageType {
   CborParser parser;
   CborValue root;
   CborValue container;
 
-  cbor_parser_init(buffer, length, kDecoderFlags, &parser, &root);
-  cbor_value_enter_container(&root, &container);
+  cbor_parser_init(reinterpret_cast<uint8_t*>(msg.data()), msg.size(),
+                   kDecoderFlags, &parser, &root);
 
   uint64_t header = 0;
-  cbor_value_get_uint64(&container, &header);
+  if (cbor_value_is_container(&root)) {
+    cbor_value_enter_container(&root, &container);
+    cbor_value_get_uint64(&container, &header);
+  } else {
+    cbor_value_get_uint64(&root, &header);
+  }
 
   return static_cast<MessageType>(header);
 }
 
+auto GetAdditionalData(cpp::span<std::byte> msg) -> cpp::span<std::byte> {
+  CborParser parser;
+  CborValue root;
+  uint8_t* cast_data = reinterpret_cast<uint8_t*>(msg.data());
+
+  cbor_parser_init(cast_data, msg.size(), kDecoderFlags, &parser, &root);
+
+  while (!cbor_value_at_end(&root)) {
+    cbor_value_advance(&root);
+  }
+
+  const uint8_t* remaining = cbor_value_get_next_byte(&root);
+  size_t header_size = remaining - cast_data;
+
+  return cpp::span<std::byte>(msg.data() + header_size,
+                              msg.size() - header_size);
+}
+
 }  // namespace audio