[*] Harden AuByteBuffer against use after free

This commit is contained in:
Reece Wilson 2024-03-18 08:36:11 +00:00
parent c54ac9d6a3
commit b65f27fa8f
8 changed files with 128 additions and 7 deletions

View File

@ -42,6 +42,8 @@ namespace Aurora::IO::Buffered
return EStreamError::eErrorStreamNotOpen;
}
__audetail::BufferLock(pBuffer.get());
if (pBuffer->HasStreamError())
{
return EStreamError::eErrorByteBuffer;

View File

@ -34,6 +34,8 @@ namespace Aurora::IO::Buffered
return EStreamError::eErrorStreamNotOpen;
}
__audetail::BufferLock(this->pBuffer_.get());
if (this->pBuffer_->empty())
{
return EStreamError::eErrorEndOfStream;

View File

@ -34,6 +34,8 @@ namespace Aurora::IO::Buffered
return EStreamError::eErrorStreamNotOpen;
}
__audetail::BufferLock(this->pBuffer_.get());
if (this->pBuffer_->empty())
{
return EStreamError::eErrorEndOfStream;

View File

@ -41,6 +41,8 @@ namespace Aurora::IO::Buffered
return EStreamError::eErrorStreamNotOpen;
}
__audetail::BufferLock(pBuffer.get());
if (pBuffer->HasStreamError())
{
return EStreamError::eErrorByteBuffer;

View File

@ -7,6 +7,87 @@
***/
#pragma once
namespace __audetail
{
struct BufferLock
{
Aurora::Memory::ByteBuffer *pThat;
AuUInt8 *pIDC {};
inline BufferLock(Aurora::Memory::ByteBuffer *pThat) :
pThat(pThat)
{
if (pThat)
{
AuAtomicAdd(&pThat->uInUseCounter, 1u);
this->StrongLoad();
pThat->PrivateUserDataToUtilityMutex()->Lock();
}
}
inline ~BufferLock()
{
if (this->pThat)
{
AuAtomicSub(&pThat->uInUseCounter, 1u);
pThat->PrivateUserDataToUtilityMutex()->Unlock();
}
}
private:
#pragma optimize("", off)
void StrongLoad()
{
// Strict load now. After this point, the compiler can lazy load with using these loads in its' emitted code or via the CPU cache
// It would be nice to make these properly volatile without impacting our other high perf paths, somehow
pIDC = AuAtomicLoad(&pThat->base);
pIDC = AuAtomicLoad(&pThat->readPtr);
pIDC = AuAtomicLoad(&pThat->writePtr);
}
#pragma optimize("", on)
};
struct BufferAllocLock
{
Aurora::Memory::ByteBuffer *pThat;
AuUInt8 *pIDC {};
inline BufferAllocLock(Aurora::Memory::ByteBuffer *pThat) :
pThat(pThat)
{
if (pThat)
{
this->StrongLoad();
pThat->PrivateUserDataToUtilityMutex()->Lock();
}
}
inline ~BufferAllocLock()
{
if (this->pThat)
{
pThat->PrivateUserDataToUtilityMutex()->Unlock();
}
}
private:
#pragma optimize("", off)
void StrongLoad()
{
// Strict load now. After this point, the compiler can lazy load with using these loads in its' emitted code or via the CPU cache
// It would be nice to make these properly volatile without impacting our other high perf paths, somehow
pIDC = AuAtomicLoad(&pThat->base);
pIDC = AuAtomicLoad(&pThat->readPtr);
pIDC = AuAtomicLoad(&pThat->writePtr);
}
#pragma optimize("", on)
};
}
#include "BlobReader.hpp"
#include "BlobSeekableReader.hpp"
#include "BlobWriter.hpp"

View File

@ -9,6 +9,7 @@
#include <Aurora/Locale/ECodePage.hpp>
#include <Aurora/Utility/PrivData.hpp>
#include <Aurora/Threading/Waitables/FutexWaitable.hpp>
namespace Aurora::Memory
{
@ -611,6 +612,9 @@ namespace Aurora::Memory
return -(pSubtrahend - this->base);
}
}
private:
inline auline bool Allocate2(AuUInt length, AuUInt alignment);
inline auline bool Allocate2(AuUInt length);
};
struct SharedByteBuffer : ByteBuffer
@ -628,7 +632,9 @@ namespace Aurora::Memory
this->memory = pWriteView;
}
inline SharedByteBuffer(AuSPtr<void> pRAIIParentOwner, MemoryViewWrite view) : ByteBuffer()
inline SharedByteBuffer(AuSPtr<void> pRAIIParentOwner, MemoryViewWrite view) :
ByteBuffer(),
__view(view)
{
this->allocSize = 0;
this->base = (AuUInt8 *)view.ptr;
@ -638,6 +644,9 @@ namespace Aurora::Memory
this->flagNoFree = true;
this->memory = pRAIIParentOwner;
}
private:
MemoryViewWrite __view;
};
struct SharableByteBuffer : ByteBuffer,

