Tricky borrow checker problems in the rustls deframer

Recently I've been working to improve efficiency on the data receive path in rustls, the TLS implementation in Rust. In a recent PR I was able to merge the deframing (parsing messages -- called records in TLS -- out of the byte stream) and joining of handshake messages. The outer records in TLS can be too short for some kinds of handshake messages, so sometimes records need to be joined together to be able to parse the handshake message. So far, rustls has allocated and copied data for decrypting incoming records and joining handshake messages and now I'd like to try to get it to do decryption and joining in-place, minimizing copies (which cannot be easily avoided for joining process) and allocations.

However, although I think my design should conceptually match what the borrow checker allows, I'm having some trouble getting the compiler to allow it in practice. The toy version of the problem is somewhat like this:

use std::borrow::Cow;

struct Buffer {
    buf: Vec<u8>,
}

impl Buffer {
    fn pop<'a>(&'a mut self) -> Option<Borrowed<'a>> {
        loop {
            let mutable = Mutable { buf: &mut self.buf[..10] };
            if mutable.buf[0] == 0 {
                return Some(mutable.borrow());
            }
        }
    }
}

struct Mutable<'a> {
    buf: &'a mut [u8],
}

impl<'a> Mutable<'a> {
    fn borrow(self) -> Borrowed<'a> {
        Borrowed {
            buf: (&*self.buf).into()
        }
    }
}

struct Borrowed<'a> {
    buf: Cow<'a, [u8]>,
}

(Playground.) -- updated with non-Iter version.

If you look at my draft PR, the problem is in the MessageDeframer::pop() method (in the last commit), where the flow should be something like this:

  • While we don't have a complete handshake message, look at the incoming record stream:
    • Try to parse a record (encrypted length-prefixed message, called OpaqueMessage here)
    • If this an unencrypted message, convert it to a PlainMessage referencing the same MessageDeframer::buf data and yield it to the caller
    • Otherwise try to decrypt it in place using the RecordLayer
    • If it's not a handshake message, yield the decrypted message to the caller
    • If it is, make note of that in the HandshakePayloadMeta data and move the payload to the correct place in the buffer if necessary for joining
  • If we have enough data for a full handshake message now, yield that to the caller

In my mind conceptually this should work while keeping a lifetime to the &mut MessageDeframer around, and yielding message types that carry the same lifetime, which can be handled in code that's downstream from the message deframer in terms of the control flow graph. However, I'm running into what I seem to recall are borrowck limitations in making sense of loops with conditional return statements, and the diagnostics aren't helping me here. I got some help suggesting that I split_mut_at() the buffer and this seemed to help a little, but I don't think that is a complete solution in the sense that, in order to do handshake joining, I need to drop the msg anyway so that I can use my knowledge of where the referenced payload was in the buffer to join it into place.

Anyone got some suggestions on how to attack this problem and/or a notion of why I'm wrong that this should be conceptually allowed in today's borrow checker?

No, why should that work? That would mean aliasing a mutable borrow with another (it doesn't matter whether the returned, aliasing borrow is mutable or immutable).

The way this restriction manifests as a concrete compiler error is that you can't return an &'a mut T to a field of &'_ mut self (where '_ indicates that the two lifetimes are unrelated). This makes sense – if this were allowed, then the example Playground that you provided would trivially allow you to return a & [_] to a subslice, and then still point a &mut [_] to the same subslice upon the next call:

let mut buf = Buffer { buf: vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9] };
let mut iter = Iter { buf: &mut buf };
let b1 = iter.pop();
let b2 = iter.pop();
dbg!(b1, b2);

If you force the lifetimes to be accepted using unsafe, and run the above code under Miri, then it clearly indicates Undefined Behavior due to aliased mutability.

In fact, here's an even worse example demonstrating mutability of an immutable slice :cold_face:

1 Like

So there's two things going on here.

Thing 1 is: you can't return an Option<Borrowed<'a>> because it's only borrowed for 'this, which is the name I'm giving the lifetime of &self: the toy code as-is is equivalent to this:

impl<'a> Iter<'a> {
    fn pop<'this>(&'this mut self) -> Option<Borrowed<'a>> {
        // cut
    }
}

