I recently found a small buffer overflow in Firefox’s SOCKS support, and came up with a nice-ish way to make it a C++ compilation error when it may happen with some template magic.

A simplfied form of the problem looks like this:

class nsSOCKSSocketInfo {
public:
    nsSOCKSSocketInfo() : mData(new uint8_t[BUFFER_SIZE]) , mDataLength(0) {}
    ~nsSOCKSSocketInfo() { delete[] mData; }

    void WriteUint8(uint8_t aValue) { mData[mDataLength++] = aValue; }

    void WriteV5AuthRequest() {
        mDataLength = 0;
        WriteUint8(0x05);
        WriteUint8(0x01);
        WriteUint8(0x00);
    }
private:
    uint8_t* mData;
    uint32_t mDataLength;
    static const size_t BUFFER_SIZE = 2;
};

Here, the problem is more or less obvious: the third WriteUint8() call in WriteV5AuthRequest() will write beyond the allocated buffer size. (The real buffer size was much larger than that, and it was a different method overflowing, but you get the point)

While growing the buffer size fixes the overflow, that doesn’t do much to prevent a similar overflow from happening again if the code changes. That got me thinking that there has to be a way to do some compile-time checking of this. The resulting solution, at its core, looks like this:

template  class Buffer {
public:
    Buffer() : mBuf(nullptr) , mLength(0) {}
    Buffer(uint8_t* aBuf, size_t aLength=0) : mBuf(aBuf), mLength(aLength) {}

    Buffer WriteUint8(uint8_t aValue) {
        static_assert(Size >= 1, "Cannot write that much");
        *mBuf = aValue;
        Buffer result(mBuf + 1, mLength + 1);
        mBuf = nullptr;
        mLength = 0;
        return result;
    }

    size_t Written() { return mLength; }
private:
    uint8_t* mBuf;
    size_t mLength;
};

Then replacing WriteV5AuthRequest() with the following:

void WriteV5AuthRequest() {
    mDataLength = Buffer(mData)
        .WriteUint8(0x05)
        .WriteUint8(0x01)
        .WriteUint8(0x00)
        .Written();
}

So, how does this work? The Buffer class is templated by size. The first thing we do is to create an instance for the complete size of the buffer:

Buffer(mData)

Then call the WriteUint8 method on that instance, to write the first byte:

.WriteUint8(0x05)

The result of that call is a Buffer (in our case, Buffer<1>) instance pointing to &mData[1] and recording that 1 byte has been written.
Then we call the WriteUint8 method on that result, to write the second byte:

.WriteUint8(0x01)

The result of that call is a Buffer (in our case, Buffer<0>) instance pointing to &mData[2] and recording that 2 bytes have been written so far.
Then we call the WriteUint8 method on that new result, to write the third byte:

.WriteUint8(0x00)

But this time, the Size template parameter being 0, it doesn’t match the Size >= 1 static assertion, and the build fails.
If we modify BUFFER_SIZE to 3, then the instance we run that last WriteUint8 call on is a Buffer<1> and we don’t hit the static assertion.

Interestingly, this also makes the compiler emit more efficient code than the original version.

Check the full patch for more about the complete solution.

http://glandium.org/blog/?p=3369