Add ByteBuf wrapper for aws_byte_buf

Signed-off-by: Daniel Carl Jones <djonesoa@amazon.com>

mountpoint-s3-crt/CHANGELOG.md

@@ -3,6 +3,7 @@
 * Update to latest CRT dependencies.
 * Checksum hashers no longer implement `std::hash::Hasher`. ([#1082](https://github.com/awslabs/mountpoint-s3/pull/1082))
 * Add bindings to remaining checksum types CRC64, SHA1, and SHA256. ([#1082](https://github.com/awslabs/mountpoint-s3/pull/1082))
+* Add wrapping type `ByteBuf` for `aws_byte_buf`. ([#1082](https://github.com/awslabs/mountpoint-s3/pull/1082))
 
 ## v0.10.0 (October 17, 2024)
 

mountpoint-s3-crt/src/checksums/sha1.rs

@@ -1,11 +1,11 @@
-use std::{mem::MaybeUninit, ptr::NonNull};
+use std::ptr::NonNull;
 
 use mountpoint_s3_crt_sys::{
-    aws_byte_buf, aws_byte_buf_init, aws_hash, aws_hash_destroy, aws_hash_finalize, aws_hash_update, aws_sha1_new,
-    AWS_SHA1_LEN,
+    aws_hash, aws_hash_destroy, aws_hash_finalize, aws_hash_update, aws_sha1_new, AWS_SHA1_LEN,
 };
 
 use crate::common::allocator::Allocator;
+use crate::common::byte_buf::ByteBuf;
 use crate::common::error::Error;
 use crate::{CrtError as _, ToAwsByteCursor};
 
@@ -60,20 +60,13 @@ impl Sha1Hasher {
 
     /// Finalize the hash state and return the computed SHA1 checksum value.
     pub fn finalize(self, allocator: &Allocator) -> Result<Sha1, Error> {
-        // SAFETY: allocator is a valid aws_allocator, and `aws_byte_buf_init` initializes the buffer on success.
-        let mut buffer = unsafe {
-            let mut buffer: MaybeUninit<aws_byte_buf> = MaybeUninit::uninit();
-            aws_byte_buf_init(buffer.as_mut_ptr(), allocator.inner.as_ptr(), Sha1::LENGTH).ok_or_last_error()?;
-            buffer.assume_init()
-        };
+        let mut buffer = ByteBuf::new(allocator, Sha1::LENGTH)?;
 
         // SAFETY: `self.inner` is a valid `aws_hash` and `buffer` was initialized above.
-        unsafe { aws_hash_finalize(self.inner.as_ptr(), &mut buffer, 0).ok_or_last_error()? };
+        unsafe { aws_hash_finalize(self.inner.as_ptr(), buffer.as_mut_ptr(), 0).ok_or_last_error()? };
 
-        // SAFETY: `buffer` holds a valid buffer.
-        let output = unsafe { std::slice::from_raw_parts(buffer.buffer, buffer.len) };
-
-        Ok(Sha1(output.try_into().unwrap()))
+        // Slice will be copied into the struct.
+        Ok(Sha1(buffer.as_slice().try_into().unwrap()))
     }
 }
 

mountpoint-s3-crt/src/checksums/sha256.rs

@@ -1,11 +1,11 @@
-use std::{mem::MaybeUninit, ptr::NonNull};
+use std::ptr::NonNull;
 
 use mountpoint_s3_crt_sys::{
-    aws_byte_buf, aws_byte_buf_init, aws_hash, aws_hash_destroy, aws_hash_finalize, aws_hash_update, aws_sha256_new,
-    AWS_SHA256_LEN,
+    aws_hash, aws_hash_destroy, aws_hash_finalize, aws_hash_update, aws_sha256_new, AWS_SHA256_LEN,
 };
 
 use crate::common::allocator::Allocator;
+use crate::common::byte_buf::ByteBuf;
 use crate::common::error::Error;
 use crate::{CrtError as _, ToAwsByteCursor};
 
@@ -60,20 +60,13 @@ impl Sha256Hasher {
 
     /// Finalize the hash state and return the computed SHA256 checksum value.
     pub fn finalize(self, allocator: &Allocator) -> Result<Sha256, Error> {
-        // SAFETY: allocator is a valid aws_allocator, and `aws_byte_buf_init` initializes the buffer on success.
-        let mut buffer = unsafe {
-            let mut buffer: MaybeUninit<aws_byte_buf> = MaybeUninit::uninit();
-            aws_byte_buf_init(buffer.as_mut_ptr(), allocator.inner.as_ptr(), Sha256::LENGTH).ok_or_last_error()?;
-            buffer.assume_init()
-        };
+        let mut buffer = ByteBuf::new(allocator, Sha256::LENGTH)?;
 
         // SAFETY: `self.inner` is a valid `aws_hash` and `buffer` was initialized above.
-        unsafe { aws_hash_finalize(self.inner.as_ptr(), &mut buffer, 0).ok_or_last_error()? };
+        unsafe { aws_hash_finalize(self.inner.as_ptr(), buffer.as_mut_ptr(), 0).ok_or_last_error()? };
 
-        // SAFETY: `buffer` holds a valid buffer.
-        let output = unsafe { std::slice::from_raw_parts(buffer.buffer, buffer.len) };
-
-        Ok(Sha256(output.try_into().unwrap()))
+        // Slice will be copied into the struct.
+        Ok(Sha256(buffer.as_slice().try_into().unwrap()))
     }
 }
 

mountpoint-s3-crt/src/common.rs

@@ -7,6 +7,7 @@ use mountpoint_s3_crt_sys::*;
 use crate::common::allocator::Allocator;
 
 pub mod allocator;
+pub mod byte_buf;
 pub mod error;
 pub mod logging;
 pub mod ref_count;

mountpoint-s3-crt/src/common/byte_buf.rs

@@ -0,0 +1,62 @@
+//! Implements safe wrappers for interacting with CRT [aws_byte_buf].
+
+use mountpoint_s3_crt_sys::*;
+
+use std::mem::MaybeUninit;
+
+use crate::common::allocator::Allocator;
+use crate::common::error::Error;
+use crate::CrtError as _;
+
+/// Rust wrapper for a CRT [aws_byte_buf].
+///
+/// This does not support pointing to constant memory at this time (as an allocator is currently required for this struct).
+/// See CRT documentation for more information.
+///
+/// This type does not implement [Copy] or [Clone].
+/// It is not safe to simply clone the struct, we must allocate a new buffer using [aws_byte_buf_init_copy].
+#[derive(Debug)]
+pub struct ByteBuf {
+    /// Inner struct, representing the buffer.
+    ///
+    /// We own this struct and must not lose it, or we lose the ability to access and free the buffer.
+    pub(crate) inner: aws_byte_buf,
+}
+
+impl ByteBuf {
+    /// Create a new [ByteBuf] (backed by [aws_byte_buf]) with the given [Allocator] and capacity.
+    pub fn new(allocator: &Allocator, capacity: usize) -> Result<Self, Error> {
+        // SAFETY: Allocator is valid, we allocate the struct and immediately ask the CRT to initialize it (or error).
+        let inner = unsafe {
+            let mut inner: MaybeUninit<aws_byte_buf> = MaybeUninit::uninit();
+            aws_byte_buf_init(inner.as_mut_ptr(), allocator.inner.as_ptr(), capacity).ok_or_last_error()?;
+            inner.assume_init()
+        };
+        let buf = Self { inner };
+        Ok(buf)
+    }
+
+    /// Get out the inner pointer to the [aws_byte_buf].
+    ///
+    /// This is useful for passing the buffer to CRT functions which take a pointer.
+    pub fn as_mut_ptr(&mut self) -> *mut aws_byte_buf {
+        &raw mut self.inner
+    }
+
+    /// Provide a slice into the underlying bytes for this [ByteBuf].
+    pub fn as_slice(&self) -> &[u8] {
+        // SAFETY: We know that the underlying buffer is valid and we have an immutable reference to it.
+        //         The slice lifetime will be tied to this struct.
+        unsafe { std::slice::from_raw_parts(self.inner.buffer, self.inner.len) }
+    }
+}
+
+impl Drop for ByteBuf {
+    fn drop(&mut self) {
+        // SAFETY: We know that the [ByteBuf] will be dropped at this point.
+        //         No other pointers should refer to the memory of the underlying buffer as we do not implement Copy or Clone.
+        unsafe {
+            aws_byte_buf_clean_up(&mut self.inner);
+        }
+    }
+}