1
// Copyright (c) 2017-2022 Cloudflare, Inc.
2
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
//     https://opensource.org/licenses/Apache-2.0
4
#include "async-hooks.h"
5
 
6
#include <workerd/io/features.h>
7
#include <workerd/io/io-context.h>
8
 
9
namespace workerd::api::node {
10
 
11
namespace {
12
// If there is a current IoContext, then it is possible/likely that the
13
// current AsyncContextFrame is storing values that are bound to that
14
// IoContext. In that case, we want to protect against the case where
15
// the returned snapshot function is called from a different IoContext.
16
// To do this we will capture a weak reference to the current IoContext
17
// and check it against the current IoContext where the snapshot function
18
// is invoked.
19
jsg::Function<void()> getValidator(jsg::Lock& js) {
20
  kj::Maybe<kj::Own<IoContext::WeakRef>> maybeIoContext;
21
  if (FeatureFlags::get(js).getBindAsyncLocalStorageSnapshot() && IoContext::hasCurrent()) {
22
    // We use a weak reference to the IoContext because the current IoContext
23
    // may be destroyed before the snapshot function is called.
24
    maybeIoContext = IoContext::current().getWeakRef();
25
  }
26
 
27
  static constexpr auto kErrorMessage =
28
      "Cannot call this AsyncLocalStorage bound function outside of the "
29
      "request in which it was created."_kj;
30
 
31
  return [maybeIoContext = kj::mv(maybeIoContext)](jsg::Lock&) {
32
    KJ_IF_SOME(originIoContext, maybeIoContext) {
33
      // We had an IoContext when we created the snapshot function.
34
      // If it is not the current IoContext, or if there is no current
35
      // IoContext, or if the captured IoContext has been destroyed,
36
      // we throw an error.
37
      JSG_REQUIRE(IoContext::hasCurrent() && originIoContext->isValid(), Error, kErrorMessage);
38
      originIoContext->runIfAlive([&](IoContext& otherContext) {
39
        JSG_REQUIRE(&otherContext == &IoContext::current(), Error, kErrorMessage);
40
      });
41
    }
42
  };
43
}
44
 
45
}  // namespace
46
 
47
jsg::Ref<AsyncLocalStorage> AsyncLocalStorage::constructor(
48
    jsg::Lock& js, jsg::Optional<AsyncLocalStorage::AsyncLocalStorageOptions> options) {
49
  return js.alloc<AsyncLocalStorage>(kj::mv(options));
50
}
51
 
52
v8::Local<v8::Value> AsyncLocalStorage::run(jsg::Lock& js,
53
    v8::Local<v8::Value> store,
54
    jsg::Function<v8::Local<v8::Value>(jsg::Arguments<jsg::Value>)> callback,
55
    jsg::Arguments<jsg::Value> args) {
56
  callback.setReceiver(js.v8Ref<v8::Value>(js.v8Context()->Global()));
57
  jsg::AsyncContextFrame::StorageScope scope(js, *key, js.v8Ref(store));
58
  return callback(js, kj::mv(args));
59
}
60
 
61
v8::Local<v8::Value> AsyncLocalStorage::exit(jsg::Lock& js,
62
    jsg::Function<v8::Local<v8::Value>(jsg::Arguments<jsg::Value>)> callback,
63
    jsg::Arguments<jsg::Value> args) {
64
  // Node.js defines exit as running "a function synchronously outside of a context".
65
  // It goes on to say that the store is not accessible within the callback or the
66
  // asynchronous operations created within the callback. Any getStore() call done
67
  // within the callback function will always return undefined... except if run() is
68
  // called which implicitly enables the context again within that scope.
69
  //
70
  // We do not have to emulate Node.js enable/disable behavior since we are not
71
  // implementing the enterWith/disable methods. We can emulate the correct
72
  // behavior simply by calling run with the store value set to undefined, which
73
  // will propagate correctly.
74
  return run(js, js.undefined(), kj::mv(callback), kj::mv(args));
75
}
76
 
77
v8::Local<v8::Value> AsyncLocalStorage::getStore(jsg::Lock& js) {
78
  KJ_IF_SOME(context, jsg::AsyncContextFrame::current(js)) {
79
    KJ_IF_SOME(value, context.get(*key)) {
80
      return value.getHandle(js);
81
    }
82
  }
83
  KJ_IF_SOME(value, defaultValue) {
84
    return value.getHandle(js);
85
  }
86
  return js.undefined();
87
}
88
 
89
kj::StringPtr AsyncLocalStorage::getName() {
90
  KJ_IF_SOME(n, name) {
91
    return n.asPtr();
92
  }
93
  return nullptr;
94
}
95
 
96
v8::Local<v8::Function> AsyncLocalStorage::bind(jsg::Lock& js, v8::Local<v8::Function> fn) {
97
  KJ_IF_SOME(frame, jsg::AsyncContextFrame::current(js)) {
98
    return frame.wrap(js, fn, getValidator(js));
99
  } else {
100
    return jsg::AsyncContextFrame::wrapRoot(js, fn);
101
  }
102
}
103
 
104
v8::Local<v8::Function> AsyncLocalStorage::snapshot(jsg::Lock& js) {
105
  return jsg::AsyncContextFrame::wrapSnapshot(js, getValidator(js));
106
}
107
 
108
namespace {
109
kj::Maybe<jsg::Ref<jsg::AsyncContextFrame>> tryGetFrameRef(jsg::Lock& js) {
110
  return jsg::AsyncContextFrame::current(js).map(
111
      [](jsg::AsyncContextFrame& frame) { return frame.addRef(); });
112
}
113
}  // namespace
114
 
115
AsyncResource::AsyncResource(jsg::Lock& js): frame(tryGetFrameRef(js)) {}
116
 
117
jsg::Ref<AsyncResource> AsyncResource::constructor(
118
    jsg::Lock& js, jsg::Optional<kj::String> type, jsg::Optional<Options> options) {
119
  // The type and options are required as part of the Node.js API compatibility
120
  // but our implementation does not currently make use of them at all. It is OK
121
  // for us to silently ignore both here.
122
  return js.alloc<AsyncResource>(js);
123
}
124
 
125
v8::Local<v8::Function> AsyncResource::staticBind(jsg::Lock& js,
126
    v8::Local<v8::Function> fn,
127
    jsg::Optional<kj::String> type,
128
    jsg::Optional<v8::Local<v8::Value>> thisArg,
129
    const jsg::TypeHandler<jsg::Ref<AsyncResource>>& handler) {
130
  return AsyncResource::constructor(js, kj::mv(type).orDefault([] {
131
    return kj::str("AsyncResource");
132
  }))->bind(js, fn, thisArg, handler);
133
}
134
 
135
kj::Maybe<jsg::AsyncContextFrame&> AsyncResource::getFrame() {
136
  return frame.map([](jsg::Ref<jsg::AsyncContextFrame>& frame) -> jsg::AsyncContextFrame& {
137
    return *(frame.get());
138
  });
139
}
140
 
141
v8::Local<v8::Function> AsyncResource::bind(jsg::Lock& js,
142
    v8::Local<v8::Function> fn,
143
    jsg::Optional<v8::Local<v8::Value>> thisArg,
144
    const jsg::TypeHandler<jsg::Ref<AsyncResource>>& handler) {
145
  v8::Local<v8::Function> bound;
146
  KJ_IF_SOME(frame, getFrame()) {
147
    bound = frame.wrap(js, fn, getValidator(js), thisArg);
148
  } else {
149
    bound = jsg::AsyncContextFrame::wrapRoot(js, fn, thisArg);
150
  }
151
 
152
  // Per Node.js documentation (https://nodejs.org/dist/latest-v19.x/docs/api/async_context.html#asyncresourcebindfn-thisarg), the returned function "will have an
153
  // asyncResource property referencing the AsyncResource to which the function
154
  // is bound".
155
  js.v8Set(bound, "asyncResource"_kj, handler.wrap(js, JSG_THIS));
156
  return bound;
157
}
158
 
159
v8::Local<v8::Value> AsyncResource::runInAsyncScope(jsg::Lock& js,
160
    jsg::Function<v8::Local<v8::Value>(jsg::Arguments<jsg::Value>)> fn,
161
    jsg::Optional<v8::Local<v8::Value>> thisArg,
162
    jsg::Arguments<jsg::Value> args) {
163
  v8::Local<v8::Value> receiver = js.v8Context()->Global();
164
  KJ_IF_SOME(arg, thisArg) {
165
    receiver = arg;
166
  }
167
  fn.setReceiver(js.v8Ref<v8::Value>(receiver));
168
  jsg::AsyncContextFrame::Scope scope(js, getFrame());
169
  return fn(js, kj::mv(args));
170
}
171
 
172
kj::Own<jsg::AsyncContextFrame::StorageKey> AsyncLocalStorage::getKey() {
173
  return kj::addRef(*key);
174
}
175
 
176
}  // namespace workerd::api::node