Changing it to this gets rid of "lifetime may not live long enough:

impl<'a> Iter<'a> {
    fn pop<'this>(&'this mut self) -> Option<Borrowed<'this>> {
        // cut
    }
}

But at this point, you can just elide lifetimes:

impl<'a> Iter<'a> {
    fn pop(&mut self) -> Option<Borrowed<'_>> {
        // cut
    }
}

Thing 2 is: the main borrow checker is a bit silly with loops. This code doesn't borrow-check under the main borrow checker, but it does with Polonius:

$ RUSTFLAGS="-Z polonius" cargo +nightly check

I don't think that's a viable option for rustls (neither is waiting around for it to become the "main" borrow checker - I believe the types team is on it, but it could take a few years). I'll comment again if I can find a workaround.

2 Likes

Here's a workaround I found, I'm not sure if it applies for the non-toy code (I haven't read it yet):

impl<'a> Iter<'a> {
    fn pop(&mut self) -> Option<Borrowed<'_>> {
        loop {
            let mutable = Mutable {
                buf: &mut self.buf.buf[..10],
            };
            if mutable.buf[0] == 0 {
                let mutable = Mutable {
                    buf: &mut self.buf.buf[..10],
                };
                return Some(mutable.borrow());
            }
        }
    }
}

With this, self.buf.buf is reborrowed properly. I have a feeling this isn't what you want, and you would like to keep 'a everywhere.

1 Like

Sorry, I've since gone back and dropped the Iter proxy, which I think can be worked around some other way, so we can borrow directly from the &mut self argument to pop() on the Buffer itself.

In general, it should be fine to hold the mutable borrow on the Buffer/MessageDeframer for the lifeftime of the yielded Borrowed/PlainMessage.

I do remember that the borrow checker is silly around loops, but I forgot again exactly what the restrictions are and/or what any workarounds are.

According to @kornel in When is transmuting lifetimes useful? - #4 by kornel

Mutable iterators are a common case where this is necessary. There's no way to explain to the borrow checker that next() returns a different, non-overlapping mutable borrow each time if they all come from the same source. This can be reduced to implementation of slice.split_at_mut() , which also requires such fudge.

And indeed, this works:

impl Buffer {
    fn pop(&mut self) -> Option<Borrowed<'_>> {
        loop {
            let mutable = Mutable {
                buf: &mut self.buf[..10],
            };
            if mutable.buf[0] == 0 {
                let mutable: Mutable<'_> = unsafe { std::mem::transmute(mutable) };
                return Some(mutable.borrow());
            }
        }
    }
}

..but I certainly wouldn't feel good shipping that (even though miri seems happy with it)


@steffahn below (I can't reply anymore, Discourse hates me because I just signed up for it) mentioned polonius-the-crab, for the curious here's how it would look for the toy case. (This doesn't actually build on the playground, I guess it's not one of the top 100 crates)

1 Like

polonius-the-crab to the rescue :slight_smile:

1 Like

We have a pretty strong no unsafe policy for rustls, so I don't think transmute()ing the lifetime away will work.

In my understanding for something like this, because the Borrowed<'_> lifetime is derived from the &mut self reference, the borrow checker won't allow calling pop() a second time while an early Borrowed<'_> is still alive, and that's not an issue in this case.

(By the way… click the second link if you haven’t already; it’s a fork of the PR with the lifetime issue fixed using said crate.)


I mean this as a serious option. The soundness story of this crate (of @Yandros) is remarkably convincing. It might seem dodgy at first because of all those macros, but those macros are mere convenience wrappers and irrelevant for the soundness. They are convenience wrappers for the following API

/// See the [top-level docs][crate] for more info.
pub trait WithLifetime<'lt> {
    // Note: the `&'lt Self` implicit bound hack is,
    // for once, unnecessary.
    type T;
}

/// See the [top-level docs][crate] for more info.
pub trait HKT
where
    Self: for<'any> WithLifetime<'any>,
{
}

impl<T: ?Sized> HKT for T where Self: for<'any> WithLifetime<'any> {}

