Skip to content

Commit

Permalink
enforce isolate lock and track allocations on v8
Browse files Browse the repository at this point in the history
  • Loading branch information
@anonrig
anonrig committed Sep 17, 2024
1 parent d1a7d6e commit 23753d2
Show file tree
Hide file tree
Showing 4 changed files with 57 additions and 44 deletions.
36 changes: 19 additions & 17 deletions src/workerd/api/node/zlib-util.c++
View file
Original file line number Diff line number Diff line change
Expand Up @@ -444,8 +444,8 @@ void ZlibContext::setOutputBuffer(kj::ArrayPtr<kj::byte> output) {

template <typename CompressionContext>
jsg::Ref<ZlibUtil::CompressionStream<CompressionContext>> ZlibUtil::CompressionStream<
CompressionContext>::constructor(ZlibModeValue mode) {
return jsg::alloc<CompressionStream>(static_cast<ZlibMode>(mode));
CompressionContext>::constructor(jsg::Lock& js, ZlibModeValue mode) {
return jsg::alloc<CompressionStream>(js, static_cast<ZlibMode>(mode));
}

template <typename CompressionContext>
Expand Down Expand Up @@ -584,8 +584,9 @@ void ZlibUtil::CompressionStream<CompressionContext>::reset(jsg::Lock& js) {
}
}

jsg::Ref<ZlibUtil::ZlibStream> ZlibUtil::ZlibStream::constructor(ZlibModeValue mode) {
return jsg::alloc<ZlibStream>(static_cast<ZlibMode>(mode));
jsg::Ref<ZlibUtil::ZlibStream> ZlibUtil::ZlibStream::constructor(
jsg::Lock& js, ZlibModeValue mode) {
return jsg::alloc<ZlibStream>(js, static_cast<ZlibMode>(mode));
}

void ZlibUtil::ZlibStream::initialize(int windowBits,
Expand Down Expand Up @@ -752,8 +753,8 @@ kj::Maybe<CompressionError> BrotliDecoderContext::getError() const {

template <typename CompressionContext>
jsg::Ref<ZlibUtil::BrotliCompressionStream<CompressionContext>> ZlibUtil::BrotliCompressionStream<
CompressionContext>::constructor(ZlibModeValue mode) {
return jsg::alloc<BrotliCompressionStream>(static_cast<ZlibMode>(mode));
CompressionContext>::constructor(jsg::Lock& js, ZlibModeValue mode) {
return jsg::alloc<BrotliCompressionStream>(js, static_cast<ZlibMode>(mode));
}

template <typename CompressionContext>
Expand Down Expand Up @@ -806,10 +807,10 @@ static kj::Array<kj::byte> syncProcessBuffer(Context& ctx, GrowableBuffer& resul
} // namespace

kj::Array<kj::byte> ZlibUtil::zlibSync(
ZlibUtil::InputSource data, ZlibContext::Options opts, ZlibModeValue mode) {
jsg::Lock& js, ZlibUtil::InputSource data, ZlibContext::Options opts, ZlibModeValue mode) {
// Any use of zlib APIs constitutes an implicit dependency on Allocator which must
// remain alive until the zlib stream is destroyed
CompressionAllocator allocator;
CompressionAllocator allocator(js);
ZlibContext ctx(static_cast<ZlibMode>(mode));
allocator.configure(ctx.getStream());

Expand Down Expand Up @@ -849,7 +850,7 @@ void ZlibUtil::zlibWithCallback(jsg::Lock& js,
CompressCallback cb) {
// Capture only relevant errors so they can be passed to the callback
auto res = js.tryCatch([&]() {
return CompressCallbackArg(zlibSync(kj::mv(data), kj::mv(options), mode));
return CompressCallbackArg(zlibSync(js, kj::mv(data), kj::mv(options), mode));
}, [&](jsg::Value&& exception) {
return CompressCallbackArg(jsg::JsValue(exception.getHandle(js)));
});
Expand All @@ -859,10 +860,11 @@ void ZlibUtil::zlibWithCallback(jsg::Lock& js,
}

template <typename Context>
kj::Array<kj::byte> ZlibUtil::brotliSync(InputSource data, BrotliContext::Options opts) {
kj::Array<kj::byte> ZlibUtil::brotliSync(
jsg::Lock& js, InputSource data, BrotliContext::Options opts) {
// Any use of brotli APIs constitutes an implicit dependency on Allocator which must
// remain alive until the brotli state is destroyed
CompressionAllocator allocator;
CompressionAllocator allocator(js);
Context ctx(Context::Mode);

auto chunkSize = opts.chunkSize.orDefault(ZLIB_PERFORMANT_CHUNK_SIZE);
Expand Down Expand Up @@ -914,7 +916,7 @@ void ZlibUtil::brotliWithCallback(
jsg::Lock& js, InputSource data, BrotliContext::Options options, CompressCallback cb) {
// Capture only relevant errors so they can be passed to the callback
auto res = js.tryCatch([&]() {
return CompressCallbackArg(brotliSync<Context>(kj::mv(data), kj::mv(options)));
return CompressCallbackArg(brotliSync<Context>(js, kj::mv(data), kj::mv(options)));
}, [&](jsg::Value&& exception) {
return CompressCallbackArg(jsg::JsValue(exception.getHandle(js)));
});
Expand All @@ -933,25 +935,25 @@ void ZlibUtil::brotliWithCallback(
jsg::Optional<kj::Array<kj::byte>> input, int inputOffset, int inputLength, \
kj::Array<kj::byte> output, int outputOffset, int outputLength); \
template jsg::Ref<ZlibUtil::CompressionStream<T>> ZlibUtil::CompressionStream<T>::constructor( \
ZlibModeValue mode);
jsg::Lock& js, ZlibModeValue mode);

CREATE_TEMPLATE(ZlibContext)
CREATE_TEMPLATE(BrotliEncoderContext)
CREATE_TEMPLATE(BrotliDecoderContext)

template jsg::Ref<ZlibUtil::BrotliCompressionStream<BrotliEncoderContext>> ZlibUtil::
BrotliCompressionStream<BrotliEncoderContext>::constructor(ZlibModeValue mode);
BrotliCompressionStream<BrotliEncoderContext>::constructor(jsg::Lock& js, ZlibModeValue mode);
template jsg::Ref<ZlibUtil::BrotliCompressionStream<BrotliDecoderContext>> ZlibUtil::
BrotliCompressionStream<BrotliDecoderContext>::constructor(ZlibModeValue mode);
BrotliCompressionStream<BrotliDecoderContext>::constructor(jsg::Lock& js, ZlibModeValue mode);
template bool ZlibUtil::BrotliCompressionStream<BrotliEncoderContext>::initialize(
jsg::Lock&, jsg::BufferSource, jsg::BufferSource, jsg::Function<void()>);
template bool ZlibUtil::BrotliCompressionStream<BrotliDecoderContext>::initialize(
jsg::Lock&, jsg::BufferSource, jsg::BufferSource, jsg::Function<void()>);

template kj::Array<kj::byte> ZlibUtil::brotliSync<BrotliEncoderContext>(
InputSource data, BrotliContext::Options opts);
jsg::Lock& js, InputSource data, BrotliContext::Options opts);
template kj::Array<kj::byte> ZlibUtil::brotliSync<BrotliDecoderContext>(
InputSource data, BrotliContext::Options opts);
jsg::Lock& js, InputSource data, BrotliContext::Options opts);
template void ZlibUtil::brotliWithCallback<BrotliEncoderContext>(
jsg::Lock& js, InputSource data, BrotliContext::Options options, CompressCallback cb);
template void ZlibUtil::brotliWithCallback<BrotliDecoderContext>(
Expand Down
20 changes: 10 additions & 10 deletions src/workerd/api/node/zlib-util.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -293,13 +293,12 @@ class ZlibUtil final: public jsg::Object {
template <class CompressionContext>
class CompressionStream: public jsg::Object {
public:
explicit CompressionStream(ZlibMode _mode): context_(_mode) {}
CompressionStream() = default;
explicit CompressionStream(jsg::Lock& js, ZlibMode _mode): allocator(js), context_(_mode) {}
// TODO(soon): Find a way to add noexcept(false) to this destructor.
~CompressionStream();
KJ_DISALLOW_COPY_AND_MOVE(CompressionStream);

static jsg::Ref<CompressionStream> constructor(ZlibModeValue mode);
static jsg::Ref<CompressionStream> constructor(jsg::Lock& js, ZlibModeValue mode);

void close();
bool checkError(jsg::Lock& js);
Expand Down Expand Up @@ -359,9 +358,9 @@ class ZlibUtil final: public jsg::Object {

class ZlibStream final: public CompressionStream<ZlibContext> {
public:
explicit ZlibStream(ZlibMode mode): CompressionStream(mode) {}
explicit ZlibStream(jsg::Lock& js, ZlibMode mode): CompressionStream(js, mode) {}
KJ_DISALLOW_COPY_AND_MOVE(ZlibStream);
static jsg::Ref<ZlibStream> constructor(ZlibModeValue mode);
static jsg::Ref<ZlibStream> constructor(jsg::Lock& js, ZlibModeValue mode);

// Instance methods
void initialize(int windowBits,
Expand All @@ -384,10 +383,10 @@ class ZlibUtil final: public jsg::Object {
template <typename CompressionContext>
class BrotliCompressionStream: public CompressionStream<CompressionContext> {
public:
explicit BrotliCompressionStream(ZlibMode _mode)
: CompressionStream<CompressionContext>(_mode) {}
explicit BrotliCompressionStream(jsg::Lock& js, ZlibMode _mode)
: CompressionStream<CompressionContext>(js, _mode) {}
KJ_DISALLOW_COPY_AND_MOVE(BrotliCompressionStream);
static jsg::Ref<BrotliCompressionStream> constructor(ZlibModeValue mode);
static jsg::Ref<BrotliCompressionStream> constructor(jsg::Lock& js, ZlibModeValue mode);

bool initialize(jsg::Lock& js,
jsg::BufferSource params,
Expand Down Expand Up @@ -423,10 +422,11 @@ class ZlibUtil final: public jsg::Object {
ZlibContext::Options options,
ZlibModeValue mode,
CompressCallback cb);
kj::Array<kj::byte> zlibSync(InputSource data, ZlibContext::Options options, ZlibModeValue mode);
kj::Array<kj::byte> zlibSync(
jsg::Lock& js, InputSource data, ZlibContext::Options options, ZlibModeValue mode);

template <typename Context>
kj::Array<kj::byte> brotliSync(InputSource data, BrotliContext::Options options);
kj::Array<kj::byte> brotliSync(jsg::Lock& js, InputSource data, BrotliContext::Options options);
template <typename Context>
void brotliWithCallback(
jsg::Lock& js, InputSource data, BrotliContext::Options options, CompressCallback cb);
Expand Down
36 changes: 26 additions & 10 deletions src/workerd/api/streams/compression.c++
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,10 @@

namespace workerd::api {

CompressionAllocator::CompressionAllocator(jsg::Lock& js) {
memoryAdjustment = js.getExternalMemoryAdjustment(0);
}

void CompressionAllocator::configure(z_stream* stream) {
stream->zalloc = AllocForZlib;
stream->zfree = FreeForZlib;
Expand All @@ -27,17 +31,28 @@ void* CompressionAllocator::AllocForZlib(void* data, uInt items, uInt size) {
// TODO(soon): Enabling V8 reporting breaks CompressionStreamImpl on edgeworker. Investigate.
// This is mostly likely due to running CompressionStream without an isolate lock.
void* CompressionAllocator::AllocForBrotli(void* opaque, size_t size) {
auto* thisAllocator = static_cast<CompressionAllocator*>(opaque);
auto data = kj::heapArray<uint8_t>(size);
auto* allocator = static_cast<CompressionAllocator*>(opaque);
auto data = kj::heapArray<kj::byte>(size);
auto begin = data.begin();
thisAllocator->allocations.insert(begin, kj::mv(data));
*reinterpret_cast<size_t*>(begin) = size;
auto isolate = v8::Isolate::GetCurrent();
KJ_REQUIRE_NONNULL(isolate, "Isolate lock is required"_kj);
auto& js = jsg::Lock::from(isolate);
allocator->memoryAdjustment.set(js, allocator->memoryAdjustment.getAmount() + size);
allocator->allocations.insert(begin, kj::mv(data));
return begin;
}

void CompressionAllocator::FreeForZlib(void* opaque, void* pointer) {
if (KJ_UNLIKELY(pointer == nullptr)) return;
auto* thisAllocator = static_cast<CompressionAllocator*>(opaque);
JSG_REQUIRE(thisAllocator->allocations.erase(pointer), Error, "Zlib allocation should exist"_kj);
auto* allocator = static_cast<CompressionAllocator*>(opaque);
auto* real_pointer = static_cast<kj::byte*>(pointer) - sizeof(size_t);
size_t real_size = *reinterpret_cast<size_t*>(real_pointer);
auto isolate = v8::Isolate::GetCurrent();
KJ_REQUIRE_NONNULL(isolate, "Isolate lock is required"_kj);
auto& js = jsg::Lock::from(isolate);
allocator->memoryAdjustment.set(js, allocator->memoryAdjustment.getAmount() - real_size);
JSG_REQUIRE(allocator->allocations.erase(real_pointer), Error, "Zlib allocation should exist"_kj);
}

namespace {
Expand All @@ -59,8 +74,9 @@ public:
kj::ArrayPtr<const byte> buffer;
};

explicit Context(Mode mode, kj::StringPtr format, ContextFlags flags)
explicit Context(jsg::Lock& js, Mode mode, kj::StringPtr format, ContextFlags flags)
: mode(mode),
allocator(js),
strictCompression(flags) {
// Configure allocator before any stream operations.
allocator.configure(&ctx);
Expand Down Expand Up @@ -225,8 +241,8 @@ class CompressionStreamImpl: public kj::Refcounted,
public ReadableStreamSource,
public WritableStreamSink {
public:
explicit CompressionStreamImpl(kj::String format, Context::ContextFlags flags)
: context(mode, format, flags) {}
explicit CompressionStreamImpl(jsg::Lock& js, kj::String format, Context::ContextFlags flags)
: context(js, mode, format, flags) {}

// WritableStreamSink implementation ---------------------------------------------------

Expand Down Expand Up @@ -472,7 +488,7 @@ jsg::Ref<CompressionStream> CompressionStream::constructor(jsg::Lock& js, kj::St
"The compression format must be either 'deflate', 'deflate-raw' or 'gzip'.");

auto readableSide = kj::refcounted<CompressionStreamImpl<Context::Mode::COMPRESS>>(
kj::mv(format), Context::ContextFlags::NONE);
js, kj::mv(format), Context::ContextFlags::NONE);
auto writableSide = kj::addRef(*readableSide);

auto& ioContext = IoContext::current();
Expand All @@ -486,7 +502,7 @@ jsg::Ref<DecompressionStream> DecompressionStream::constructor(jsg::Lock& js, kj
"The compression format must be either 'deflate', 'deflate-raw' or 'gzip'.");

auto readableSide =
kj::refcounted<CompressionStreamImpl<Context::Mode::DECOMPRESS>>(kj::mv(format),
kj::refcounted<CompressionStreamImpl<Context::Mode::DECOMPRESS>>(js, kj::mv(format),
FeatureFlags::get(js).getStrictCompression() ? Context::ContextFlags::STRICT
: Context::ContextFlags::NONE);
auto writableSide = kj::addRef(*readableSide);
Expand Down
9 changes: 2 additions & 7 deletions src/workerd/api/streams/compression.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,20 +16,15 @@ namespace workerd::api {
// isolate pointer and use that to get the external memory adjustment.
class CompressionAllocator final {
public:
CompressionAllocator() = default;
CompressionAllocator(jsg::Lock& js);
void configure(z_stream* stream);

static void* AllocForZlib(void* data, uInt items, uInt size);
static void* AllocForBrotli(void* data, size_t size);
static void FreeForZlib(void* data, void* pointer);

private:
// TODO(soon): Use this instead of array<byte> after fixing edgeworker.
// struct Allocation {
// kj::Array<kj::byte> data;
// kj::Maybe<jsg::ExternalMemoryAdjustment> memoryAdjustment;
// };

jsg::ExternalMemoryAdjustment memoryAdjustment;
kj::HashMap<void*, kj::Array<kj::byte>> allocations;
};

Expand Down

0 comments on commit 23753d2

Please sign in to comment.