Figure out a way to express PacketBuffer ownership in the type system

[bzbarsky-apple]

Problem

If a function takes a PacketBuffer *, it's not obvious from the function signature whether it's expected to free the buffer or not. We can document the semantics, but that still leaves a lot of room for error, and for the function failing to actually implement the semantics properly.

Proposed Solution

Express "this function will free the buffer" in the type system. It seems to me that we want to be able to express the following concepts:

1. A function is handed "ownership" of a PacketBuffer and is expected to free it.
2. We are now handing "ownership" of our PacketBuffer to someone else and should not reference it anymore.
3. We are now incrementing the refcount on a PacketBuffer, after which we will hand over one of the references to someone else and be responsible for releasing the other ourselves.

We could perhaps use std::shared_ptr for this, with std::move at the ownership transfers. Functions expecting ownership would take shared_ptr<PacketBuffer>&&. It's not completely ideal because nothing prevents use of a "moved out of" shared_ptr from being accessed, but at least it would be a null deref, not a UAF.

We could also just pass something like shared_ptr by value, if we think our refcounting is cheap enough.

Ideally, of course, we would have some sort of clonable PacketBuffer handle so we can model the increment explicitly as a clone and we'd have lifetimes a la Rust so when you hand the handle to someone the compiler won't let you touch it again. Sadly, I don't see a way to do that in C++; you can always mess with something like a unique_ptr you called release on...

@sagar-apple @pan-apple @andy31415 @gerickson @turon @rwalker-apple any thoughts or suggestions?

[issue-label-bot]

Issue-Label Bot is automatically applying the label feature_request to this issue, with a confidence of 0.64. Please mark this comment with 👍 or 👎 to give our bot feedback!

Links: app homepage, dashboard and code for this bot.

[mspang]

Shouldn't it be std::unique_ptr not std::shared_ptr ? Unique ownership captures the main use cases we're considering here; shared ownership makes lifetime much harder to nail down.

It seems to be that std::unique_ptr, with a single return statement should be comparable to what we have now in terms of code size, while being much safer. So that seems like a good idea to me.

[bzbarsky-apple]

Shouldn't it be std::unique_ptr not std::shared_ptr ? Unique ownership captures the main use cases we're considering here; shared ownership makes lifetime much harder to nail down.

The whole issue is that not all ownership here is unique. There are several places where we do a manual addref before passing a buffer down into code that "consumes" it, because that code doesn't actually consume everything we have in the buffer.... This is item 3 in my list above.

[mspang]

std::shared_ptr will introduce atomics, weak ptr support, and other things we probably don't want to pay for.

[bzbarsky-apple]

I am perfectly comfortable with us inventing a PacketBufferHandle class that has the minimal set of semantics that we want, if people are ok with us doing our own thing.

I think we want:

1. Construction from a PacketBuffer* addrefs the PacketBuffer.
2. Construction from a PacketBufferHandle&& steals the value.
3. Default-construction is allowed.
4. All other constructors (specifically the copy ctor) are deleted.
5. Construction from PacketBuffer* is protected; handles are vended by either PacketBuffer factory methods or by cloning an existing handle.
6. Assignment operators other than one taking PacketBufferHandler&& are deleted.
7. Destruction frees the PacketBuffer.

I think that more or less satisfies our use case, right?

[rwalker-apple]

I am perfectly comfortable with us inventing a PacketBufferHandle class that has the minimal set of semantics that we want, if people are ok with us doing our own thing.

I think we want:

1. Construction from a PacketBuffer* addrefs the PacketBuffer.
2. Construction from a PacketBufferHandle&& steals the value.
3. Default-construction is allowed.
4. All other constructors (specifically the copy ctor) are deleted.
5. Construction from PacketBuffer* is protected; handles are vended by either PacketBuffer factory methods or by cloning an existing handle.
6. Assignment operators other than one taking PacketBufferHandler&& are deleted.
7. Destruction frees the PacketBuffer.

I think that more or less satisfies our use case, right?

yeah, I think that's complete, and would be a big improvement

[bzbarsky-apple]

Here's what I have so far:

// Implementation inlined to make this zero-cost as much as possible.
class DLL_EXPORT PacketBufferHandle
{
public:
    PacketBufferHandle() : buffer(nullptr) {}
    PacketBufferHandle(decltype(nullptr)) : buffer(nullptr) {}
    PacketBufferHandle(PacketBufferHandle && aOther) 
    {
        buffer        = aOther.buffer;
        aOther.buffer = nullptr;
    }
    ~PacketBufferHandle()
    {
        // PacketBufferHandle::Free is null-safe, but maybe we should check for
        // null here to avoid the pool locking it does?
        PacketBuffer::Free(buffer);
    }
    PacketBufferHandle & operator=(PacketBufferHandle && aOther)
    {
        if (buffer)
        {
            PacketBuffer::Free(buffer);
        }
        buffer        = aOther.buffer;
        aOther.buffer = nullptr;
        return *this;
    }
    PacketBuffer * operator->() { return buffer; }
    // FIXME: or an explicit empty() or something?
    explicit operator bool() { return !!buffer; }