/// See the [top-level docs][crate] for more info.
pub fn polonius<Ret: ?Sized + HKT, State: ?Sized, Err, F>(
    state: &mut State,
    branch: F,
) -> Result<<Ret as WithLifetime<'_>>::T, (&'_ mut State, Err)>
where
    F: FnOnce(&'_ mut State) -> Result<<Ret as WithLifetime<'_>>::T, Err>,
{
    let err = {
        #[cfg(not(feature = "polonius"))]
        let state = unsafe {
            // SAFETY:
            // > Though this be `unsafe`, there is soundness in 't.
            //
            // More seriously, read the docs, I've detailed the reasoning there
            // in great length. And/or check the `tests/soundness.rs` test,
            // which `cargo check`s this very snippet without this `unsafe`.
            &mut *(state as *mut _)
        };
        match branch(state) {
            Ok(ret) => return Ok(ret),
            Err(err) => err,
        }
    };
    Err((state, err))
}

That’s a sufficiently generic to cover a lot of use-cases of polonius. And the funciton, with the unsafe lifetime-extending code removed, passes under -Z polonius, so it’s essentially safe Rust (unless -Z polonius is behaving unsoundly here).

1 Like

Wow, very cool! Now, I do wonder if the CPS-style code mentioned in the polonius-the-crab README could work for this use case, which seems like a leaner abstraction if it works.

1 Like

Looking at the use-cases, AFAICT that might actually work.


Edit: It’s not quite as straighforward. The record_layer: &mut RecordLayer argument would be kept alive during a callback to a continuation, and the use-sites want to mutably borrow the value the RecordLayer is contained in. Maybe solvable with some refactoring, and passing the borrow along to the continuation.

Also one of the use-sites wants to more out of and then back into a local variable after the call to pop. Which is not possible directly in a closure; thought the variable could be wrapped in Option, I guess, and then Option::take used. Or moved in and out of the closure…

I implemented the callback for the fun of it. It is a bit tedious to call indeed.

Comparing rustls:7207ef732d74c4a408cbca24983b4714107910cb...steffahn:callback · rustls/rustls · GitHub

At that point you could reach for another macro-sugar-providing crate such as with_locals - Rust :laughing:


More seriously, regarding

I'd welcome ideas to make the macros easier to use: indeed, as is usually the case with macros, the diagnostics are quite bad when misused, and the API itself is complex enough that this time it is very likely for user to end up misusing it

  • (heck, for more complex loop cases I can struggle to use them properly myself! :sweat_smile:)

One simple idea already would be to use proc-macros to replace the 'polonius lifetime with the classic elided one (like lending-iterator's HKT! macro does), but that would require proc-macro processing and thus involving a proc-macro crate, which I'm not very fond of using (although maybe a macro_rules! polyfill would not be that hard to write/maintain?).

Another embryo of an idea would be to do like the aforementioned ::with_locals crate does, and allow for some attribute that would automagically replace returns and breaks with macros doing it.

Combining both ideas, we'd end up with:

//                 &'_ String     [alternative]
polonius!(|map| -> &String { // <- rather than `&'polonius String`
    if let Some(v) = map.get(&22) {
        return v; // <- rather than `polonius_return!`
    }
});

You can use rustexplorer.com which uses way more crates from crates.io (enough for polonius-the-crab to pass the bar): demo

2 Likes

I tried polonius_loop first before realizing that I need access to buf later in the loop body that would need to happen after the polonius! incovation. The downside was then that I needed to re-implement the convenient control-flow enum stuff (for continue) myself; and passing around some local variables was necessary. I don’t quite see any good idea how a general macro could help better here. Maybe with the option to offer support for continue etc…, though I also like the minimalism of not using what I don’t need. And w.r.t. local variables, I had to choose which ones should be dropped (msg in particular) and which ones kept, anyways.

The diff necessary in the end was quite small anyways, and some first attempts were necessary to correctly understand the code anyways. (In case you look at the GitHub commit, note that the changes involving payload.clone() are fixing an independent compilation error.)

This is great! Probably going to spend tonight to build a cleaner way of doing this, which I think should definitely be feasible.