View File

@ -10,6 +10,12 @@
namespace Aurora::Memory
{
bool ByteBuffer::Allocate(AuUInt length)
{
__audetail::BufferAllocLock lock(this);
return Allocate2(length);
}
bool ByteBuffer::Allocate2(AuUInt length)
{
if (this->flagNoFree ||
this->flagNoRealloc ||
@ -54,6 +60,12 @@ namespace Aurora::Memory
}
bool ByteBuffer::Allocate(AuUInt length, AuUInt alignment)
{
__audetail::BufferAllocLock lock(this);
return Allocate2(length, alignment);
}
bool ByteBuffer::Allocate2(AuUInt length, AuUInt alignment)
{
if (this->flagNoFree ||
this->flagNoRealloc ||
@ -93,6 +105,8 @@ namespace Aurora::Memory
bool ByteBuffer::SetBuffer(MemoryViewRead readView, bool bMoveWriteHeadForReaders)
{
__audetail::BufferAllocLock lock(this);
if (this->flagNoFree ||
this->flagNoRealloc ||
this->uInUseCounter)
@ -100,7 +114,7 @@ namespace Aurora::Memory
return false;
}
if (!Allocate(readView.length))
if (!Allocate2(readView.length))
{
return false;
}
@ -123,6 +137,8 @@ namespace Aurora::Memory
void ByteBuffer::Reset()
{
__audetail::BufferAllocLock lock(this);
if (this->uInUseCounter)
{
return;
@ -145,16 +161,18 @@ namespace Aurora::Memory
void ByteBuffer::Reserve(AuUInt length)
{
auto oldLength = this->length;
auto oldLength = AuAtomicLoad(&this->length);
if (length > this->allocSize)
{
this->Resize(length);
}
this->length = AuMin(oldLength, this->length);
this->length = AuMin(oldLength, AuAtomicLoad(&this->length));
}
void ByteBuffer::GC()
{
__audetail::BufferAllocLock lock(this);
if (this->flagNoFree ||
this->flagNoRealloc ||
this->uInUseCounter)
@ -214,6 +232,8 @@ namespace Aurora::Memory
bool ByteBuffer::Resize(AuUInt length)
{
__audetail::BufferAllocLock lock(this);
if (this->flagNoFree ||
this->flagNoRealloc ||
this->uInUseCounter)
@ -231,7 +251,7 @@ namespace Aurora::Memory
if (!this->base)
{
return Allocate(length);
return Allocate2(length);
}
auto uOldHead = this->readPtr - this->base;

View File

@ -229,6 +229,8 @@ namespace Aurora::Memory
MemoryViewRead ByteBuffer::GetNextLinearRead()
{
__audetail::BufferAllocLock lock(this);
AuUInt8 *pBase {};
AuUInt uCount {};
@ -273,6 +275,7 @@ namespace Aurora::Memory
MemoryViewWrite ByteBuffer::GetNextLinearWrite()
{
__audetail::BufferAllocLock lock(this);
AuUInt8 *pBase {};
AuUInt uCount {};
@ -336,7 +339,7 @@ namespace Aurora::Memory
MemoryViewRead ByteBuffer::GetLinearReadableForAtleast(AuUInt length)
{
auto view = this->GetNextLinearRead();
return view.length < length ? MemoryViewRead {} : MemoryViewRead(view.ptr, view.length);
return view.length < length ? MemoryViewRead {} : MemoryViewRead(view.ptr, view.length, &this->uInUseCounter);
}
MemoryViewWrite ByteBuffer::GetLinearWriteableForAtleast(AuUInt length)
@ -347,7 +350,7 @@ namespace Aurora::Memory
{
return this->GetOrAllocateLinearWriteable(length);
}
return view.length < length ? MemoryViewWrite {} : MemoryViewWrite(view.ptr, view.length);
return view.length < length ? MemoryViewWrite {} : MemoryViewWrite(view.ptr, view.length, &this->uInUseCounter);
}
MemoryViewRead ByteBuffer::GetLinearReadableForNoMore(AuUInt length)