Add UninitSlice::split_at_mut

[roblabla]

Adds UninitSlice::split_at_mut, a thin wrapper around slice::split_at_mut. This can make certain types of code much easier to write.

[carllerche]

Thanks for the PR. Could you provide the context in which you needed this?

[roblabla]

The use-case comes from writing a Tonic Codec implementation that encrypts the ProtoBuf payloads before sending them.

The Tonic Encoder trait has this method we need to implement:

pub fn encode(
    &mut self,
    item: Self::Item, // Our protobuf message, impls prost::Message
    dst: &mut EncodeBuf<'_> // impls BufMut
) -> Result<(), Self::Error>;

What we want to do is send three pieces of information: A nonce/IV, the actual encrypted payload, and a MAC. To make this easier and less error-proof, my implementation starts by splitting the chunk into three different mutable sections of the right size:

let msg_size = item.encodded_len();
let ciphered_msg_size = msg_size + MAC_SIZE + NONCE_SIZE;
dst.reserve(ciphered_msg_size);
let bytes = dst.chunk();

let (bytes_nonce, bytes) = bytes.split_at_mut(NONCE_SIZE);
let (bytes_payload, bytes) = bytes.split_at_mut(msg_size);
let (bytes_mac, _unused_bytes) = bytes.split_at_mut(MAC_SIZE);

// Generate nonce, write it to bytes.
let nonce = todo!("Generate nonce");
bytes_nonce.copy_from_slice(nonce);

// Generate payload. This requires https://github.com/roblabla/bytes/commit/d7e288d0e2974ac85521697a9365257903ad3176
{
    let mut bytes_msg = &mut bytes_payload[..];
    item.encode(&mut bytes_msg).unwrap(); // Takes a &mut impl BufMut
    // Assert that the whole bytes_payload slice got initialized, so we can safely turn it into a [u8].
    assert!(bytes_msg.len() == 0);
}

let bytes_msg_slice = unsafe {
    // Safety: We can turn bytes_payload into a slice safely, as it has
    // been fully initialized by item.encode.
    slice::from_raw_parts_mut(bytes_payload.as_mut_ptr(), bytes_payload.len())
};
let mac = todo!("Encrypt bytes_msg_slice in-place");
bytes_mac.copy_from_slice(mac); 
unsafe { dst.advance(ciphered_msg_size); }

I find this easier to reason about and make sure that it all works. I could technically make do without it, but split_at_mut makes it easier to understand what each region of the slice contains, and it means the assertion necessary to turn the UninitSlice into a [u8] is straightforward to verify.

---

Separately from this MR, I also implemented BufMut on &mut UninitSlice in https://github.com/tokio-rs/bytes/commit/15a3834dfae6abb7cb3e58a80374d12bfa8cb3ed which I haven't submitted as a PR because it depends on this method (though that's an implementation detail that could be worked around).

[Darksonn]

let bytes_msg_slice = unsafe {
    // Safety: We can turn bytes_payload into a slice safely, as it has
    // been fully initialized by item.encode.
    slice::from_raw_parts_mut(bytes_payload.as_mut_ptr(), bytes_payload.len())
};

We might want to add an unsafe method that does this for us too, since what you are doing here detaches the lifetimes, giving up some compile-time checks.

[seanmonstar]

I've been wondering if we need to flesh out UninitSlice more, when the BufMut trait is hopefully the way that you fill in most bytes. My feeling is that UninitSlice's purpose is to pass it to FFI like read(2) or similar. If just using Rust code, I'd think you could use the methods from BufMut easily.

Such as in your example, instead of UninitSlice::copy_from_slice, you could just use BufMut::put_slice, right?

[roblabla]

@seanmonstar the problem with put_slice is that BufMut doesn't allow me to further mutate the slice after it's been written. In my case, I first serialize a type to a slice, then encrypt that slice, preferably in-place.

BufMut doesn't allow accessing or modifying the bytes that are initialized, so there is no simple way to do what I want to do:

• If I let serialize use put_slice, I can't modify it afterwards, so can't encrypt.
• The other solution is to first serialize to an intermediate Vec before encrypting this and finally copying it to our BufMut, but that kinda stinks: It requires having a heap allocator, and is mostly needed to work around an overly restrictive API.

The only solution I see to the above is to go through chunk, but that is currently very unpleasant due to the restrictive API.

[seanmonstar]

BufMut doesn't allow me to further mutate the slice after it's been written

Ah, sure, that makes sense, thanks for clarifying!