    void Clear()
    {
        if (buffer)
        {
            PacketBuffer::Free(buffer);
            buffer = nullptr;
        }
    }

    PacketBufferHandle Clone() const { return PacketBufferHandle(buffer); }

private:
    PacketBufferHandle(const PacketBufferHandle &) = delete;
    PacketBufferHandle & operator=(const PacketBufferHandle &) = delete;
    friend class PacketBuffer;
    PacketBufferHandle(PacketBuffer * aBuffer) : buffer(aBuffer) { buffer->AddRef(); }
    PacketBuffer * buffer;
};

I would propose that this be used as follows:

• Functions that conceptually "consume" the buffer would take a PacketBufferHandle by value. This would (1) force callers to std::move at the callsite and (2) guarantee that after the call returns the caller has a nulled-out handle no matter what codepath the callee took. Another option here would be to take PacketBufferHandle &&, which would allow zero-cost forwarding of the argument by just passing the reference along instead of running constructor/destructor code at every call, but it would lead to a slightly more complex mental model in terms of who "owns" the value, especially on error returns. I guess a function that guarantees that it always calls a consuming function could still take PacketBufferHandle&& and maintain the "I always consume it" invariant.
• Places that need to store a PacketBuffer would store a handle, and assign to it via the assignment operator taking an rvalue reference.
• Places that need to return a PacketBuffer would return a PacketBufferHandle by value.

I suspect that there would need to be one big PR converting the whole tree to handles or something like that, though if we want to avoid that we could add a temporary forget() method that would allow us to go from a handle to a raw PacketBuffer* that someone else would then have to free. Then we could convert the tree incrementally, starting with SystemPacketBuffer.h/cpp with forget() at all the PacketBuffer::New* callsites and then working out from there until there are no more forget() calls in the tree.

[andy31415]

We may want some Aquire and Release semantics as well (I added such things in #2987)

It would be nice if we could gradually change code rather than an 'all-or-nothing approach'. Usually one would move classes without change into some 'deprecated' (maybe have using DeprecatedPacketBuffe = PacketBuffer) and then work towards updating things.

I would like ownership to be better for sure - I had bugs because of this.

[bzbarsky-apple]

We may want some Aquire and Release semantics as well

Given that we always allocate PacketBuffers ourselves (so don't need to Acquire things coming from outside our code), the one case I can think that might need Acquire/Release is if we have to pass a PacketBuffer to some third-party code that basically takes a void*, with ownership being enforced by some out-of-band mechanism (i.e. we pass it a pointer, it eventually calls some callback with that pointer and we free in the callback). If we run into that situation, we can add something to enable it, but I'd rather wait to see whether this becomes an issue in practice.

[rwalker-apple]

note that we don't always allocate PacketBuffers "ourselves", many come up from lwip.

[bzbarsky-apple]

Ah, in that case we'll need an Acquire/Adopt thing, yes.

[tcarmelveilleux]

note that we don't always allocate PacketBuffers "ourselves", many come up from lwip.

This is spot-on. Note, as well, that PacketBuffers are implemented on top of lwip pbuf, and there is thus a reference count. PacketBuffers should not be considered to be like "heap". They are kinda more like shared pointers, but they are not direct payload buffer either (not uint8_t*, but rather header + payload). In the underlying pbuf implementation, the pbuf may be linearly allocated header + payload from a heap, or may be passed back from a memory pool, with the payload itself being a reference, or coming from a heap (lwip internal or external allocator) or buffer pools. Lots of companies (including Google Nest) have optimized pbuf allocation code to fit different constraints over time. It's important not to consider a "one size fits all, PacketBuffer means dynamic byte buffer" as this break royally in subtle ways when mixed with different flavors of lwip.

[bzbarsky-apple]

Right, to be clear, what we are trying to build here is something to do automated refcounting of PacketBuffers. We are not making any assumptions about heaps or allocation strategies or anything like that. Just making sure that refcounting happens in the right places with minimal space for getting it wrong.

[bzbarsky-apple]

@kpschoedel do you have any specific planned timeframe for this? This is continuing to cause pain in new code we write.....

[kpschoedel]

@kpschoedel do you have any specific planned timeframe for this? This is continuing to cause pain in new code we write.....

I hadn't, as I was focusing on collection memory use information, but I can start to prioritize this.